v2-0001-Allow-plugins-to-Jumble-an-expression.patch

From 698a4d0c8e2230fdb08b31f154a3b04fa264e8ab Mon Sep 17 00:00:00 2001
From: Sami Imseih <simseih@amazon.com>
Date: Mon, 31 Mar 2025 09:59:16 -0700
Subject: [PATCH v2] Allow plugins to Jumble an expression.

This change makes _jumbleNode available to plugins that wish to
perform jumbling on nodes other than Query. To facilitate this
capability, two new routines to initialize a jumble state and to
produce a 64-bit hash from the jumble are also made available.

It may also be a good idea to separate these aforementioned
routines into a separate C file, as they can be used for more
than query jumbling. That could be done in a future change.

Discussion: https://www.postgresql.org/message-id/Z9khOo14yzZHCnmn%40paquier.xyz
---
 src/backend/nodes/queryjumblefuncs.c | 114 +++++++++++++--------------
 src/include/nodes/queryjumble.h      |  49 ++++++++++++
 2 files changed, 106 insertions(+), 57 deletions(-)

diff --git a/src/backend/nodes/queryjumblefuncs.c b/src/backend/nodes/queryjumblefuncs.c
index 513cf92d357..41e8cb9e5be 100644
--- a/src/backend/nodes/queryjumblefuncs.c
+++ b/src/backend/nodes/queryjumblefuncs.c
@@ -55,14 +55,11 @@ int			compute_query_id = COMPUTE_QUERY_ID_AUTO;
  */
 bool		query_id_enabled = false;
 
-static JumbleState *InitJumble(void);
+static JumbleState *InitJumbleInternal(bool record_clocations);
 static uint64 DoJumble(JumbleState *jstate, Node *node);
-static void AppendJumble(JumbleState *jstate,
-						 const unsigned char *item, Size size);
 static void FlushPendingNulls(JumbleState *jstate);
 static void RecordConstLocation(JumbleState *jstate,
 								int location, bool merged);
-static void _jumbleNode(JumbleState *jstate, Node *node);
 static void _jumbleElements(JumbleState *jstate, List *elements);
 static void _jumbleA_Const(JumbleState *jstate, Node *node);
 static void _jumbleList(JumbleState *jstate, Node *node);
@@ -133,7 +130,7 @@ JumbleQuery(Query *query)
 
 	Assert(IsQueryIdEnabled());
 
-	jstate = InitJumble();
+	jstate = InitJumbleInternal(true);
 
 	query->queryId = DoJumble(jstate, (Node *) query);
 
@@ -166,11 +163,11 @@ EnableQueryId(void)
 }
 
 /*
- * InitJumble
+ * InitJumbleInternal
  *		Allocate a JumbleState object and make it ready to jumble.
  */
 static JumbleState *
-InitJumble(void)
+InitJumbleInternal(bool record_clocations)
 {
 	JumbleState *jstate;
 
@@ -179,9 +176,19 @@ InitJumble(void)
 	/* Set up workspace for query jumbling */
 	jstate->jumble = (unsigned char *) palloc(JUMBLE_SIZE);
 	jstate->jumble_len = 0;
-	jstate->clocations_buf_size = 32;
-	jstate->clocations = (LocationLen *) palloc(jstate->clocations_buf_size *
-												sizeof(LocationLen));
+
+	if (record_clocations)
+	{
+		jstate->clocations_buf_size = 32;
+		jstate->clocations = (LocationLen *)
+			palloc(jstate->clocations_buf_size * sizeof(LocationLen));
+	}
+	else
+	{
+		jstate->clocations_buf_size = 0;
+		jstate->clocations = NULL;
+	}
+
 	jstate->clocations_count = 0;
 	jstate->highest_extern_param_id = 0;
 	jstate->pending_nulls = 0;
@@ -193,16 +200,21 @@ InitJumble(void)
 }
 
 /*
- * DoJumble
- *		Jumble the given Node using the given JumbleState and return the resulting
- *		jumble hash.
+ * Exported initializer for jumble state that allows plugins to hash values and
+ * nodes, but does not record constant locations, for now.
  */
-static uint64
-DoJumble(JumbleState *jstate, Node *node)
+JumbleState *
+InitJumble()
 {
-	/* Jumble the given node */
-	_jumbleNode(jstate, node);
+	return InitJumbleInternal(false);
+}
 
+/*
+ * Produce a 64-bit hash from a jumble state.
+ */
+uint64
+HashJumbleState(JumbleState *jstate)
+{
 	/* Flush any pending NULLs before doing the final hash */
 	if (jstate->pending_nulls > 0)
 		FlushPendingNulls(jstate);
@@ -213,6 +225,20 @@ DoJumble(JumbleState *jstate, Node *node)
 											0));
 }
 
+/*
+ * DoJumble
+ *		Jumble the given Node using the given JumbleState and return the resulting
+ *		jumble hash.
+ */
+static uint64
+DoJumble(JumbleState *jstate, Node *node)
+{
+	/* Jumble the given node */
+	JumbleNode(jstate, node);
+
+	return HashJumbleState(jstate);
+}
+
 /*
  * AppendJumbleInternal: Internal function for appending to the jumble buffer
  *
@@ -281,7 +307,7 @@ AppendJumbleInternal(JumbleState *jstate, const unsigned char *item,
  * AppendJumble
  *		Add 'size' bytes of the given jumble 'value' to the jumble state
  */
-static pg_noinline void
+pg_noinline void
 AppendJumble(JumbleState *jstate, const unsigned char *value, Size size)
 {
 	if (jstate->pending_nulls > 0)
@@ -290,21 +316,11 @@ AppendJumble(JumbleState *jstate, const unsigned char *value, Size size)
 	AppendJumbleInternal(jstate, value, size);
 }
 
-/*
- * AppendJumbleNull
- *		For jumbling NULL pointers
- */
-static pg_attribute_always_inline void
-AppendJumbleNull(JumbleState *jstate)
-{
-	jstate->pending_nulls++;
-}
-
 /*
  * AppendJumble8
  *		Add the first byte from the given 'value' pointer to the jumble state
  */
-static pg_noinline void
+pg_noinline void
 AppendJumble8(JumbleState *jstate, const unsigned char *value)
 {
 	if (jstate->pending_nulls > 0)
@@ -318,7 +334,7 @@ AppendJumble8(JumbleState *jstate, const unsigned char *value)
  *		Add the first 2 bytes from the given 'value' pointer to the jumble
  *		state.
  */
-static pg_noinline void
+pg_noinline void
 AppendJumble16(JumbleState *jstate, const unsigned char *value)
 {
 	if (jstate->pending_nulls > 0)
@@ -332,7 +348,7 @@ AppendJumble16(JumbleState *jstate, const unsigned char *value)
  *		Add the first 4 bytes from the given 'value' pointer to the jumble
  *		state.
  */
-static pg_noinline void
+pg_noinline void
 AppendJumble32(JumbleState *jstate, const unsigned char *value)
 {
 	if (jstate->pending_nulls > 0)
@@ -346,7 +362,7 @@ AppendJumble32(JumbleState *jstate, const unsigned char *value)
  *		Add the first 8 bytes from the given 'value' pointer to the jumble
  *		state.
  */
-static pg_noinline void
+pg_noinline void
 AppendJumble64(JumbleState *jstate, const unsigned char *value)
 {
 	if (jstate->pending_nulls > 0)
@@ -485,31 +501,15 @@ IsSquashableConstList(List *elements, Node **firstExpr, Node **lastExpr)
 }
 
 #define JUMBLE_NODE(item) \
-	_jumbleNode(jstate, (Node *) expr->item)
+	JumbleNode(jstate, (Node *) expr->item)
+#define JUMBLE_FIELD(item) \
+	JUMBLE_VALUE(expr->item)
+#define JUMBLE_STRING(str) \
+	JUMBLE_VALUE_STRING(expr->str)
 #define JUMBLE_ELEMENTS(list) \
 	_jumbleElements(jstate, (List *) expr->list)
 #define JUMBLE_LOCATION(location) \
 	RecordConstLocation(jstate, expr->location, false)
-#define JUMBLE_FIELD(item) \
-do { \
-	if (sizeof(expr->item) == 8) \
-		AppendJumble64(jstate, (const unsigned char *) &(expr->item)); \
-	else if (sizeof(expr->item) == 4) \
-		AppendJumble32(jstate, (const unsigned char *) &(expr->item)); \
-	else if (sizeof(expr->item) == 2) \
-		AppendJumble16(jstate, (const unsigned char *) &(expr->item)); \
-	else if (sizeof(expr->item) == 1) \
-		AppendJumble8(jstate, (const unsigned char *) &(expr->item)); \
-	else \
-		AppendJumble(jstate, (const unsigned char *) &(expr->item), sizeof(expr->item)); \
-} while (0)
-#define JUMBLE_STRING(str) \
-do { \
-	if (expr->str) \
-		AppendJumble(jstate, (const unsigned char *) (expr->str), strlen(expr->str) + 1); \
-	else \
-		AppendJumbleNull(jstate); \
-} while(0)
 /* Function name used for the node field attribute custom_query_jumble. */
 #define JUMBLE_CUSTOM(nodetype, item) \
 	_jumble##nodetype##_##item(jstate, expr, expr->item)
@@ -548,12 +548,12 @@ _jumbleElements(JumbleState *jstate, List *elements)
 	}
 	else
 	{
-		_jumbleNode(jstate, (Node *) elements);
+		JumbleNode(jstate, (Node *) elements);
 	}
 }
 
-static void
-_jumbleNode(JumbleState *jstate, Node *node)
+void
+JumbleNode(JumbleState *jstate, Node *node)
 {
 	Node	   *expr = node;
 #ifdef USE_ASSERT_CHECKING
@@ -627,7 +627,7 @@ _jumbleList(JumbleState *jstate, Node *node)
 	{
 		case T_List:
 			foreach(l, expr)
-				_jumbleNode(jstate, lfirst(l));
+				JumbleNode(jstate, lfirst(l));
 			break;
 		case T_IntList:
 			foreach(l, expr)
diff --git a/src/include/nodes/queryjumble.h b/src/include/nodes/queryjumble.h
index da7c7abed2e..47f49a2f831 100644
--- a/src/include/nodes/queryjumble.h
+++ b/src/include/nodes/queryjumble.h
@@ -81,6 +81,55 @@ enum ComputeQueryIdType
 extern PGDLLIMPORT int compute_query_id;
 
 
+/*
+ * Generic routines for expression jumbling.
+ *
+ * XXX: It may be better to separate these routines in a separate
+ * file.
+ */
+extern JumbleState *InitJumble(void);
+extern void JumbleNode(JumbleState *jstate, Node *node);
+extern uint64 HashJumbleState(JumbleState *jstate);
+
+extern void AppendJumble(JumbleState *jstate,
+						 const unsigned char *item, Size size);
+extern void AppendJumble8(JumbleState *jstate, const unsigned char *value);
+extern void AppendJumble16(JumbleState *jstate, const unsigned char *value);
+extern void AppendJumble32(JumbleState *jstate, const unsigned char *value);
+extern void AppendJumble64(JumbleState *jstate, const unsigned char *value);
+
+/*
+ * AppendJumbleNull
+ *		For jumbling NULL pointers
+ */
+static pg_attribute_always_inline void
+AppendJumbleNull(JumbleState *jstate)
+{
+	jstate->pending_nulls++;
+}
+
+#define JUMBLE_VALUE(val) \
+do { \
+	if (sizeof(val) == 8) \
+		AppendJumble64(jstate, (const unsigned char *) &(val)); \
+	else if (sizeof(val) == 4) \
+		AppendJumble32(jstate, (const unsigned char *) &(val)); \
+	else if (sizeof(val) == 2) \
+		AppendJumble16(jstate, (const unsigned char *) &(val)); \
+	else if (sizeof(val) == 1) \
+		AppendJumble8(jstate, (const unsigned char *) &(val)); \
+	else \
+		AppendJumble(jstate, (const unsigned char *) &(val), sizeof(val)); \
+} while (0)
+#define JUMBLE_VALUE_STRING(str) \
+do { \
+	if (str) \
+		AppendJumble(jstate, (const unsigned char *) (str), strlen(str) + 1); \
+	else \
+		AppendJumbleNull(jstate); \
+} while(0)
+
+/* Query jumbling routines */
 extern const char *CleanQuerytext(const char *query, int *location, int *len);
 extern JumbleState *JumbleQuery(Query *query);
 extern void EnableQueryId(void);
-- 
2.47.1