From 20b030628d3eacb24d451bd0856378308d1e67ec Mon Sep 17 00:00:00 2001
From: Matthias van de Meent <boekewurm+postgres@gmail.com>
Date: Wed, 31 Jan 2024 02:21:00 +0100
Subject: [PATCH v15 6/6] btree: Specialize various performance-sensitive
 functions on btree key shape.

nbtree keys are not all made the same, and a significant amount of time is
spent in code that exists only to deal with various key shapes. By specializing
function calls based on the key shape, we can remove or reduce these causes
of overhead.

This commit adds the basic infrastructure for specializing specific hot code
in the nbtree AM to certain shapes of keys, and splits the code that can
benefit from attribute offset optimizations into separate files. This does
NOT yet update the code itself - it just makes the code compile cleanly.
The performance should be comparable if not the same.
---
 src/include/access/itup_attiter.h | 199 ++++++++++++++++++++++++++++++
 src/include/access/nbtree.h       |  20 ++-
 src/include/access/nbtree_spec.h  |  36 ++++++
 3 files changed, 252 insertions(+), 3 deletions(-)
 create mode 100644 src/include/access/itup_attiter.h

diff --git a/src/include/access/itup_attiter.h b/src/include/access/itup_attiter.h
new file mode 100644
index 0000000000..abf52761e7
--- /dev/null
+++ b/src/include/access/itup_attiter.h
@@ -0,0 +1,199 @@
+/*-------------------------------------------------------------------------
+ *
+ * itup_attiter.h
+ *	  POSTGRES index tuple attribute iterator definitions.
+ *
+ *
+ * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * src/include/access/itup_attiter.h
+ *
+ *-------------------------------------------------------------------------
+ */
+#ifndef ITUP_ATTITER_H
+#define ITUP_ATTITER_H
+
+#include "access/itup.h"
+#include "varatt.h"
+
+typedef struct IAttrIterStateData
+{
+	int			offset;
+	bool		slow;
+	bool		isNull;
+} IAttrIterStateData;
+
+typedef IAttrIterStateData * IAttrIterState;
+
+/* ----------------
+ *		index_attiterinit
+ *
+ *		This gets called many times, so we macro the cacheable and NULL
+ *		lookups, and call nocache_index_attiterinit() for the rest.
+ *
+ *		tup - the tuple being iterated on
+ *		attnum - the attribute number that we start the iteration with
+ *				 in the first index_attiternext call
+ *		tupdesc - the tuple description
+ *
+ * ----------------
+ */
+#define index_attiterinit(tup, attnum, tupleDesc, iter) \
+do { \
+	if ((attnum) == 1) \
+	{ \
+		*(iter) = ((IAttrIterStateData) { \
+			0 /* Offset of attribute 1 is always 0 */, \
+			false /* slow */, \
+			false /* isNull */ \
+		}); \
+	} \
+	else if (!IndexTupleHasNulls(tup) && \
+			 TupleDescAttr((tupleDesc), (attnum)-1)->attcacheoff >= 0) \
+	{ \
+		*(iter) = ((IAttrIterStateData) { \
+			TupleDescAttr((tupleDesc), (attnum)-1)->attcacheoff, /* offset */ \
+			false, /* slow */ \
+			false /* isNull */ \
+		}); \
+	} \
+	else \
+		nocache_index_attiterinit((tup), (attnum) - 1, (tupleDesc), (iter)); \
+} while (false);
+
+/*
+ * Initiate an index attribute iterator to attribute attnum,
+ * and return the corresponding datum.
+ *
+ * This is nearly the same as index_deform_tuple, except that this
+ * returns the internal state up to attnum, instead of populating the
+ * datum- and isnull-arrays
+ */
+static inline void
+nocache_index_attiterinit(IndexTuple tup, AttrNumber attnum, TupleDesc tupleDesc, IAttrIterState iter)
+{
+	bool		hasnulls = IndexTupleHasNulls(tup);
+	int			curatt;
+	char	   *tp;				/* ptr to tuple data */
+	int			off;			/* offset in tuple data */
+	bits8	   *bp;				/* ptr to null bitmap in tuple */
+	bool		slow = false;	/* can we use/set attcacheoff? */
+	bool		null = false;
+
+	/* Assert to protect callers */
+	Assert(PointerIsValid(iter));
+	Assert(tupleDesc->natts <= INDEX_MAX_KEYS);
+	Assert(attnum <= tupleDesc->natts);
+	Assert(attnum > 0);
+
+	/* XXX "knows" t_bits are just after fixed tuple header! */
+	bp = (bits8 *) ((char *) tup + sizeof(IndexTupleData));
+
+	tp = (char *) tup + IndexInfoFindDataOffset(tup->t_info);
+	off = 0;
+
+	for (curatt = 0; curatt < attnum; curatt++)
+	{
+		Form_pg_attribute thisatt = TupleDescAttr(tupleDesc, curatt);
+
+		if (hasnulls && att_isnull(curatt, bp))
+		{
+			null = true;
+			slow = true;		/* can't use attcacheoff anymore */
+			continue;
+		}
+
+		null = false;
+
+		if (!slow && thisatt->attcacheoff >= 0)
+			off = thisatt->attcacheoff;
+		else if (thisatt->attlen == -1)
+		{
+			off = att_align_pointer(off, thisatt->attalign, -1,
+									tp + off);
+			slow = true;
+		}
+		else
+		{
+			/* not varlena, so safe to use att_align_nominal */
+			off = att_align_nominal(off, thisatt->attalign);
+		}
+
+		off = att_addlength_pointer(off, thisatt->attlen, tp + off);
+
+		if (thisatt->attlen <= 0)
+			slow = true;		/* can't use attcacheoff anymore */
+	}
+
+	iter->isNull = null;
+	iter->offset = off;
+	iter->slow = slow;
+}
+
+/* ----------------
+ *		index_attiternext() - get the next attribute of an index tuple
+ *
+ *		This gets called many times, so we do the least amount of work
+ *		possible.
+ *
+ *		The code does not attempt to update attcacheoff; as it is unlikely
+ *		to reach a situation where the cached offset matters a lot.
+ *		If the cached offset do matter, the caller should make sure that
+ *		PopulateTupleDescCacheOffsets() was called on the tuple descriptor
+ *		to populate the attribute offset cache.
+ *
+ * ----------------
+ */
+static inline Datum
+index_attiternext(IndexTuple tup, AttrNumber attnum, TupleDesc tupleDesc, IAttrIterState iter)
+{
+	bool		hasnulls = IndexTupleHasNulls(tup);
+	char	   *tp;				/* ptr to tuple data */
+	bits8	   *bp;				/* ptr to null bitmap in tuple */
+	Datum		datum;
+	Form_pg_attribute thisatt = TupleDescAttr(tupleDesc, attnum - 1);
+
+	Assert(PointerIsValid(iter));
+	Assert(tupleDesc->natts <= INDEX_MAX_KEYS);
+	Assert(attnum <= tupleDesc->natts);
+	Assert(attnum > 0);
+
+	bp = (bits8 *) ((char *) tup + sizeof(IndexTupleData));
+
+	tp = (char *) tup + IndexInfoFindDataOffset(tup->t_info);
+
+	if (hasnulls && att_isnull(attnum - 1, bp))
+	{
+		iter->isNull = true;
+		iter->slow = true;
+		return (Datum) 0;
+	}
+
+	iter->isNull = false;
+
+	if (!iter->slow && thisatt->attcacheoff >= 0)
+		iter->offset = thisatt->attcacheoff;
+	else if (thisatt->attlen == -1)
+	{
+		iter->offset = att_align_pointer(iter->offset, thisatt->attalign, -1,
+										 tp + iter->offset);
+		iter->slow = true;
+	}
+	else
+	{
+		/* not varlena, so safe to use att_align_nominal */
+		iter->offset = att_align_nominal(iter->offset, thisatt->attalign);
+	}
+
+	datum = fetchatt(thisatt, tp + iter->offset);
+
+	iter->offset = att_addlength_pointer(iter->offset, thisatt->attlen, tp + iter->offset);
+
+	if (thisatt->attlen <= 0)
+		iter->slow = true;		/* can't use attcacheoff anymore */
+
+	return datum;
+}
+
+#endif /* ITUP_ATTITER_H */
diff --git a/src/include/access/nbtree.h b/src/include/access/nbtree.h
index 9579fed492..58d8a0fdc3 100644
--- a/src/include/access/nbtree.h
+++ b/src/include/access/nbtree.h
@@ -16,6 +16,7 @@
 
 #include "access/amapi.h"
 #include "access/itup.h"
+#include "access/itup_attiter.h"
 #include "access/sdir.h"
 #include "access/tableam.h"
 #include "access/xlogreader.h"
@@ -1120,6 +1121,7 @@ typedef struct BTOptions
 
 #define NBTS_TYPE_CACHED CACHED
 #define NBTS_TYPE_SINGLE_KEYATT SINGLE_KEYATT
+#define NBTS_TYPE_UNCACHED UNCACHED
 
 /* Generate the specialized function headers */
 #define NBT_FILE "access/nbtree_specfuncs.h"
@@ -1132,15 +1134,23 @@ typedef struct BTOptions
 typedef enum NBTS_CTX {
 	NBTS_CTX_CACHED = 1, /* Equivalent to unspecialized code */
 	NBTS_CTX_SINGLE_KEYATT, /* Single key column index; cannot handle more */
+	NBTS_CTX_UNCACHED, /* Last key attribute can't use attcachoff */
 } NBTS_CTX;
 
 static inline NBTS_CTX _nbt_spec_context(Relation irel)
 {
+	AttrNumber	nKeyAtts;
+
 	Assert(PointerIsValid(irel));
 
-	if (IndexRelationGetNumberOfKeyAttributes(irel) == 1)
+	nKeyAtts = IndexRelationGetNumberOfKeyAttributes(irel);
+
+	if (nKeyAtts == 1)
 		return NBTS_CTX_SINGLE_KEYATT;
 
+	if (TupleDescAttr(irel->rd_att, nKeyAtts - 1)->attcacheoff < -1)
+		return NBTS_CTX_UNCACHED;
+
 	return NBTS_CTX_CACHED;
 }
 
@@ -1183,8 +1193,12 @@ static inline Datum _bt_getfirstatt(IndexTuple tuple, TupleDesc tupleDesc,
 	(NBTS_CTX_NAME == NBTS_CTX_SINGLE_KEYATT) ? ( \
 		NBTS_MAKE_NAME(name, NBTS_TYPE_SINGLE_KEYATT) \
 	) : ( \
-		AssertMacro((NBTS_CTX_NAME) == NBTS_CTX_CACHED), \
-		NBTS_MAKE_NAME(name, NBTS_TYPE_CACHED) \
+		((NBTS_CTX_NAME) == NBTS_CTX_UNCACHED) ? ( \
+			NBTS_MAKE_NAME(name, NBTS_TYPE_UNCACHED) \
+		) : ( \
+			AssertMacro((NBTS_CTX_NAME) == NBTS_CTX_CACHED), \
+			NBTS_MAKE_NAME(name, NBTS_TYPE_CACHED) \
+		) \
 	) \
 )
 
diff --git a/src/include/access/nbtree_spec.h b/src/include/access/nbtree_spec.h
index 4919f95e79..2f3f615ea4 100644
--- a/src/include/access/nbtree_spec.h
+++ b/src/include/access/nbtree_spec.h
@@ -153,6 +153,42 @@
 #undef nbts_attiter_nextattdatum
 #undef nbts_attiter_curattisnull
 
+/*
+ * Specialization 3: UNCACHED
+ *
+ * Multiple key columns, but attcacheoff -optimization doesn't apply because
+ * some intermediate or prefixing key attributes that don't have a fixed size
+ * at the type level (i.e. their attlen < 0).
+ */
+#define NBTS_TYPE NBTS_TYPE_UNCACHED
+
+#define nbts_attiterdeclare(itup) \
+	IAttrIterStateData	NBTS_MAKE_NAME(itup, iter)
+
+#define nbts_attiterinit(itup, initAttNum, tupDesc) \
+	index_attiterinit((itup), (initAttNum), (tupDesc), &(NBTS_MAKE_NAME(itup, iter)))
+
+#define nbts_foreachattr(initAttNum, endAttNum) \
+	for (int spec_i = (initAttNum); spec_i <= (endAttNum); spec_i++)
+
+#define nbts_attiter_attnum spec_i
+
+#define nbts_attiter_nextattdatum(itup, tupDesc) \
+	index_attiternext((itup), spec_i, (tupDesc), &(NBTS_MAKE_NAME(itup, iter)))
+
+#define nbts_attiter_curattisnull(itup) \
+	NBTS_MAKE_NAME(itup, iter).isNull
+
+#include NBT_SPECIALIZE_FILE
+
+#undef NBTS_TYPE
+#undef nbts_attiterdeclare
+#undef nbts_attiterinit
+#undef nbts_foreachattr
+#undef nbts_attiter_attnum
+#undef nbts_attiter_nextattdatum
+#undef nbts_attiter_curattisnull
+
 /*
  * We're done with templating, so restore or create the macros which can be
  * used in non-template sources.
-- 
2.42.1

