From be9ff3d0cd8f7ed2c2f1fa479db14a24c12fb3c2 Mon Sep 17 00:00:00 2001
From: Nikita Glukhov <n.gluhov@postgrespro.ru>
Date: Thu, 21 Feb 2019 19:17:38 +0300
Subject: [PATCH v3 4/5] Extract findJsonbKeyInObject()

---
 src/backend/utils/adt/jsonb_util.c | 147 +++++++++++++++++------------
 src/backend/utils/adt/jsonfuncs.c  |  51 ++++------
 src/include/utils/jsonb.h          |   3 +
 3 files changed, 108 insertions(+), 93 deletions(-)

diff --git a/src/backend/utils/adt/jsonb_util.c b/src/backend/utils/adt/jsonb_util.c
index ac04c4a57b..afde5148aa 100644
--- a/src/backend/utils/adt/jsonb_util.c
+++ b/src/backend/utils/adt/jsonb_util.c
@@ -57,6 +57,8 @@ 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	lengthCompareJsonbPair(const void *a, const void *b, void *arg);
+static int lengthCompareJsonbString(const char *val1, int len1,
+						 const char *val2, int len2);
 static void uniqueifyJsonbObject(JsonbValue *object);
 static JsonbValue *pushJsonbValueScalar(JsonbParseState **pstate,
 										JsonbIteratorToken seq,
@@ -297,6 +299,71 @@ compareJsonbContainers(JsonbContainer *a, JsonbContainer *b)
 	return res;
 }
 
+/* Find value by key in Jsonb object and fetch it into 'res'. */
+JsonbValue *
+findJsonbKeyInObject(JsonbContainer *container, const char *keyVal, int keyLen,
+					 JsonbValue *res)
+{
+	JEntry	   *children = container->children;
+	JsonbValue *result = res;
+	int			count = JsonContainerSize(container);
+	/* Since this is an object, account for *Pairs* of Jentrys */
+	char	   *baseAddr = (char *) (children + count * 2);
+	uint32		stopLow = 0,
+				stopHigh = count;
+
+	Assert(JsonContainerIsObject(container));
+
+	/* Quick out without a palloc cycle if object is empty */
+	if (count <= 0)
+		return NULL;
+
+	if (!result)
+		result = palloc(sizeof(JsonbValue));
+
+	/* Binary search on object/pair keys *only* */
+	while (stopLow < stopHigh)
+	{
+		uint32		stopMiddle;
+		int			difference;
+		const char *candidateVal;
+		int			candidateLen;
+
+		stopMiddle = stopLow + (stopHigh - stopLow) / 2;
+
+		candidateVal = baseAddr + getJsonbOffset(container, stopMiddle);
+		candidateLen = getJsonbLength(container, stopMiddle);
+
+		difference = lengthCompareJsonbString(candidateVal, candidateLen,
+											  keyVal, keyLen);
+
+		if (difference == 0)
+		{
+			/* Found our key, return corresponding value */
+			int			index = stopMiddle + count;
+
+			fillJsonbValue(container, index, baseAddr,
+						   getJsonbOffset(container, index),
+						   result);
+
+			return result;
+		}
+		else
+		{
+			if (difference < 0)
+				stopLow = stopMiddle + 1;
+			else
+				stopHigh = stopMiddle;
+		}
+	}
+
+	/* Not found */
+	if (!res)
+		pfree(result);
+
+	return NULL;
+}
+
 /*
  * Find value in object (i.e. the "value" part of some key/value pair in an
  * object), or find a matching element if we're looking through an array.  Do
@@ -329,7 +396,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 +403,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 +422,19 @@ 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;
-
-			stopMiddle = stopLow + (stopHigh - stopLow) / 2;
-
-			candidate.type = jbvString;
-			candidate.val.string.val =
-				base_addr + getJsonbOffset(container, stopMiddle);
-			candidate.val.string.len = getJsonbLength(container, stopMiddle);
-
-			difference = lengthCompareJsonbStringValue(&candidate, key);
-
-			if (difference == 0)
-			{
-				/* Found our key, return corresponding value */
-				int			index = stopMiddle + count;
-
-				fillJsonbValue(container, index, base_addr,
-							   getJsonbOffset(container, index),
-							   result);
-
-				return result;
-			}
-			else
-			{
-				if (difference < 0)
-					stopLow = stopMiddle + 1;
-				else
-					stopHigh = stopMiddle;
-			}
-		}
+		return findJsonbKeyInObject(container, key->val.string.val,
+									key->val.string.len, NULL);
 	}
 
 	/* Not found */
-	pfree(result);
 	return NULL;
 }
 
@@ -1009,6 +1037,7 @@ JsonbDeepContains(JsonbIterator **val, JsonbIterator **mContained)
 		for (;;)
 		{
 			JsonbValue *lhsVal; /* lhsVal is from pair in lhs object */
+			JsonbValue	lhsValBuf;
 
 			rcont = JsonbIteratorNext(mContained, &vcontained, false);
 
@@ -1021,11 +1050,13 @@ 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 = findJsonbKeyInObject((*val)->container,
+										  vcontained.val.string.val,
+										  vcontained.val.string.len,
+										  &lhsValBuf);
 
 			if (!lhsVal)
 				return false;
@@ -1754,6 +1785,15 @@ convertJsonbScalar(StringInfo buffer, JEntry *jentry, JsonbValue *scalarVal)
 	}
 }
 
+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;
+}
+
 /*
  * Compare two jbvString JsonbValue values, a and b.
  *
@@ -1771,21 +1811,12 @@ 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 res;
+	return lengthCompareJsonbString(va->val.string.val, va->val.string.len,
+									vb->val.string.val, vb->val.string.len);
 }
 
 /*
diff --git a/src/backend/utils/adt/jsonfuncs.c b/src/backend/utils/adt/jsonfuncs.c
index 753d691d75..3732fb5e6b 100644
--- a/src/backend/utils/adt/jsonfuncs.c
+++ b/src/backend/utils/adt/jsonfuncs.c
@@ -453,12 +453,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);
@@ -717,13 +711,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 = findJsonbKeyInObject(&jb->root,
+							 VARDATA_ANY(key),
+							 VARSIZE_ANY_EXHDR(key),
+							 &vbuf);
 
 	if (v != NULL)
 		PG_RETURN_JSONB_P(JsonbValueToJsonb(v));
@@ -794,14 +790,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 = findJsonbKeyInObject(&jb->root,
+							 VARDATA_ANY(key),
+							 VARSIZE_ANY_EXHDR(key),
+							 &vbuf);
 
 	if (v != NULL && v->type != jbvNull)
 		PG_RETURN_TEXT_P(JsonbValueAsText(v));
@@ -1433,10 +1430,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 = findJsonbKeyInObject(container,
+										VARDATA(pathtext[i]),
+										VARSIZE(pathtext[i]) - VARHDRSZ,
+										&jbvbuf);
 		}
 		else if (have_array)
 		{
@@ -3019,8 +3016,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));
+			findJsonbKeyInObject(obj->val.jsonb_cont, field, strlen(field),
+								 NULL);
 
 		return jsv->val.jsonb != NULL;
 	}
@@ -3844,22 +3841,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.
  *
diff --git a/src/include/utils/jsonb.h b/src/include/utils/jsonb.h
index ac52b75f51..bddf272b53 100644
--- a/src/include/utils/jsonb.h
+++ b/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 *findJsonbKeyInObject(JsonbContainer *container,
+				  						const char *keyVal, int keyLen,
+										JsonbValue *res);
 extern JsonbValue *getIthJsonbValueFromContainer(JsonbContainer *sheader,
 												 uint32 i);
 extern JsonbValue *pushJsonbValue(JsonbParseState **pstate,
-- 
2.17.1

