From e590cd2be54923aa608d0e575f272fde84b7a193 Mon Sep 17 00:00:00 2001
From: Lukas Fittl <lukas@fittl.com>
Date: Tue, 31 Dec 2024 15:05:39 -0800
Subject: [PATCH v4 1/3] Allow using jumbling logic outside of query jumble
 unit file

This can be useful either for jumbling expressions in other contexts
(e.g. to calculate a plan jumble), or to allow extensions to use
a modified jumbling logic more easily.

This intentionally supports the use case where a separate jumbling logic
does not care about recording constants, as the query jumble does.
---
 src/backend/nodes/queryjumblefuncs.c | 59 +++++++++++++++++-----------
 src/include/nodes/queryjumble.h      |  6 +++
 2 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/src/backend/nodes/queryjumblefuncs.c b/src/backend/nodes/queryjumblefuncs.c
index b103a281936..545d8edcae2 100644
--- a/src/backend/nodes/queryjumblefuncs.c
+++ b/src/backend/nodes/queryjumblefuncs.c
@@ -51,10 +51,7 @@ int			compute_query_id = COMPUTE_QUERY_ID_AUTO;
  */
 bool		query_id_enabled = false;
 
-static void AppendJumble(JumbleState *jstate,
-						 const unsigned char *item, Size size);
 static void RecordConstLocation(JumbleState *jstate, int location);
-static void _jumbleNode(JumbleState *jstate, Node *node);
 static void _jumbleA_Const(JumbleState *jstate, Node *node);
 static void _jumbleList(JumbleState *jstate, Node *node);
 static void _jumbleVariableSetStmt(JumbleState *jstate, Node *node);
@@ -109,28 +106,42 @@ CleanQuerytext(const char *query, int *location, int *len)
 }
 
 JumbleState *
-JumbleQuery(Query *query)
+InitializeJumbleState(bool record_clocations)
 {
-	JumbleState *jstate = NULL;
-
-	Assert(IsQueryIdEnabled());
-
-	jstate = (JumbleState *) palloc(sizeof(JumbleState));
+	JumbleState *jstate = (JumbleState *) palloc0(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;
+
+	if (record_clocations)
+	{
+		jstate->clocations_buf_size = 32;
+		jstate->clocations = (LocationLen *)
+			palloc(jstate->clocations_buf_size * sizeof(LocationLen));
+	}
+
+	return jstate;
+}
+
+uint64
+HashJumbleState(JumbleState *jstate)
+{
+	return DatumGetUInt64(hash_any_extended(jstate->jumble,
+											jstate->jumble_len,
+											0));
+}
+
+JumbleState *
+JumbleQuery(Query *query)
+{
+	JumbleState *jstate = InitializeJumbleState(true);
+
+	Assert(IsQueryIdEnabled());
 
 	/* 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
@@ -164,7 +175,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;
@@ -205,7 +216,7 @@ static void
 RecordConstLocation(JumbleState *jstate, int location)
 {
 	/* -1 indicates unknown or undefined location */
-	if (location >= 0)
+	if (location >= 0 && jstate->clocations_buf_size > 0)
 	{
 		/* enlarge array if needed */
 		if (jstate->clocations_count >= jstate->clocations_buf_size)
@@ -224,7 +235,7 @@ RecordConstLocation(JumbleState *jstate, int location)
 }
 
 #define JUMBLE_NODE(item) \
-	_jumbleNode(jstate, (Node *) expr->item)
+	JumbleNode(jstate, (Node *) expr->item)
 #define JUMBLE_LOCATION(location) \
 	RecordConstLocation(jstate, expr->location)
 #define JUMBLE_FIELD(item) \
@@ -239,8 +250,8 @@ do { \
 
 #include "queryjumblefuncs.funcs.c"
 
-static void
-_jumbleNode(JumbleState *jstate, Node *node)
+void
+JumbleNode(JumbleState *jstate, Node *node)
 {
 	Node	   *expr = node;
 
@@ -305,7 +316,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 50eb9566587..5afa6f3605f 100644
--- a/src/include/nodes/queryjumble.h
+++ b/src/include/nodes/queryjumble.h
@@ -83,4 +83,10 @@ IsQueryIdEnabled(void)
 	return query_id_enabled;
 }
 
+/* Functions intended for other users of jumbling (e.g. plan jumbling) */
+extern JumbleState *InitializeJumbleState(bool record_clocations);
+extern void AppendJumble(JumbleState *jstate, const unsigned char *item, Size size);
+extern void JumbleNode(JumbleState *jstate, Node *node);
+extern uint64 HashJumbleState(JumbleState *jstate);
+
 #endif							/* QUERYJUMBLE_H */
-- 
2.47.1