From 81df5424a9c69045f85da6782ae7485c0f66d03e Mon Sep 17 00:00:00 2001
From: mahendra <mahendra@localhost.localdomain>
Date: Tue, 5 Apr 2022 07:48:56 +0300
Subject: [PATCH 4/6] Add helper jsonb functions and
 macros

---
 src/backend/utils/adt/jsonb_util.c |  27 ++++++
 src/backend/utils/adt/jsonfuncs.c  |  10 ++-
 src/include/utils/jsonb.h          | 165 +++++++++++++++++++++++++++++++++++--
 3 files changed, 194 insertions(+), 8 deletions(-)

diff --git a/src/backend/utils/adt/jsonb_util.c b/src/backend/utils/adt/jsonb_util.c
index aa151a5..5a7831d 100644
--- a/src/backend/utils/adt/jsonb_util.c
+++ b/src/backend/utils/adt/jsonb_util.c
@@ -391,6 +391,22 @@ findJsonbValueFromContainer(JsonbContainer *container, uint32 flags,
 }
 
 /*
+ * findJsonbValueFromContainer() wrapper that sets up JsonbValue key string.
+ */
+JsonbValue *
+findJsonbValueFromContainerLen(JsonbContainer *container, uint32 flags,
+							   const char *key, uint32 keylen)
+{
+	JsonbValue	k;
+
+	k.type = jbvString;
+	k.val.string.val = key;
+	k.val.string.len = keylen;
+
+	return findJsonbValueFromContainer(container, flags, &k);
+}
+
+/*
  * Find value by key in Jsonb object and fetch it into 'res', which is also
  * returned.
  *
@@ -604,6 +620,17 @@ pushJsonbValue(JsonbParseState **pstate, JsonbIteratorToken seq,
 		return pushJsonbValueScalar(pstate, seq, jbval);
 	}
 
+	/*
+	 * XXX I'm not quite sure why we actually do this? Why do we need to change
+	 * how JsonbValue is converted to Jsonb for the statistics patch?
+	 */
+	/* push value from scalar container without its enclosing array */
+	if (*pstate && JsonbExtractScalar(jbval->val.binary.data, &v))
+	{
+		Assert(IsAJsonbScalar(&v));
+		return pushJsonbValueScalar(pstate, seq, &v);
+	}
+
 	/* unpack the binary and add each piece to the pstate */
 	it = JsonbIteratorInit(jbval->val.binary.data);
 
diff --git a/src/backend/utils/adt/jsonfuncs.c b/src/backend/utils/adt/jsonfuncs.c
index a682d9c..ed5dd17 100644
--- a/src/backend/utils/adt/jsonfuncs.c
+++ b/src/backend/utils/adt/jsonfuncs.c
@@ -5305,7 +5305,8 @@ iterate_jsonb_values(Jsonb *jb, uint32 flags, void *state,
 		if (type == WJB_KEY)
 		{
 			if (flags & jtiKey)
-				action(state, v.val.string.val, v.val.string.len);
+				action(state, unconstify(char *, v.val.string.val),
+					   v.val.string.len);
 
 			continue;
 		}
@@ -5320,7 +5321,8 @@ iterate_jsonb_values(Jsonb *jb, uint32 flags, void *state,
 		{
 			case jbvString:
 				if (flags & jtiString)
-					action(state, v.val.string.val, v.val.string.len);
+					action(state, unconstify(char *, v.val.string.val),
+						   v.val.string.len);
 				break;
 			case jbvNumeric:
 				if (flags & jtiNumeric)
@@ -5442,7 +5444,9 @@ transform_jsonb_string_values(Jsonb *jsonb, void *action_state,
 	{
 		if ((type == WJB_VALUE || type == WJB_ELEM) && v.type == jbvString)
 		{
-			out = transform_action(action_state, v.val.string.val, v.val.string.len);
+			out = transform_action(action_state,
+								   unconstify(char *, v.val.string.val),
+								   v.val.string.len);
 			v.val.string.val = VARDATA_ANY(out);
 			v.val.string.len = VARSIZE_ANY_EXHDR(out);
 			res = pushJsonbValue(&st, type, type < WJB_BEGIN_ARRAY ? &v : NULL);
diff --git a/src/include/utils/jsonb.h b/src/include/utils/jsonb.h
index bae466b..158c9ae 100644
--- a/src/include/utils/jsonb.h
+++ b/src/include/utils/jsonb.h
@@ -14,6 +14,7 @@
 
 #include "lib/stringinfo.h"
 #include "utils/array.h"
+#include "utils/builtins.h"
 #include "utils/numeric.h"
 
 /* Tokens used when sequentially processing a jsonb value */
@@ -229,8 +230,7 @@ typedef struct
 #define JB_ROOT_IS_OBJECT(jbp_) ((*(uint32 *) VARDATA(jbp_) & JB_FOBJECT) != 0)
 #define JB_ROOT_IS_ARRAY(jbp_)	((*(uint32 *) VARDATA(jbp_) & JB_FARRAY) != 0)
 
-
-enum jbvType
+typedef enum jbvType
 {
 	/* Scalar types */
 	jbvNull = 0x0,
@@ -250,7 +250,7 @@ enum jbvType
 	 * into JSON strings when outputted to json/jsonb.
 	 */
 	jbvDatetime = 0x20,
-};
+} JsonbValueType;
 
 /*
  * JsonbValue:	In-memory representation of Jsonb.  This is a convenient
@@ -269,7 +269,7 @@ struct JsonbValue
 		struct
 		{
 			int			len;
-			char	   *val;	/* Not necessarily null-terminated */
+			const char *val;	/* Not necessarily null-terminated */
 		}			string;		/* String primitive type */
 
 		struct
@@ -400,6 +400,10 @@ extern int	compareJsonbContainers(JsonbContainer *a, JsonbContainer *b);
 extern JsonbValue *findJsonbValueFromContainer(JsonbContainer *sheader,
 											   uint32 flags,
 											   JsonbValue *key);
+extern JsonbValue *findJsonbValueFromContainerLen(JsonbContainer *container,
+												  uint32 flags,
+												  const char *key,
+												  uint32 keylen);
 extern JsonbValue *getKeyJsonValueFromContainer(JsonbContainer *container,
 												const char *keyVal, int keyLen,
 												JsonbValue *res);
@@ -444,5 +448,156 @@ extern Datum jsonb_build_object_worker(int nargs, Datum *args, bool *nulls,
 									   bool unique_keys);
 extern Datum jsonb_build_array_worker(int nargs, Datum *args, bool *nulls,
 									  Oid *types, bool absent_on_null);
+/*
+ *  * XXX Not sure we want to add these functions to jsonb.h, which is the
+ *   * public API. Maybe it belongs rather to jsonb_typanalyze.c or elsewhere,
+ *    * closer to how it's used?
+ *     */
+
+/* helper inline functions for JsonbValue initialization */
+static inline JsonbValue *
+JsonValueInitObject(JsonbValue *val, int nPairs, int nPairsAllocated)
+{
+	val->type = jbvObject;
+	val->val.object.nPairs = nPairs;
+	val->val.object.pairs = nPairsAllocated ?
+							palloc(sizeof(JsonbPair) * nPairsAllocated) : NULL;
+
+	return val;
+}
+
+static inline JsonbValue *
+JsonValueInitArray(JsonbValue *val, int nElems, int nElemsAllocated,
+				   bool rawScalar)
+{
+	val->type = jbvArray;
+	val->val.array.nElems = nElems;
+	val->val.array.elems = nElemsAllocated ?
+							palloc(sizeof(JsonbValue) * nElemsAllocated) : NULL;
+	val->val.array.rawScalar = rawScalar;
+
+	return val;
+}
+
+static inline JsonbValue *
+JsonValueInitBinary(JsonbValue *val, Jsonb *jb)
+{
+	val->type = jbvBinary;
+	val->val.binary.data = &(jb)->root;
+	val->val.binary.len = VARSIZE_ANY_EXHDR(jb);
+	return val;
+}
+
+
+static inline JsonbValue *
+JsonValueInitString(JsonbValue *jbv, const char *str)
+{
+	jbv->type = jbvString;
+	jbv->val.string.len = strlen(str);
+	jbv->val.string.val = memcpy(palloc(jbv->val.string.len + 1), str,
+								 jbv->val.string.len + 1);
+	return jbv;
+}
+
+static inline JsonbValue *
+JsonValueInitStringWithLen(JsonbValue *jbv, const char *str, int len)
+{
+	jbv->type = jbvString;
+	jbv->val.string.val = str;
+	jbv->val.string.len = len;
+	return jbv;
+}
+
+static inline JsonbValue *
+JsonValueInitText(JsonbValue *jbv, text *txt)
+{
+	jbv->type = jbvString;
+	jbv->val.string.val = VARDATA_ANY(txt);
+	jbv->val.string.len = VARSIZE_ANY_EXHDR(txt);
+	return jbv;
+}
+
+static inline JsonbValue *
+JsonValueInitNumeric(JsonbValue *jbv, Numeric num)
+{
+	jbv->type = jbvNumeric;
+	jbv->val.numeric = num;
+	return jbv;
+}
+
+static inline JsonbValue *
+JsonValueInitInteger(JsonbValue *jbv, int64 i)
+{
+	jbv->type = jbvNumeric;
+	jbv->val.numeric = DatumGetNumeric(DirectFunctionCall1(
+											int8_numeric, Int64GetDatum(i)));
+	return jbv;
+}
+
+static inline JsonbValue *
+JsonValueInitFloat(JsonbValue *jbv, float4 f)
+{
+	jbv->type = jbvNumeric;
+	jbv->val.numeric = DatumGetNumeric(DirectFunctionCall1(
+											float4_numeric, Float4GetDatum(f)));
+	return jbv;
+}
+
+static inline JsonbValue *
+JsonValueInitDouble(JsonbValue *jbv, float8 f)
+{
+	jbv->type = jbvNumeric;
+	jbv->val.numeric = DatumGetNumeric(DirectFunctionCall1(
+											float8_numeric, Float8GetDatum(f)));
+	return jbv;
+}
+
+/* helper macros for jsonb building */
+#define pushJsonbKey(pstate, jbv, key) \
+		pushJsonbValue(pstate, WJB_KEY, JsonValueInitString(jbv, key))
+
+#define pushJsonbValueGeneric(Type, pstate, jbv, val) \
+		pushJsonbValue(pstate, WJB_VALUE, JsonValueInit##Type(jbv, val))
+
+#define pushJsonbElemGeneric(Type, pstate, jbv, val) \
+		pushJsonbValue(pstate, WJB_ELEM, JsonValueInit##Type(jbv, val))
+
+#define pushJsonbValueInteger(pstate, jbv, i) \
+		pushJsonbValueGeneric(Integer, pstate, jbv, i)
+
+#define pushJsonbValueFloat(pstate, jbv, f) \
+		pushJsonbValueGeneric(Float, pstate, jbv, f)
+
+#define pushJsonbElemFloat(pstate, jbv, f) \
+		pushJsonbElemGeneric(Float, pstate, jbv, f)
+
+#define pushJsonbElemString(pstate, jbv, txt) \
+		pushJsonbElemGeneric(String, pstate, jbv, txt)
+
+#define pushJsonbElemText(pstate, jbv, txt) \
+		pushJsonbElemGeneric(Text, pstate, jbv, txt)
+
+#define pushJsonbElemNumeric(pstate, jbv, num) \
+		pushJsonbElemGeneric(Numeric, pstate, jbv, num)
+
+#define pushJsonbElemInteger(pstate, jbv, num) \
+		pushJsonbElemGeneric(Integer, pstate, jbv, num)
+
+#define pushJsonbElemBinary(pstate, jbv, jbcont) \
+		pushJsonbElemGeneric(Binary, pstate, jbv, jbcont)
+
+#define pushJsonbKeyValueGeneric(Type, pstate, jbv, key, val) ( \
+		pushJsonbKey(pstate, jbv, key), \
+		pushJsonbValueGeneric(Type, pstate, jbv, val) \
+	)
+
+#define pushJsonbKeyValueString(pstate, jbv, key, val) \
+		pushJsonbKeyValueGeneric(String, pstate, jbv, key, val)
+
+#define pushJsonbKeyValueFloat(pstate, jbv, key, val) \
+		pushJsonbKeyValueGeneric(Float, pstate, jbv, key, val)
+
+#define pushJsonbKeyValueInteger(pstate, jbv, key, val) \
+		pushJsonbKeyValueGeneric(Integer, pstate, jbv, key, val)
 
-#endif							/* __JSONB_H__ */
+#endif                                                 /* __JSONB_H__ */
-- 
1.8.3.1

