Split out code into new getKeyJsonValueFromContainer()

The new function stashes its output value in a JsonbValue that can be
passed in by the caller, which enables some of them to pass
stack-allocated structs -- saving palloc cycles.  It also allows some
callers that know they are handling a jsonb object to use this new jsonb
object-specific API, instead of going through generic container
findJsonbValueFromContainer.

Author: Nikita Glukhov
Discussion: https://postgr.es/m/7c417f90-f95f-247e-ba63-d95e39c0ad14@postgrespro.ru

src/backend/utils/adt/jsonb_util.c

@@ -56,6 +56,8 @@ static void appendKey(JsonbParseState *pstate, JsonbValue *scalarVal);
 static void appendValue(JsonbParseState *pstate, JsonbValue *scalarVal);
 static void appendElement(JsonbParseState *pstate, JsonbValue *scalarVal);
 static int	lengthCompareJsonbStringValue(const void *a, const void *b);
+static int	lengthCompareJsonbString(const char *val1, int len1,
+									 const char *val2, int len2);
 static int	lengthCompareJsonbPair(const void *a, const void *b, void *arg);
 static void uniqueifyJsonbObject(JsonbValue *object);
 static JsonbValue *pushJsonbValueScalar(JsonbParseState **pstate,
@@ -329,7 +331,6 @@ findJsonbValueFromContainer(JsonbContainer *container, uint32 flags,
 {
 	JEntry	   *children = container->children;
 	int			count = JsonContainerSize(container);
-	JsonbValue *result;
 
 	Assert((flags & ~(JB_FARRAY | JB_FOBJECT)) == 0);
 
@@ -337,10 +338,9 @@ findJsonbValueFromContainer(JsonbContainer *container, uint32 flags,
 	if (count <= 0)
 		return NULL;
 
-	result = palloc(sizeof(JsonbValue));
-
 	if ((flags & JB_FARRAY) && JsonContainerIsArray(container))
 	{
+		JsonbValue *result = palloc(sizeof(JsonbValue));
 		char	   *base_addr = (char *) (children + count);
 		uint32		offset = 0;
 		int			i;
@@ -357,56 +357,90 @@ findJsonbValueFromContainer(JsonbContainer *container, uint32 flags,
 
 			JBE_ADVANCE_OFFSET(offset, children[i]);
 		}
+
+		pfree(result);
 	}
 	else if ((flags & JB_FOBJECT) && JsonContainerIsObject(container))
 	{
-		/* Since this is an object, account for *Pairs* of Jentrys */
-		char	   *base_addr = (char *) (children + count * 2);
-		uint32		stopLow = 0,
-					stopHigh = count;
-
 		/* Object key passed by caller must be a string */
 		Assert(key->type == jbvString);
 
-		/* Binary search on object/pair keys *only* */
-		while (stopLow < stopHigh)
-		{
-			uint32		stopMiddle;
-			int			difference;
-			JsonbValue	candidate;
+		return getKeyJsonValueFromContainer(container, key->val.string.val,
+											key->val.string.len, NULL);
+	}
+
+	/* Not found */
+	return NULL;
+}
+
+/*
+ * Find value by key in Jsonb object and fetch it into 'res', which is also
+ * returned.
+ *
+ * 'res' can be passed in as NULL, in which case it's newly palloc'ed here.
+ */
+JsonbValue *
+getKeyJsonValueFromContainer(JsonbContainer *container,
+							 const char *keyVal, int keyLen, JsonbValue *res)
+{
+	JEntry	   *children = container->children;
+	int			count = JsonContainerSize(container);
+	char	   *baseAddr;
+	uint32		stopLow,
+				stopHigh;
 
-			stopMiddle = stopLow + (stopHigh - stopLow) / 2;
+	Assert(JsonContainerIsObject(container));
 
-			candidate.type = jbvString;
-			candidate.val.string.val =
-				base_addr + getJsonbOffset(container, stopMiddle);
-			candidate.val.string.len = getJsonbLength(container, stopMiddle);
+	/* Quick out without a palloc cycle if object is empty */
+	if (count <= 0)
+		return NULL;
 
-			difference = lengthCompareJsonbStringValue(&candidate, key);
+	/*
+	 * Binary search the container. Since we know this is an object, account
+	 * for *Pairs* of Jentrys
+	 */
+	baseAddr = (char *) (children + count * 2);
+	stopLow = 0;
+	stopHigh = count;
+	while (stopLow < stopHigh)
+	{
+		uint32		stopMiddle;
+		int			difference;
+		const char *candidateVal;
+		int			candidateLen;
 
-			if (difference == 0)
-			{
-				/* Found our key, return corresponding value */
-				int			index = stopMiddle + count;
+		stopMiddle = stopLow + (stopHigh - stopLow) / 2;
 
-				fillJsonbValue(container, index, base_addr,
-							   getJsonbOffset(container, index),
-							   result);
+		candidateVal = baseAddr + getJsonbOffset(container, stopMiddle);
+		candidateLen = getJsonbLength(container, stopMiddle);
 
-				return result;
-			}
+		difference = lengthCompareJsonbString(candidateVal, candidateLen,
+											  keyVal, keyLen);
+
+		if (difference == 0)
+		{
+			/* Found our key, return corresponding value */
+			int			index = stopMiddle + count;
+
+			if (!res)
+				res = palloc(sizeof(JsonbValue));
+
+			fillJsonbValue(container, index, baseAddr,
+						   getJsonbOffset(container, index),
+						   res);
+
+			return res;
+		}
+		else
+		{
+			if (difference < 0)
+				stopLow = stopMiddle + 1;
 			else
-			{
-				if (difference < 0)
-					stopLow = stopMiddle + 1;
-				else
-					stopHigh = stopMiddle;
-			}
+				stopHigh = stopMiddle;
 		}
 	}
 
 	/* Not found */
-	pfree(result);
 	return NULL;
 }
 
@@ -1009,6 +1043,7 @@ JsonbDeepContains(JsonbIterator **val, JsonbIterator **mContained)
 		for (;;)
 		{
 			JsonbValue *lhsVal; /* lhsVal is from pair in lhs object */
+			JsonbValue	lhsValBuf;
 
 			rcont = JsonbIteratorNext(mContained, &vcontained, false);
 
@@ -1021,12 +1056,14 @@ JsonbDeepContains(JsonbIterator **val, JsonbIterator **mContained)
 				return true;
 
 			Assert(rcont == WJB_KEY);
+			Assert(vcontained.type == jbvString);
 
 			/* First, find value by key... */
-			lhsVal = findJsonbValueFromContainer((*val)->container,
-												 JB_FOBJECT,
-												 &vcontained);
-
+			lhsVal =
+				getKeyJsonValueFromContainer((*val)->container,
+											 vcontained.val.string.val,
+											 vcontained.val.string.len,
+											 &lhsValBuf);
 			if (!lhsVal)
 				return false;
 
@@ -1771,21 +1808,27 @@ lengthCompareJsonbStringValue(const void *a, const void *b)
 {
 	const JsonbValue *va = (const JsonbValue *) a;
 	const JsonbValue *vb = (const JsonbValue *) b;
-	int			res;
 
 	Assert(va->type == jbvString);
 	Assert(vb->type == jbvString);
 
-	if (va->val.string.len == vb->val.string.len)
-	{
-		res = memcmp(va->val.string.val, vb->val.string.val, va->val.string.len);
-	}
-	else
-	{
-		res = (va->val.string.len > vb->val.string.len) ? 1 : -1;
-	}
+	return lengthCompareJsonbString(va->val.string.val, va->val.string.len,
+									vb->val.string.val, vb->val.string.len);
+}
 
-	return res;
+/*
+ * Subroutine for lengthCompareJsonbStringValue
+ *
+ * This is also useful separately to implement binary search on
+ * JsonbContainers.
+ */
+static int
+lengthCompareJsonbString(const char *val1, int len1, const char *val2, int len2)
+{
+	if (len1 == len2)
+		return memcmp(val1, val2, len1);
+	else
+		return len1 > len2 ? 1 : -1;
 }
 
 /*

src/backend/utils/adt/jsonfuncs.c

@@ -454,12 +454,6 @@ static Datum populate_array(ArrayIOData *aio, const char *colname,
 static Datum populate_domain(DomainIOData *io, Oid typid, const char *colname,
 							 MemoryContext mcxt, JsValue *jsv, bool isnull);
 
-/* Worker that takes care of common setup for us */
-static JsonbValue *findJsonbValueFromContainerLen(JsonbContainer *container,
-												  uint32 flags,
-												  char *key,
-												  uint32 keylen);
-
 /* functions supporting jsonb_delete, jsonb_set and jsonb_concat */
 static JsonbValue *IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
 								  JsonbParseState **state);
@@ -718,13 +712,15 @@ jsonb_object_field(PG_FUNCTION_ARGS)
 	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
 	text	   *key = PG_GETARG_TEXT_PP(1);
 	JsonbValue *v;
+	JsonbValue	vbuf;
 
 	if (!JB_ROOT_IS_OBJECT(jb))
 		PG_RETURN_NULL();
 
-	v = findJsonbValueFromContainerLen(&jb->root, JB_FOBJECT,
-									   VARDATA_ANY(key),
-									   VARSIZE_ANY_EXHDR(key));
+	v = getKeyJsonValueFromContainer(&jb->root,
+									 VARDATA_ANY(key),
+									 VARSIZE_ANY_EXHDR(key),
+									 &vbuf);
 
 	if (v != NULL)
 		PG_RETURN_JSONB_P(JsonbValueToJsonb(v));
@@ -754,14 +750,15 @@ jsonb_object_field_text(PG_FUNCTION_ARGS)
 	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
 	text	   *key = PG_GETARG_TEXT_PP(1);
 	JsonbValue *v;
+	JsonbValue	vbuf;
 
 	if (!JB_ROOT_IS_OBJECT(jb))
 		PG_RETURN_NULL();
 
-	v = findJsonbValueFromContainerLen(&jb->root, JB_FOBJECT,
-									   VARDATA_ANY(key),
-									   VARSIZE_ANY_EXHDR(key));
-
+	v = getKeyJsonValueFromContainer(&jb->root,
+									 VARDATA_ANY(key),
+									 VARSIZE_ANY_EXHDR(key),
+									 &vbuf);
 
 	if (v != NULL && v->type != jbvNull)
 		PG_RETURN_TEXT_P(JsonbValueAsText(v));
@@ -1336,6 +1333,7 @@ get_jsonb_path_all(FunctionCallInfo fcinfo, bool as_text)
 	bool		have_object = false,
 				have_array = false;
 	JsonbValue *jbvp = NULL;
+	JsonbValue	jbvbuf;
 	JsonbContainer *container;
 
 	/*
@@ -1393,10 +1391,10 @@ get_jsonb_path_all(FunctionCallInfo fcinfo, bool as_text)
 	{
 		if (have_object)
 		{
-			jbvp = findJsonbValueFromContainerLen(container,
-												  JB_FOBJECT,
-												  VARDATA(pathtext[i]),
-												  VARSIZE(pathtext[i]) - VARHDRSZ);
+			jbvp = getKeyJsonValueFromContainer(container,
+												VARDATA(pathtext[i]),
+												VARSIZE(pathtext[i]) - VARHDRSZ,
+												&jbvbuf);
 		}
 		else if (have_array)
 		{
@@ -3023,8 +3021,8 @@ JsObjectGetField(JsObject *obj, char *field, JsValue *jsv)
 	else
 	{
 		jsv->val.jsonb = !obj->val.jsonb_cont ? NULL :
-			findJsonbValueFromContainerLen(obj->val.jsonb_cont, JB_FOBJECT,
-										   field, strlen(field));
+			getKeyJsonValueFromContainer(obj->val.jsonb_cont, field, strlen(field),
+										 NULL);
 
 		return jsv->val.jsonb != NULL;
 	}
@@ -3848,22 +3846,6 @@ populate_recordset_object_field_end(void *state, char *fname, bool isnull)
 	}
 }
 
-/*
- * findJsonbValueFromContainer() wrapper that sets up JsonbValue key string.
- */
-static JsonbValue *
-findJsonbValueFromContainerLen(JsonbContainer *container, uint32 flags,
-							   char *key, uint32 keylen)
-{
-	JsonbValue	k;
-
-	k.type = jbvString;
-	k.val.string.val = key;
-	k.val.string.len = keylen;
-
-	return findJsonbValueFromContainer(container, flags, &k);
-}
-
 /*
  * Semantic actions for json_strip_nulls.
  *

src/include/utils/jsonb.h

@@ -364,6 +364,9 @@ extern int	compareJsonbContainers(JsonbContainer *a, JsonbContainer *b);
 extern JsonbValue *findJsonbValueFromContainer(JsonbContainer *sheader,
 											   uint32 flags,
 											   JsonbValue *key);
+extern JsonbValue *getKeyJsonValueFromContainer(JsonbContainer *container,
+												const char *keyVal, int keyLen,
+												JsonbValue *res);
 extern JsonbValue *getIthJsonbValueFromContainer(JsonbContainer *sheader,
 												 uint32 i);
 extern JsonbValue *pushJsonbValue(JsonbParseState **pstate,