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

From 5f93ef8dfb2ce56ade2e438ff983a2e162efd811 Mon Sep 17 00:00:00 2001
From: Sami Imseih <simseih@amazon.com>
Date: Mon, 24 Mar 2025 16:07:28 -0500
Subject: [PATCH v6 1/1] 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 | 71 ++++++++++++++++++----------
 src/include/nodes/queryjumble.h      | 11 +++++
 2 files changed, 58 insertions(+), 24 deletions(-)

diff --git a/src/backend/nodes/queryjumblefuncs.c b/src/backend/nodes/queryjumblefuncs.c
index 189bfda610a..99b85e2765c 100644
--- a/src/backend/nodes/queryjumblefuncs.c
+++ b/src/backend/nodes/queryjumblefuncs.c
@@ -62,7 +62,6 @@ static void AppendJumble(JumbleState *jstate,
 						 const unsigned char *item, Size size);
 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);
@@ -120,26 +119,13 @@ CleanQuerytext(const char *query, int *location, int *len)
 JumbleState *
 JumbleQuery(Query *query)
 {
-	JumbleState *jstate = NULL;
+	JumbleState *jstate = InitializeJumbleState(true);
 
 	Assert(IsQueryIdEnabled());
 
-	jstate = (JumbleState *) palloc(sizeof(JumbleState));
-
-	/* 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));
-	jstate->clocations_count = 0;
-	jstate->highest_extern_param_id = 0;
-
 	/* Compute query ID and mark the Query node with it */
-	_jumbleNode(jstate, (Node *) query);
-	query->queryId = DatumGetUInt64(hash_any_extended(jstate->jumble,
-													  jstate->jumble_len,
-													  0));
+	JumbleNode(jstate, (Node *) query);
+	query->queryId = HashJumbleState(jstate);
 
 	/*
 	 * If we are unlucky enough to get a hash of zero, use 1 instead for
@@ -173,7 +159,7 @@ EnableQueryId(void)
  * AppendJumble: Append a value that is substantive in a given query to
  * the current jumble.
  */
-static void
+void
 AppendJumble(JumbleState *jstate, const unsigned char *item, Size size)
 {
 	unsigned char *jumble = jstate->jumble;
@@ -214,9 +200,11 @@ AppendJumble(JumbleState *jstate, const unsigned char *item, Size size)
  * element in a series of merged constants, and everything but the first/last
  * element contributes nothing to the jumble hash.
  */
-static void
+void
 RecordConstLocation(JumbleState *jstate, int location, bool squashed)
 {
+	Assert(jstate->clocations);
+
 	/* -1 indicates unknown or undefined location */
 	if (location >= 0)
 	{
@@ -237,6 +225,41 @@ RecordConstLocation(JumbleState *jstate, int location, bool squashed)
 	}
 }
 
+/*
+ * Initialize a jumble state.
+ */
+JumbleState *
+InitializeJumbleState(bool record_clocations)
+{
+	JumbleState *jstate = (JumbleState *) palloc0(sizeof(JumbleState));
+
+	/* Set up a workspace for expression jumbling */
+	jstate->jumble = (unsigned char *) palloc(JUMBLE_SIZE);
+	jstate->jumble_len = 0;
+
+	if (record_clocations)
+	{
+		jstate->clocations_count = 0;
+		jstate->highest_extern_param_id = 0;
+		jstate->clocations_buf_size = 32;
+		jstate->clocations = (LocationLen *)
+			palloc(jstate->clocations_buf_size * sizeof(LocationLen));
+	}
+
+	return jstate;
+}
+
+/*
+ * Produce a 64-bit hash from a jumble state.
+ */
+uint64
+HashJumbleState(JumbleState *jstate)
+{
+	return DatumGetUInt64(hash_any_extended(jstate->jumble,
+											jstate->jumble_len,
+											0));
+}
+
 /*
  * Subroutine for _jumbleElements: Verify a few simple cases where we can
  * deduce that the expression is a constant:
@@ -319,7 +342,7 @@ IsSquashableConstList(List *elements, Node **firstExpr, Node **lastExpr)
 }
 
 #define JUMBLE_NODE(item) \
-	_jumbleNode(jstate, (Node *) expr->item)
+	JumbleNode(jstate, (Node *) expr->item)
 #define JUMBLE_ELEMENTS(list) \
 	_jumbleElements(jstate, (List *) expr->list)
 #define JUMBLE_LOCATION(location) \
@@ -371,12 +394,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;
 
@@ -441,7 +464,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 905f66bc0bd..620df0d0666 100644
--- a/src/include/nodes/queryjumble.h
+++ b/src/include/nodes/queryjumble.h
@@ -69,7 +69,18 @@ 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 const char *CleanQuerytext(const char *query, int *location, int *len);
+extern JumbleState *InitializeJumbleState(bool record_clocations);
+extern void JumbleNode(JumbleState *jstate, Node *node);
+extern uint64 HashJumbleState(JumbleState *jstate);
+
+/* Query jumbling routines */
 extern JumbleState *JumbleQuery(Query *query);
 extern void EnableQueryId(void);
 
-- 
2.47.1