Revise APIs for pushJsonbValue() and associated routines.

Instead of passing "JsonbParseState **" to pushJsonbValue(),
pass a pointer to a JsonbInState, which will contain the
parseState stack pointer as well as other useful fields.
Also, instead of returning a JsonbValue pointer that is often
meaningless/ignored, return the top-level JsonbValue pointer
in the "result" field of the JsonbInState.

This involves a lot of (mostly mechanical) edits, but I think
the results are notationally cleaner and easier to understand.
Certainly the business with sometimes capturing the result of
pushJsonbValue() and sometimes not was bug-prone and incapable of
mechanical verification.  In the new arrangement, JsonbInState.result
remains null until we've completed a valid sequence of pushes, so
that an incorrect sequence will result in a null-pointer dereference,
not mistaken use of a partial result.

However, this isn't simply an exercise in prettier notation.
The real reason for doing it is to provide a mechanism whereby
pushJsonbValue() can be told to construct the JsonbValue tree
in a context that is not CurrentMemoryContext.  That happens
when a non-null "outcontext" is specified in the JsonbInState.
No callers exercise that option in this patch, but the next
patch in the series will make use of it.

I tried to improve the comments in this area too.

Author: Tom Lane <tgl@sss.pgh.pa.us>
Reviewed-by: jian he <jian.universality@gmail.com>
Reviewed-by: Chao Li <li.evan.chao@gmail.com>
Discussion: https://postgr.es/m/1060917.1753202222@sss.pgh.pa.us

contrib/hstore/hstore_io.c

@@ -1439,10 +1439,9 @@ hstore_to_jsonb(PG_FUNCTION_ARGS)
 	int			count = HS_COUNT(in);
 	char	   *base = STRPTR(in);
 	HEntry	   *entries = ARRPTR(in);
-	JsonbParseState *state = NULL;
-	JsonbValue *res;
+	JsonbInState state = {0};
 
-	(void) pushJsonbValue(&state, WJB_BEGIN_OBJECT, NULL);
+	pushJsonbValue(&state, WJB_BEGIN_OBJECT, NULL);
 
 	for (i = 0; i < count; i++)
 	{
@@ -1453,7 +1452,7 @@ hstore_to_jsonb(PG_FUNCTION_ARGS)
 		key.val.string.len = HSTORE_KEYLEN(entries, i);
 		key.val.string.val = HSTORE_KEY(entries, base, i);
 
-		(void) pushJsonbValue(&state, WJB_KEY, &key);
+		pushJsonbValue(&state, WJB_KEY, &key);
 
 		if (HSTORE_VALISNULL(entries, i))
 		{
@@ -1465,12 +1464,12 @@ hstore_to_jsonb(PG_FUNCTION_ARGS)
 			val.val.string.len = HSTORE_VALLEN(entries, i);
 			val.val.string.val = HSTORE_VAL(entries, base, i);
 		}
-		(void) pushJsonbValue(&state, WJB_VALUE, &val);
+		pushJsonbValue(&state, WJB_VALUE, &val);
 	}
 
-	res = pushJsonbValue(&state, WJB_END_OBJECT, NULL);
+	pushJsonbValue(&state, WJB_END_OBJECT, NULL);
 
-	PG_RETURN_POINTER(JsonbValueToJsonb(res));
+	PG_RETURN_POINTER(JsonbValueToJsonb(state.result));
 }
 
 PG_FUNCTION_INFO_V1(hstore_to_jsonb_loose);
@@ -1482,13 +1481,12 @@ hstore_to_jsonb_loose(PG_FUNCTION_ARGS)
 	int			count = HS_COUNT(in);
 	char	   *base = STRPTR(in);
 	HEntry	   *entries = ARRPTR(in);
-	JsonbParseState *state = NULL;
-	JsonbValue *res;
+	JsonbInState state = {0};
 	StringInfoData tmp;
 
 	initStringInfo(&tmp);
 
-	(void) pushJsonbValue(&state, WJB_BEGIN_OBJECT, NULL);
+	pushJsonbValue(&state, WJB_BEGIN_OBJECT, NULL);
 
 	for (i = 0; i < count; i++)
 	{
@@ -1499,7 +1497,7 @@ hstore_to_jsonb_loose(PG_FUNCTION_ARGS)
 		key.val.string.len = HSTORE_KEYLEN(entries, i);
 		key.val.string.val = HSTORE_KEY(entries, base, i);
 
-		(void) pushJsonbValue(&state, WJB_KEY, &key);
+		pushJsonbValue(&state, WJB_KEY, &key);
 
 		if (HSTORE_VALISNULL(entries, i))
 		{
@@ -1541,10 +1539,10 @@ hstore_to_jsonb_loose(PG_FUNCTION_ARGS)
 				val.val.string.val = HSTORE_VAL(entries, base, i);
 			}
 		}
-		(void) pushJsonbValue(&state, WJB_VALUE, &val);
+		pushJsonbValue(&state, WJB_VALUE, &val);
 	}
 
-	res = pushJsonbValue(&state, WJB_END_OBJECT, NULL);
+	pushJsonbValue(&state, WJB_END_OBJECT, NULL);
 
-	PG_RETURN_POINTER(JsonbValueToJsonb(res));
+	PG_RETURN_POINTER(JsonbValueToJsonb(state.result));
 }

contrib/jsonb_plperl/jsonb_plperl.c

@@ -13,7 +13,7 @@ PG_MODULE_MAGIC_EXT(
 );
 
 static SV  *Jsonb_to_SV(JsonbContainer *jsonb);
-static JsonbValue *SV_to_JsonbValue(SV *obj, JsonbParseState **ps, bool is_elem);
+static void SV_to_JsonbValue(SV *obj, JsonbInState *ps, bool is_elem);
 
 
 static SV  *
@@ -127,8 +127,8 @@ Jsonb_to_SV(JsonbContainer *jsonb)
 	}
 }
 
-static JsonbValue *
-AV_to_JsonbValue(AV *in, JsonbParseState **jsonb_state)
+static void
+AV_to_JsonbValue(AV *in, JsonbInState *jsonb_state)
 {
 	dTHX;
 	SSize_t		pcount = av_len(in) + 1;
@@ -141,14 +141,14 @@ AV_to_JsonbValue(AV *in, JsonbParseState **jsonb_state)
 		SV		  **value = av_fetch(in, i, FALSE);
 
 		if (value)
-			(void) SV_to_JsonbValue(*value, jsonb_state, true);
+			SV_to_JsonbValue(*value, jsonb_state, true);
 	}
 
-	return pushJsonbValue(jsonb_state, WJB_END_ARRAY, NULL);
+	pushJsonbValue(jsonb_state, WJB_END_ARRAY, NULL);
 }
 
-static JsonbValue *
-HV_to_JsonbValue(HV *obj, JsonbParseState **jsonb_state)
+static void
+HV_to_JsonbValue(HV *obj, JsonbInState *jsonb_state)
 {
 	dTHX;
 	JsonbValue	key;
@@ -167,14 +167,14 @@ HV_to_JsonbValue(HV *obj, JsonbParseState **jsonb_state)
 		key.val.string.val = pnstrdup(kstr, klen);
 		key.val.string.len = klen;
 		pushJsonbValue(jsonb_state, WJB_KEY, &key);
-		(void) SV_to_JsonbValue(val, jsonb_state, false);
+		SV_to_JsonbValue(val, jsonb_state, false);
 	}
 
-	return pushJsonbValue(jsonb_state, WJB_END_OBJECT, NULL);
+	pushJsonbValue(jsonb_state, WJB_END_OBJECT, NULL);
 }
 
-static JsonbValue *
-SV_to_JsonbValue(SV *in, JsonbParseState **jsonb_state, bool is_elem)
+static void
+SV_to_JsonbValue(SV *in, JsonbInState *jsonb_state, bool is_elem)
 {
 	dTHX;
 	JsonbValue	out;			/* result */
@@ -186,10 +186,12 @@ SV_to_JsonbValue(SV *in, JsonbParseState **jsonb_state, bool is_elem)
 	switch (SvTYPE(in))
 	{
 		case SVt_PVAV:
-			return AV_to_JsonbValue((AV *) in, jsonb_state);
+			AV_to_JsonbValue((AV *) in, jsonb_state);
+			return;
 
 		case SVt_PVHV:
-			return HV_to_JsonbValue((HV *) in, jsonb_state);
+			HV_to_JsonbValue((HV *) in, jsonb_state);
+			return;
 
 		default:
 			if (!SvOK(in))
@@ -259,14 +261,24 @@ SV_to_JsonbValue(SV *in, JsonbParseState **jsonb_state, bool is_elem)
 				ereport(ERROR,
 						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 						 errmsg("cannot transform this Perl type to jsonb")));
-				return NULL;
 			}
 	}
 
-	/* Push result into 'jsonb_state' unless it is a raw scalar. */
-	return *jsonb_state
-		? pushJsonbValue(jsonb_state, is_elem ? WJB_ELEM : WJB_VALUE, &out)
-		: memcpy(palloc_object(JsonbValue), &out, sizeof(JsonbValue));
+	if (jsonb_state->parseState)
+	{
+		/* We're in an array or object, so push value as element or field. */
+		pushJsonbValue(jsonb_state, is_elem ? WJB_ELEM : WJB_VALUE, &out);
+	}
+	else
+	{
+		/*
+		 * We are at top level, so it's a raw scalar.  If we just shove the
+		 * scalar value into jsonb_state->result, JsonbValueToJsonb will take
+		 * care of wrapping it into a dummy array.
+		 */
+		jsonb_state->result = palloc_object(JsonbValue);
+		memcpy(jsonb_state->result, &out, sizeof(JsonbValue));
+	}
 }
 
 
@@ -289,10 +301,9 @@ Datum
 plperl_to_jsonb(PG_FUNCTION_ARGS)
 {
 	dTHX;
-	JsonbParseState *jsonb_state = NULL;
 	SV		   *in = (SV *) PG_GETARG_POINTER(0);
-	JsonbValue *out = SV_to_JsonbValue(in, &jsonb_state, true);
-	Jsonb	   *result = JsonbValueToJsonb(out);
+	JsonbInState jsonb_state = {0};
 
-	PG_RETURN_JSONB_P(result);
+	SV_to_JsonbValue(in, &jsonb_state, true);
+	PG_RETURN_JSONB_P(JsonbValueToJsonb(jsonb_state.result));
 }

contrib/jsonb_plpython/jsonb_plpython.c

@@ -26,8 +26,8 @@ static PLy_elog_impl_t PLy_elog_impl_p;
 static PyObject *decimal_constructor;
 
 static PyObject *PLyObject_FromJsonbContainer(JsonbContainer *jsonb);
-static JsonbValue *PLyObject_ToJsonbValue(PyObject *obj,
-										  JsonbParseState **jsonb_state, bool is_elem);
+static void PLyObject_ToJsonbValue(PyObject *obj,
+								   JsonbInState *jsonb_state, bool is_elem);
 
 typedef PyObject *(*PLyUnicode_FromStringAndSize_t)
 			(const char *s, Py_ssize_t size);
@@ -261,12 +261,11 @@ PLyObject_FromJsonbContainer(JsonbContainer *jsonb)
  *
  * Transform Python dict to JsonbValue.
  */
-static JsonbValue *
-PLyMapping_ToJsonbValue(PyObject *obj, JsonbParseState **jsonb_state)
+static void
+PLyMapping_ToJsonbValue(PyObject *obj, JsonbInState *jsonb_state)
 {
 	Py_ssize_t	pcount;
 	PyObject   *volatile items;
-	JsonbValue *volatile out;
 
 	pcount = PyMapping_Size(obj);
 	items = PyMapping_Items(obj);
@@ -297,19 +296,17 @@ PLyMapping_ToJsonbValue(PyObject *obj, JsonbParseState **jsonb_state)
 				PLyUnicode_ToJsonbValue(key, &jbvKey);
 			}
 
-			(void) pushJsonbValue(jsonb_state, WJB_KEY, &jbvKey);
-			(void) PLyObject_ToJsonbValue(value, jsonb_state, false);
+			pushJsonbValue(jsonb_state, WJB_KEY, &jbvKey);
+			PLyObject_ToJsonbValue(value, jsonb_state, false);
 		}
 
-		out = pushJsonbValue(jsonb_state, WJB_END_OBJECT, NULL);
+		pushJsonbValue(jsonb_state, WJB_END_OBJECT, NULL);
 	}
 	PG_FINALLY();
 	{
 		Py_DECREF(items);
 	}
 	PG_END_TRY();
-
-	return out;
 }
 
 /*
@@ -318,8 +315,8 @@ PLyMapping_ToJsonbValue(PyObject *obj, JsonbParseState **jsonb_state)
  * Transform python list to JsonbValue. Expects transformed PyObject and
  * a state required for jsonb construction.
  */
-static JsonbValue *
-PLySequence_ToJsonbValue(PyObject *obj, JsonbParseState **jsonb_state)
+static void
+PLySequence_ToJsonbValue(PyObject *obj, JsonbInState *jsonb_state)
 {
 	Py_ssize_t	i;
 	Py_ssize_t	pcount;
@@ -336,7 +333,7 @@ PLySequence_ToJsonbValue(PyObject *obj, JsonbParseState **jsonb_state)
 			value = PySequence_GetItem(obj, i);
 			Assert(value);
 
-			(void) PLyObject_ToJsonbValue(value, jsonb_state, true);
+			PLyObject_ToJsonbValue(value, jsonb_state, true);
 			Py_XDECREF(value);
 			value = NULL;
 		}
@@ -348,7 +345,7 @@ PLySequence_ToJsonbValue(PyObject *obj, JsonbParseState **jsonb_state)
 	}
 	PG_END_TRY();
 
-	return pushJsonbValue(jsonb_state, WJB_END_ARRAY, NULL);
+	pushJsonbValue(jsonb_state, WJB_END_ARRAY, NULL);
 }
 
 /*
@@ -406,17 +403,23 @@ PLyNumber_ToJsonbValue(PyObject *obj, JsonbValue *jbvNum)
  *
  * Transform python object to JsonbValue.
  */
-static JsonbValue *
-PLyObject_ToJsonbValue(PyObject *obj, JsonbParseState **jsonb_state, bool is_elem)
+static void
+PLyObject_ToJsonbValue(PyObject *obj, JsonbInState *jsonb_state, bool is_elem)
 {
 	JsonbValue *out;
 
 	if (!PyUnicode_Check(obj))
 	{
 		if (PySequence_Check(obj))
-			return PLySequence_ToJsonbValue(obj, jsonb_state);
+		{
+			PLySequence_ToJsonbValue(obj, jsonb_state);
+			return;
+		}
 		else if (PyMapping_Check(obj))
-			return PLyMapping_ToJsonbValue(obj, jsonb_state);
+		{
+			PLyMapping_ToJsonbValue(obj, jsonb_state);
+			return;
+		}
 	}
 
 	out = palloc_object(JsonbValue);
@@ -443,10 +446,20 @@ PLyObject_ToJsonbValue(PyObject *obj, JsonbParseState **jsonb_state, bool is_ele
 				 errmsg("Python type \"%s\" cannot be transformed to jsonb",
 						PLyObject_AsString((PyObject *) obj->ob_type))));
 
-	/* Push result into 'jsonb_state' unless it is raw scalar value. */
-	return (*jsonb_state ?
-			pushJsonbValue(jsonb_state, is_elem ? WJB_ELEM : WJB_VALUE, out) :
-			out);
+	if (jsonb_state->parseState)
+	{
+		/* We're in an array or object, so push value as element or field. */
+		pushJsonbValue(jsonb_state, is_elem ? WJB_ELEM : WJB_VALUE, out);
+	}
+	else
+	{
+		/*
+		 * We are at top level, so it's a raw scalar.  If we just shove the
+		 * scalar value into jsonb_state->result, JsonbValueToJsonb will take
+		 * care of wrapping it into a dummy array.
+		 */
+		jsonb_state->result = out;
+	}
 }
 
 /*
@@ -458,13 +471,11 @@ PG_FUNCTION_INFO_V1(plpython_to_jsonb);
 Datum
 plpython_to_jsonb(PG_FUNCTION_ARGS)
 {
-	PyObject   *obj;
-	JsonbValue *out;
-	JsonbParseState *jsonb_state = NULL;
+	PyObject   *obj = (PyObject *) PG_GETARG_POINTER(0);
+	JsonbInState jsonb_state = {0};
 
-	obj = (PyObject *) PG_GETARG_POINTER(0);
-	out = PLyObject_ToJsonbValue(obj, &jsonb_state, true);
-	PG_RETURN_POINTER(JsonbValueToJsonb(out));
+	PLyObject_ToJsonbValue(obj, &jsonb_state, true);
+	PG_RETURN_POINTER(JsonbValueToJsonb(jsonb_state.result));
 }
 
 /*

src/backend/utils/adt/jsonb.c

@@ -25,14 +25,6 @@
 #include "utils/lsyscache.h"
 #include "utils/typcache.h"
 
-typedef struct JsonbInState
-{
-	JsonbParseState *parseState;
-	JsonbValue *res;
-	bool		unique_keys;
-	Node	   *escontext;
-} JsonbInState;
-
 typedef struct JsonbAggState
 {
 	JsonbInState *res;
@@ -270,8 +262,8 @@ jsonb_from_cstring(char *json, int len, bool unique_keys, Node *escontext)
 	if (!pg_parse_json_or_errsave(&lex, &sem, escontext))
 		return (Datum) 0;
 
-	/* after parsing, the item member has the composed jsonb structure */
-	PG_RETURN_POINTER(JsonbValueToJsonb(state.res));
+	/* after parsing, the result field has the composed jsonb structure */
+	PG_RETURN_POINTER(JsonbValueToJsonb(state.result));
 }
 
 static bool
@@ -292,7 +284,7 @@ jsonb_in_object_start(void *pstate)
 {
 	JsonbInState *_state = (JsonbInState *) pstate;
 
-	_state->res = pushJsonbValue(&_state->parseState, WJB_BEGIN_OBJECT, NULL);
+	pushJsonbValue(_state, WJB_BEGIN_OBJECT, NULL);
 	_state->parseState->unique_keys = _state->unique_keys;
 
 	return JSON_SUCCESS;
@@ -303,7 +295,7 @@ jsonb_in_object_end(void *pstate)
 {
 	JsonbInState *_state = (JsonbInState *) pstate;
 
-	_state->res = pushJsonbValue(&_state->parseState, WJB_END_OBJECT, NULL);
+	pushJsonbValue(_state, WJB_END_OBJECT, NULL);
 
 	return JSON_SUCCESS;
 }
@@ -313,7 +305,7 @@ jsonb_in_array_start(void *pstate)
 {
 	JsonbInState *_state = (JsonbInState *) pstate;
 
-	_state->res = pushJsonbValue(&_state->parseState, WJB_BEGIN_ARRAY, NULL);
+	pushJsonbValue(_state, WJB_BEGIN_ARRAY, NULL);
 
 	return JSON_SUCCESS;
 }
@@ -323,7 +315,7 @@ jsonb_in_array_end(void *pstate)
 {
 	JsonbInState *_state = (JsonbInState *) pstate;
 
-	_state->res = pushJsonbValue(&_state->parseState, WJB_END_ARRAY, NULL);
+	pushJsonbValue(_state, WJB_END_ARRAY, NULL);
 
 	return JSON_SUCCESS;
 }
@@ -341,7 +333,7 @@ jsonb_in_object_field_start(void *pstate, char *fname, bool isnull)
 		return JSON_SEM_ACTION_FAILED;
 	v.val.string.val = fname;
 
-	_state->res = pushJsonbValue(&_state->parseState, WJB_KEY, &v);
+	pushJsonbValue(_state, WJB_KEY, &v);
 
 	return JSON_SUCCESS;
 }
@@ -435,9 +427,9 @@ jsonb_in_scalar(void *pstate, char *token, JsonTokenType tokentype)
 		va.val.array.rawScalar = true;
 		va.val.array.nElems = 1;
 
-		_state->res = pushJsonbValue(&_state->parseState, WJB_BEGIN_ARRAY, &va);
-		_state->res = pushJsonbValue(&_state->parseState, WJB_ELEM, &v);
-		_state->res = pushJsonbValue(&_state->parseState, WJB_END_ARRAY, NULL);
+		pushJsonbValue(_state, WJB_BEGIN_ARRAY, &va);
+		pushJsonbValue(_state, WJB_ELEM, &v);
+		pushJsonbValue(_state, WJB_END_ARRAY, NULL);
 	}
 	else
 	{
@@ -446,10 +438,10 @@ jsonb_in_scalar(void *pstate, char *token, JsonTokenType tokentype)
 		switch (o->type)
 		{
 			case jbvArray:
-				_state->res = pushJsonbValue(&_state->parseState, WJB_ELEM, &v);
+				pushJsonbValue(_state, WJB_ELEM, &v);
 				break;
 			case jbvObject:
-				_state->res = pushJsonbValue(&_state->parseState, WJB_VALUE, &v);
+				pushJsonbValue(_state, WJB_VALUE, &v);
 				break;
 			default:
 				elog(ERROR, "unexpected parent of nested structure");
@@ -795,11 +787,9 @@ datum_to_jsonb_internal(Datum val, bool is_null, JsonbInState *result,
 						{
 							if (type == WJB_END_ARRAY || type == WJB_END_OBJECT ||
 								type == WJB_BEGIN_ARRAY || type == WJB_BEGIN_OBJECT)
-								result->res = pushJsonbValue(&result->parseState,
-															 type, NULL);
+								pushJsonbValue(result, type, NULL);
 							else
-								result->res = pushJsonbValue(&result->parseState,
-															 type, &jb);
+								pushJsonbValue(result, type, &jb);
 						}
 					}
 				}
@@ -830,9 +820,9 @@ datum_to_jsonb_internal(Datum val, bool is_null, JsonbInState *result,
 		va.val.array.rawScalar = true;
 		va.val.array.nElems = 1;
 
-		result->res = pushJsonbValue(&result->parseState, WJB_BEGIN_ARRAY, &va);
-		result->res = pushJsonbValue(&result->parseState, WJB_ELEM, &jb);
-		result->res = pushJsonbValue(&result->parseState, WJB_END_ARRAY, NULL);
+		pushJsonbValue(result, WJB_BEGIN_ARRAY, &va);
+		pushJsonbValue(result, WJB_ELEM, &jb);
+		pushJsonbValue(result, WJB_END_ARRAY, NULL);
 	}
 	else
 	{
@@ -841,10 +831,10 @@ datum_to_jsonb_internal(Datum val, bool is_null, JsonbInState *result,
 		switch (o->type)
 		{
 			case jbvArray:
-				result->res = pushJsonbValue(&result->parseState, WJB_ELEM, &jb);
+				pushJsonbValue(result, WJB_ELEM, &jb);
 				break;
 			case jbvObject:
-				result->res = pushJsonbValue(&result->parseState,
+				pushJsonbValue(result,
 							   key_scalar ? WJB_KEY : WJB_VALUE,
 							   &jb);
 				break;
@@ -868,7 +858,7 @@ array_dim_to_jsonb(JsonbInState *result, int dim, int ndims, int *dims, const Da
 
 	Assert(dim < ndims);
 
-	result->res = pushJsonbValue(&result->parseState, WJB_BEGIN_ARRAY, NULL);
+	pushJsonbValue(result, WJB_BEGIN_ARRAY, NULL);
 
 	for (i = 1; i <= dims[dim]; i++)
 	{
@@ -885,7 +875,7 @@ array_dim_to_jsonb(JsonbInState *result, int dim, int ndims, int *dims, const Da
 		}
 	}
 
-	result->res = pushJsonbValue(&result->parseState, WJB_END_ARRAY, NULL);
+	pushJsonbValue(result, WJB_END_ARRAY, NULL);
 }
 
 /*
@@ -914,8 +904,8 @@ array_to_jsonb_internal(Datum array, JsonbInState *result)
 
 	if (nitems <= 0)
 	{
-		result->res = pushJsonbValue(&result->parseState, WJB_BEGIN_ARRAY, NULL);
-		result->res = pushJsonbValue(&result->parseState, WJB_END_ARRAY, NULL);
+		pushJsonbValue(result, WJB_BEGIN_ARRAY, NULL);
+		pushJsonbValue(result, WJB_END_ARRAY, NULL);
 		return;
 	}
 
@@ -962,7 +952,7 @@ composite_to_jsonb(Datum composite, JsonbInState *result)
 	tmptup.t_data = td;
 	tuple = &tmptup;
 
-	result->res = pushJsonbValue(&result->parseState, WJB_BEGIN_OBJECT, NULL);
+	pushJsonbValue(result, WJB_BEGIN_OBJECT, NULL);
 
 	for (i = 0; i < tupdesc->natts; i++)
 	{
@@ -984,7 +974,7 @@ composite_to_jsonb(Datum composite, JsonbInState *result)
 		v.val.string.len = strlen(attname);
 		v.val.string.val = attname;
 
-		result->res = pushJsonbValue(&result->parseState, WJB_KEY, &v);
+		pushJsonbValue(result, WJB_KEY, &v);
 
 		val = heap_getattr(tuple, i + 1, tupdesc, &isnull);
 
@@ -1001,7 +991,7 @@ composite_to_jsonb(Datum composite, JsonbInState *result)
 								false);
 	}
 
-	result->res = pushJsonbValue(&result->parseState, WJB_END_OBJECT, NULL);
+	pushJsonbValue(result, WJB_END_OBJECT, NULL);
 	ReleaseTupleDesc(tupdesc);
 }
 
@@ -1119,7 +1109,7 @@ datum_to_jsonb(Datum val, JsonTypeCategory tcategory, Oid outfuncoid)
 	datum_to_jsonb_internal(val, false, &result, tcategory, outfuncoid,
 							false);
 
-	return JsonbPGetDatum(JsonbValueToJsonb(result.res));
+	return JsonbPGetDatum(JsonbValueToJsonb(result.result));
 }
 
 Datum
@@ -1139,7 +1129,7 @@ jsonb_build_object_worker(int nargs, const Datum *args, const bool *nulls, const
 
 	memset(&result, 0, sizeof(JsonbInState));
 
-	result.res = pushJsonbValue(&result.parseState, WJB_BEGIN_OBJECT, NULL);
+	pushJsonbValue(&result, WJB_BEGIN_OBJECT, NULL);
 	result.parseState->unique_keys = unique_keys;
 	result.parseState->skip_nulls = absent_on_null;
 
@@ -1166,9 +1156,9 @@ jsonb_build_object_worker(int nargs, const Datum *args, const bool *nulls, const
 		add_jsonb(args[i + 1], nulls[i + 1], &result, types[i + 1], false);
 	}
 
-	result.res = pushJsonbValue(&result.parseState, WJB_END_OBJECT, NULL);
+	pushJsonbValue(&result, WJB_END_OBJECT, NULL);
 
-	return JsonbPGetDatum(JsonbValueToJsonb(result.res));
+	return JsonbPGetDatum(JsonbValueToJsonb(result.result));
 }
 
 /*
@@ -1201,10 +1191,10 @@ jsonb_build_object_noargs(PG_FUNCTION_ARGS)
 
 	memset(&result, 0, sizeof(JsonbInState));
 
-	(void) pushJsonbValue(&result.parseState, WJB_BEGIN_OBJECT, NULL);
-	result.res = pushJsonbValue(&result.parseState, WJB_END_OBJECT, NULL);
+	pushJsonbValue(&result, WJB_BEGIN_OBJECT, NULL);
+	pushJsonbValue(&result, WJB_END_OBJECT, NULL);
 
-	PG_RETURN_POINTER(JsonbValueToJsonb(result.res));
+	PG_RETURN_POINTER(JsonbValueToJsonb(result.result));
 }
 
 Datum
@@ -1216,7 +1206,7 @@ jsonb_build_array_worker(int nargs, const Datum *args, const bool *nulls, const
 
 	memset(&result, 0, sizeof(JsonbInState));
 
-	result.res = pushJsonbValue(&result.parseState, WJB_BEGIN_ARRAY, NULL);
+	pushJsonbValue(&result, WJB_BEGIN_ARRAY, NULL);
 
 	for (i = 0; i < nargs; i++)
 	{
@@ -1226,9 +1216,9 @@ jsonb_build_array_worker(int nargs, const Datum *args, const bool *nulls, const
 		add_jsonb(args[i], nulls[i], &result, types[i], false);
 	}
 
-	result.res = pushJsonbValue(&result.parseState, WJB_END_ARRAY, NULL);
+	pushJsonbValue(&result, WJB_END_ARRAY, NULL);
 
-	return JsonbPGetDatum(JsonbValueToJsonb(result.res));
+	return JsonbPGetDatum(JsonbValueToJsonb(result.result));
 }
 
 /*
@@ -1262,10 +1252,10 @@ jsonb_build_array_noargs(PG_FUNCTION_ARGS)
 
 	memset(&result, 0, sizeof(JsonbInState));
 
-	(void) pushJsonbValue(&result.parseState, WJB_BEGIN_ARRAY, NULL);
-	result.res = pushJsonbValue(&result.parseState, WJB_END_ARRAY, NULL);
+	pushJsonbValue(&result, WJB_BEGIN_ARRAY, NULL);
+	pushJsonbValue(&result, WJB_END_ARRAY, NULL);
 
-	PG_RETURN_POINTER(JsonbValueToJsonb(result.res));
+	PG_RETURN_POINTER(JsonbValueToJsonb(result.result));
 }
 
 
@@ -1290,7 +1280,7 @@ jsonb_object(PG_FUNCTION_ARGS)
 
 	memset(&result, 0, sizeof(JsonbInState));
 
-	(void) pushJsonbValue(&result.parseState, WJB_BEGIN_OBJECT, NULL);
+	pushJsonbValue(&result, WJB_BEGIN_OBJECT, NULL);
 
 	switch (ndims)
 	{
@@ -1341,7 +1331,7 @@ jsonb_object(PG_FUNCTION_ARGS)
 		v.val.string.len = len;
 		v.val.string.val = str;
 
-		(void) pushJsonbValue(&result.parseState, WJB_KEY, &v);
+		pushJsonbValue(&result, WJB_KEY, &v);
 
 		if (in_nulls[i * 2 + 1])
 		{
@@ -1358,16 +1348,16 @@ jsonb_object(PG_FUNCTION_ARGS)
 			v.val.string.val = str;
 		}
 
-		(void) pushJsonbValue(&result.parseState, WJB_VALUE, &v);
+		pushJsonbValue(&result, WJB_VALUE, &v);
 	}
 
 	pfree(in_datums);
 	pfree(in_nulls);
 
 close_object:
-	result.res = pushJsonbValue(&result.parseState, WJB_END_OBJECT, NULL);
+	pushJsonbValue(&result, WJB_END_OBJECT, NULL);
 
-	PG_RETURN_POINTER(JsonbValueToJsonb(result.res));
+	PG_RETURN_POINTER(JsonbValueToJsonb(result.result));
 }
 
 /*
@@ -1394,7 +1384,7 @@ jsonb_object_two_arg(PG_FUNCTION_ARGS)
 
 	memset(&result, 0, sizeof(JsonbInState));
 
-	(void) pushJsonbValue(&result.parseState, WJB_BEGIN_OBJECT, NULL);
+	pushJsonbValue(&result, WJB_BEGIN_OBJECT, NULL);
 
 	if (nkdims > 1 || nkdims != nvdims)
 		ereport(ERROR,
@@ -1431,7 +1421,7 @@ jsonb_object_two_arg(PG_FUNCTION_ARGS)
 		v.val.string.len = len;
 		v.val.string.val = str;
 
-		(void) pushJsonbValue(&result.parseState, WJB_KEY, &v);
+		pushJsonbValue(&result, WJB_KEY, &v);
 
 		if (val_nulls[i])
 		{
@@ -1448,7 +1438,7 @@ jsonb_object_two_arg(PG_FUNCTION_ARGS)
 			v.val.string.val = str;
 		}
 
-		(void) pushJsonbValue(&result.parseState, WJB_VALUE, &v);
+		pushJsonbValue(&result, WJB_VALUE, &v);
 	}
 
 	pfree(key_datums);
@@ -1457,9 +1447,9 @@ jsonb_object_two_arg(PG_FUNCTION_ARGS)
 	pfree(val_nulls);
 
 close_object:
-	result.res = pushJsonbValue(&result.parseState, WJB_END_OBJECT, NULL);
+	pushJsonbValue(&result, WJB_END_OBJECT, NULL);
 
-	PG_RETURN_POINTER(JsonbValueToJsonb(result.res));
+	PG_RETURN_POINTER(JsonbValueToJsonb(result.result));
 }
 
 
@@ -1534,8 +1524,7 @@ jsonb_agg_transfn_worker(FunctionCallInfo fcinfo, bool absent_on_null)
 		state = palloc(sizeof(JsonbAggState));
 		result = palloc0(sizeof(JsonbInState));
 		state->res = result;
-		result->res = pushJsonbValue(&result->parseState,
-									 WJB_BEGIN_ARRAY, NULL);
+		pushJsonbValue(result, WJB_BEGIN_ARRAY, NULL);
 		MemoryContextSwitchTo(oldcontext);
 
 		json_categorize_type(arg_type, true, &state->val_category,
@@ -1559,7 +1548,7 @@ jsonb_agg_transfn_worker(FunctionCallInfo fcinfo, bool absent_on_null)
 	datum_to_jsonb_internal(val, PG_ARGISNULL(1), &elem, state->val_category,
 							state->val_output_func, false);
 
-	jbelem = JsonbValueToJsonb(elem.res);
+	jbelem = JsonbValueToJsonb(elem.result);
 
 	/* switch to the aggregate context for accumulation operations */
 
@@ -1575,18 +1564,15 @@ jsonb_agg_transfn_worker(FunctionCallInfo fcinfo, bool absent_on_null)
 				if (v.val.array.rawScalar)
 					single_scalar = true;
 				else
-					result->res = pushJsonbValue(&result->parseState,
-												 type, NULL);
+					pushJsonbValue(result, type, NULL);
 				break;
 			case WJB_END_ARRAY:
 				if (!single_scalar)
-					result->res = pushJsonbValue(&result->parseState,
-												 type, NULL);
+					pushJsonbValue(result, type, NULL);
 				break;
 			case WJB_BEGIN_OBJECT:
 			case WJB_END_OBJECT:
-				result->res = pushJsonbValue(&result->parseState,
-											 type, NULL);
+				pushJsonbValue(result, type, NULL);
 				break;
 			case WJB_ELEM:
 			case WJB_KEY:
@@ -1606,8 +1592,7 @@ jsonb_agg_transfn_worker(FunctionCallInfo fcinfo, bool absent_on_null)
 						DatumGetNumeric(DirectFunctionCall1(numeric_uplus,
 															NumericGetDatum(v.val.numeric)));
 				}
-				result->res = pushJsonbValue(&result->parseState,
-											 type, &v);
+				pushJsonbValue(result, type, &v);
 				break;
 			default:
 				elog(ERROR, "unknown jsonb iterator token type");
@@ -1662,10 +1647,9 @@ jsonb_agg_finalfn(PG_FUNCTION_ARGS)
 
 	result.parseState = clone_parse_state(arg->res->parseState);
 
-	result.res = pushJsonbValue(&result.parseState,
-								WJB_END_ARRAY, NULL);
+	pushJsonbValue(&result, WJB_END_ARRAY, NULL);
 
-	out = JsonbValueToJsonb(result.res);
+	out = JsonbValueToJsonb(result.result);
 
 	PG_RETURN_POINTER(out);
 }
@@ -1704,8 +1688,7 @@ jsonb_object_agg_transfn_worker(FunctionCallInfo fcinfo,
 		state = palloc(sizeof(JsonbAggState));
 		result = palloc0(sizeof(JsonbInState));
 		state->res = result;
-		result->res = pushJsonbValue(&result->parseState,
-									 WJB_BEGIN_OBJECT, NULL);
+		pushJsonbValue(result, WJB_BEGIN_OBJECT, NULL);
 		result->parseState->unique_keys = unique_keys;
 		result->parseState->skip_nulls = absent_on_null;
 
@@ -1760,7 +1743,7 @@ jsonb_object_agg_transfn_worker(FunctionCallInfo fcinfo,
 	datum_to_jsonb_internal(val, false, &elem, state->key_category,
 							state->key_output_func, true);
 
-	jbkey = JsonbValueToJsonb(elem.res);
+	jbkey = JsonbValueToJsonb(elem.result);
 
 	val = PG_ARGISNULL(2) ? (Datum) 0 : PG_GETARG_DATUM(2);
 
@@ -1769,7 +1752,7 @@ jsonb_object_agg_transfn_worker(FunctionCallInfo fcinfo,
 	datum_to_jsonb_internal(val, PG_ARGISNULL(2), &elem, state->val_category,
 							state->val_output_func, false);
 
-	jbval = JsonbValueToJsonb(elem.res);
+	jbval = JsonbValueToJsonb(elem.result);
 
 	it = JsonbIteratorInit(&jbkey->root);
 
@@ -1806,14 +1789,12 @@ jsonb_object_agg_transfn_worker(FunctionCallInfo fcinfo,
 							(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 							 errmsg("object keys must be strings")));
 				}
-				result->res = pushJsonbValue(&result->parseState,
-											 WJB_KEY, &v);
+				pushJsonbValue(result, WJB_KEY, &v);
 
 				if (skip)
 				{
 					v.type = jbvNull;
-					result->res = pushJsonbValue(&result->parseState,
-												 WJB_VALUE, &v);
+					pushJsonbValue(result, WJB_VALUE, &v);
 					MemoryContextSwitchTo(oldcontext);
 					PG_RETURN_POINTER(state);
 				}
@@ -1845,18 +1826,15 @@ jsonb_object_agg_transfn_worker(FunctionCallInfo fcinfo,
 				if (v.val.array.rawScalar)
 					single_scalar = true;
 				else
-					result->res = pushJsonbValue(&result->parseState,
-												 type, NULL);
+					pushJsonbValue(result, type, NULL);
 				break;
 			case WJB_END_ARRAY:
 				if (!single_scalar)
-					result->res = pushJsonbValue(&result->parseState,
-												 type, NULL);
+					pushJsonbValue(result, type, NULL);
 				break;
 			case WJB_BEGIN_OBJECT:
 			case WJB_END_OBJECT:
-				result->res = pushJsonbValue(&result->parseState,
-											 type, NULL);
+				pushJsonbValue(result, type, NULL);
 				break;
 			case WJB_ELEM:
 			case WJB_KEY:
@@ -1876,9 +1854,7 @@ jsonb_object_agg_transfn_worker(FunctionCallInfo fcinfo,
 						DatumGetNumeric(DirectFunctionCall1(numeric_uplus,
 															NumericGetDatum(v.val.numeric)));
 				}
-				result->res = pushJsonbValue(&result->parseState,
-											 single_scalar ? WJB_VALUE : type,
-											 &v);
+				pushJsonbValue(result, single_scalar ? WJB_VALUE : type, &v);
 				break;
 			default:
 				elog(ERROR, "unknown jsonb iterator token type");
@@ -1953,10 +1929,9 @@ jsonb_object_agg_finalfn(PG_FUNCTION_ARGS)
 
 	result.parseState = clone_parse_state(arg->res->parseState);
 
-	result.res = pushJsonbValue(&result.parseState,
-								WJB_END_OBJECT, NULL);
+	pushJsonbValue(&result, WJB_END_OBJECT, NULL);
 
-	out = JsonbValueToJsonb(result.res);
+	out = JsonbValueToJsonb(result.result);
 
 	PG_RETURN_POINTER(out);
 }

src/backend/utils/adt/jsonb_util.c

@@ -14,10 +14,13 @@
 #include "postgres.h"
 
 #include "catalog/pg_collation.h"
+#include "catalog/pg_type.h"
 #include "common/hashfn.h"
 #include "miscadmin.h"
 #include "port/pg_bitutils.h"
+#include "utils/date.h"
 #include "utils/datetime.h"
+#include "utils/datum.h"
 #include "utils/fmgrprotos.h"
 #include "utils/json.h"
 #include "utils/jsonb.h"
@@ -54,17 +57,18 @@ static short padBufferToInt(StringInfo buffer);
 
 static JsonbIterator *iteratorFromContainer(JsonbContainer *container, JsonbIterator *parent);
 static JsonbIterator *freeAndGetParent(JsonbIterator *it);
-static JsonbParseState *pushState(JsonbParseState **pstate);
-static void appendKey(JsonbParseState *pstate, JsonbValue *string);
-static void appendValue(JsonbParseState *pstate, JsonbValue *scalarVal);
-static void appendElement(JsonbParseState *pstate, JsonbValue *scalarVal);
+static JsonbParseState *pushState(JsonbInState *pstate);
+static void appendKey(JsonbInState *pstate, JsonbValue *string, bool needCopy);
+static void appendValue(JsonbInState *pstate, JsonbValue *scalarVal, bool needCopy);
+static void appendElement(JsonbInState *pstate, JsonbValue *scalarVal, bool needCopy);
+static void copyScalarSubstructure(JsonbValue *v, MemoryContext outcontext);
 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 *binequal);
 static void uniqueifyJsonbObject(JsonbValue *object, bool unique_keys,
 								 bool skip_nulls);
-static JsonbValue *pushJsonbValueScalar(JsonbParseState **pstate,
+static void pushJsonbValueScalar(JsonbInState *pstate,
 								 JsonbIteratorToken seq,
 								 JsonbValue *scalarVal);
 
@@ -95,9 +99,8 @@ JsonbValueToJsonb(JsonbValue *val)
 
 	if (IsAJsonbScalar(val))
 	{
-		/* Scalar value */
-		JsonbParseState *pstate = NULL;
-		JsonbValue *res;
+		/* Scalar value, so wrap it in an array */
+		JsonbInState pstate = {0};
 		JsonbValue	scalarArray;
 
 		scalarArray.type = jbvArray;
@@ -106,9 +109,9 @@ JsonbValueToJsonb(JsonbValue *val)
 
 		pushJsonbValue(&pstate, WJB_BEGIN_ARRAY, &scalarArray);
 		pushJsonbValue(&pstate, WJB_ELEM, val);
-		res = pushJsonbValue(&pstate, WJB_END_ARRAY, NULL);
+		pushJsonbValue(&pstate, WJB_END_ARRAY, NULL);
 
-		out = convertToJsonb(res);
+		out = convertToJsonb(pstate.result);
 	}
 	else if (val->type == jbvObject || val->type == jbvArray)
 	{
@@ -547,13 +550,23 @@ fillJsonbValue(JsonbContainer *container, int index,
 }
 
 /*
- * Push JsonbValue into JsonbParseState.
+ * Push JsonbValue into JsonbInState.
  *
- * Used when parsing JSON tokens to form Jsonb, or when converting an in-memory
- * JsonbValue to a Jsonb.
+ * Used, for example, when parsing JSON input.
  *
- * Initial state of *JsonbParseState is NULL, since it'll be allocated here
- * originally (caller will get JsonbParseState back by reference).
+ * *pstate is typically initialized to all-zeroes, except that the caller
+ * may provide outcontext and/or escontext.  (escontext is ignored by this
+ * function and its subroutines, however.)
+ *
+ * "seq" tells what is being pushed (start/end of array or object, key,
+ * value, etc).  WJB_DONE is not used here, but the other values of
+ * JsonbIteratorToken are.  We assume the caller passes a valid sequence
+ * of values.
+ *
+ * The passed "jbval" is typically transient storage, such as a local variable.
+ * We will copy it into the outcontext (CurrentMemoryContext by default).
+ * If outcontext isn't NULL, we will also make copies of any pass-by-reference
+ * scalar values.
  *
  * Only sequential tokens pertaining to non-container types should pass a
  * JsonbValue.  There is one exception -- WJB_BEGIN_ARRAY callers may pass a
@@ -562,18 +575,32 @@ fillJsonbValue(JsonbContainer *container, int index,
  *
  * Values of type jbvBinary, which are rolled up arrays and objects,
  * are unpacked before being added to the result.
+ *
+ * At the end of construction of a JsonbValue, pstate->result will reference
+ * the top-level JsonbValue object.
  */
-JsonbValue *
-pushJsonbValue(JsonbParseState **pstate, JsonbIteratorToken seq,
+void
+pushJsonbValue(JsonbInState *pstate, JsonbIteratorToken seq,
 			   JsonbValue *jbval)
 {
 	JsonbIterator *it;
-	JsonbValue *res = NULL;
 	JsonbValue	v;
 	JsonbIteratorToken tok;
 	int			i;
 
-	if (jbval && (seq == WJB_ELEM || seq == WJB_VALUE) && jbval->type == jbvObject)
+	/*
+	 * pushJsonbValueScalar handles all cases not involving pushing a
+	 * container object as an ELEM or VALUE.
+	 */
+	if (!jbval || IsAJsonbScalar(jbval) ||
+		(seq != WJB_ELEM && seq != WJB_VALUE))
+	{
+		pushJsonbValueScalar(pstate, seq, jbval);
+		return;
+	}
+
+	/* If an object or array is pushed, recursively push its contents */
+	if (jbval->type == jbvObject)
 	{
 		pushJsonbValue(pstate, WJB_BEGIN_OBJECT, NULL);
 		for (i = 0; i < jbval->val.object.nPairs; i++)
@@ -581,32 +608,29 @@ pushJsonbValue(JsonbParseState **pstate, JsonbIteratorToken seq,
 			pushJsonbValue(pstate, WJB_KEY, &jbval->val.object.pairs[i].key);
 			pushJsonbValue(pstate, WJB_VALUE, &jbval->val.object.pairs[i].value);
 		}
-
-		return pushJsonbValue(pstate, WJB_END_OBJECT, NULL);
+		pushJsonbValue(pstate, WJB_END_OBJECT, NULL);
+		return;
 	}
 
-	if (jbval && (seq == WJB_ELEM || seq == WJB_VALUE) && jbval->type == jbvArray)
+	if (jbval->type == jbvArray)
 	{
 		pushJsonbValue(pstate, WJB_BEGIN_ARRAY, NULL);
 		for (i = 0; i < jbval->val.array.nElems; i++)
 		{
 			pushJsonbValue(pstate, WJB_ELEM, &jbval->val.array.elems[i]);
 		}
-
-		return pushJsonbValue(pstate, WJB_END_ARRAY, NULL);
+		pushJsonbValue(pstate, WJB_END_ARRAY, NULL);
+		return;
 	}
 
-	if (!jbval || (seq != WJB_ELEM && seq != WJB_VALUE) ||
-		jbval->type != jbvBinary)
-	{
-		/* drop through */
-		return pushJsonbValueScalar(pstate, seq, jbval);
-	}
+	/* Else it must be a jbvBinary value; push its contents */
+	Assert(jbval->type == jbvBinary);
 
-	/* unpack the binary and add each piece to the pstate */
 	it = JsonbIteratorInit(jbval->val.binary.data);
 
-	if ((jbval->val.binary.data->header & JB_FSCALAR) && *pstate)
+	/* ... with a special case for pushing a raw scalar */
+	if ((jbval->val.binary.data->header & JB_FSCALAR) &&
+		pstate->parseState != NULL)
 	{
 		tok = JsonbIteratorNext(&it, &v, true);
 		Assert(tok == WJB_BEGIN_ARRAY);
@@ -615,197 +639,290 @@ pushJsonbValue(JsonbParseState **pstate, JsonbIteratorToken seq,
 		tok = JsonbIteratorNext(&it, &v, true);
 		Assert(tok == WJB_ELEM);
 
-		res = pushJsonbValueScalar(pstate, seq, &v);
+		pushJsonbValueScalar(pstate, seq, &v);
 
 		tok = JsonbIteratorNext(&it, &v, true);
 		Assert(tok == WJB_END_ARRAY);
 		Assert(it == NULL);
 
-		return res;
+		return;
 	}
 
 	while ((tok = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
-		res = pushJsonbValueScalar(pstate, tok,
+		pushJsonbValueScalar(pstate, tok,
 							 tok < WJB_BEGIN_ARRAY ||
 							 (tok == WJB_BEGIN_ARRAY &&
 							  v.val.array.rawScalar) ? &v : NULL);
-
-	return res;
 }
 
 /*
  * Do the actual pushing, with only scalar or pseudo-scalar-array values
  * accepted.
  */
-static JsonbValue *
-pushJsonbValueScalar(JsonbParseState **pstate, JsonbIteratorToken seq,
+static void
+pushJsonbValueScalar(JsonbInState *pstate, JsonbIteratorToken seq,
 					 JsonbValue *scalarVal)
 {
-	JsonbValue *result = NULL;
+	JsonbParseState *ppstate;
+	JsonbValue *val;
+	MemoryContext outcontext;
 
 	switch (seq)
 	{
 		case WJB_BEGIN_ARRAY:
 			Assert(!scalarVal || scalarVal->val.array.rawScalar);
-			*pstate = pushState(pstate);
-			result = &(*pstate)->contVal;
-			(*pstate)->contVal.type = jbvArray;
-			(*pstate)->contVal.val.array.nElems = 0;
-			(*pstate)->contVal.val.array.rawScalar = (scalarVal &&
+			ppstate = pushState(pstate);
+			val = &ppstate->contVal;
+			val->type = jbvArray;
+			val->val.array.nElems = 0;
+			val->val.array.rawScalar = (scalarVal &&
 										scalarVal->val.array.rawScalar);
 			if (scalarVal && scalarVal->val.array.nElems > 0)
 			{
 				/* Assume that this array is still really a scalar */
 				Assert(scalarVal->type == jbvArray);
-				(*pstate)->size = scalarVal->val.array.nElems;
+				ppstate->size = scalarVal->val.array.nElems;
 			}
 			else
 			{
-				(*pstate)->size = 4;
+				ppstate->size = 4;	/* initial guess at array size */
 			}
-			(*pstate)->contVal.val.array.elems = palloc(sizeof(JsonbValue) *
-														(*pstate)->size);
+			outcontext = pstate->outcontext ? pstate->outcontext : CurrentMemoryContext;
+			val->val.array.elems = MemoryContextAlloc(outcontext,
+													  sizeof(JsonbValue) *
+													  ppstate->size);
 			break;
 		case WJB_BEGIN_OBJECT:
 			Assert(!scalarVal);
-			*pstate = pushState(pstate);
-			result = &(*pstate)->contVal;
-			(*pstate)->contVal.type = jbvObject;
-			(*pstate)->contVal.val.object.nPairs = 0;
-			(*pstate)->size = 4;
-			(*pstate)->contVal.val.object.pairs = palloc(sizeof(JsonbPair) *
-														 (*pstate)->size);
+			ppstate = pushState(pstate);
+			val = &ppstate->contVal;
+			val->type = jbvObject;
+			val->val.object.nPairs = 0;
+			ppstate->size = 4;	/* initial guess at object size */
+			outcontext = pstate->outcontext ? pstate->outcontext : CurrentMemoryContext;
+			val->val.object.pairs = MemoryContextAlloc(outcontext,
+													   sizeof(JsonbPair) *
+													   ppstate->size);
 			break;
 		case WJB_KEY:
 			Assert(scalarVal->type == jbvString);
-			appendKey(*pstate, scalarVal);
+			appendKey(pstate, scalarVal, true);
 			break;
 		case WJB_VALUE:
 			Assert(IsAJsonbScalar(scalarVal));
-			appendValue(*pstate, scalarVal);
+			appendValue(pstate, scalarVal, true);
 			break;
 		case WJB_ELEM:
 			Assert(IsAJsonbScalar(scalarVal));
-			appendElement(*pstate, scalarVal);
+			appendElement(pstate, scalarVal, true);
 			break;
 		case WJB_END_OBJECT:
-			uniqueifyJsonbObject(&(*pstate)->contVal,
-								 (*pstate)->unique_keys,
-								 (*pstate)->skip_nulls);
+			ppstate = pstate->parseState;
+			uniqueifyJsonbObject(&ppstate->contVal,
+								 ppstate->unique_keys,
+								 ppstate->skip_nulls);
 			/* fall through! */
 		case WJB_END_ARRAY:
 			/* Steps here common to WJB_END_OBJECT case */
 			Assert(!scalarVal);
-			result = &(*pstate)->contVal;
+			ppstate = pstate->parseState;
+			val = &ppstate->contVal;
 
 			/*
 			 * Pop stack and push current array/object as value in parent
-			 * array/object
+			 * array/object, or return it as the final result.  We don't need
+			 * to re-copy any scalars that are in the data structure.
 			 */
-			*pstate = (*pstate)->next;
-			if (*pstate)
+			pstate->parseState = ppstate = ppstate->next;
+			if (ppstate)
 			{
-				switch ((*pstate)->contVal.type)
+				switch (ppstate->contVal.type)
 				{
 					case jbvArray:
-						appendElement(*pstate, result);
+						appendElement(pstate, val, false);
 						break;
 					case jbvObject:
-						appendValue(*pstate, result);
+						appendValue(pstate, val, false);
 						break;
 					default:
 						elog(ERROR, "invalid jsonb container type");
 				}
 			}
+			else
+				pstate->result = val;
 			break;
 		default:
 			elog(ERROR, "unrecognized jsonb sequential processing token");
 	}
-
-	return result;
 }
 
 /*
- * pushJsonbValue() worker:  Iteration-like forming of Jsonb
+ * Push a new JsonbParseState onto the JsonbInState's stack
+ *
+ * As a notational convenience, the new state's address is returned.
+ * The caller must initialize the new state's contVal and size fields.
  */
 static JsonbParseState *
-pushState(JsonbParseState **pstate)
+pushState(JsonbInState *pstate)
 {
-	JsonbParseState *ns = palloc(sizeof(JsonbParseState));
+	MemoryContext outcontext = pstate->outcontext ? pstate->outcontext : CurrentMemoryContext;
+	JsonbParseState *ns = MemoryContextAlloc(outcontext,
+											 sizeof(JsonbParseState));
 
-	ns->next = *pstate;
+	ns->next = pstate->parseState;
+	/* This module never changes these fields, but callers can: */
 	ns->unique_keys = false;
 	ns->skip_nulls = false;
 
+	pstate->parseState = ns;
 	return ns;
 }
 
 /*
- * pushJsonbValue() worker:  Append a pair key to state when generating a Jsonb
+ * pushJsonbValue() worker:  Append a pair key to pstate
  */
 static void
-appendKey(JsonbParseState *pstate, JsonbValue *string)
+appendKey(JsonbInState *pstate, JsonbValue *string, bool needCopy)
 {
-	JsonbValue *object = &pstate->contVal;
+	JsonbParseState *ppstate = pstate->parseState;
+	JsonbValue *object = &ppstate->contVal;
+	JsonbPair  *pair;
 
 	Assert(object->type == jbvObject);
 	Assert(string->type == jbvString);
 
-	if (object->val.object.nPairs >= JSONB_MAX_PAIRS)
+	if (object->val.object.nPairs >= ppstate->size)
+	{
+		if (unlikely(object->val.object.nPairs >= JSONB_MAX_PAIRS))
 			ereport(ERROR,
 					(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
 					 errmsg("number of jsonb object pairs exceeds the maximum allowed (%zu)",
 							JSONB_MAX_PAIRS)));
-
-	if (object->val.object.nPairs >= pstate->size)
-	{
-		pstate->size *= 2;
+		ppstate->size = Min(ppstate->size * 2, JSONB_MAX_PAIRS);
 		object->val.object.pairs = repalloc(object->val.object.pairs,
-											sizeof(JsonbPair) * pstate->size);
+											sizeof(JsonbPair) * ppstate->size);
 	}
 
-	object->val.object.pairs[object->val.object.nPairs].key = *string;
-	object->val.object.pairs[object->val.object.nPairs].order = object->val.object.nPairs;
+	pair = &object->val.object.pairs[object->val.object.nPairs];
+	pair->key = *string;
+	pair->order = object->val.object.nPairs;
+
+	if (needCopy)
+		copyScalarSubstructure(&pair->key, pstate->outcontext);
 }
 
 /*
- * pushJsonbValue() worker:  Append a pair value to state when generating a
- * Jsonb
+ * pushJsonbValue() worker:  Append a pair value to pstate
  */
 static void
-appendValue(JsonbParseState *pstate, JsonbValue *scalarVal)
+appendValue(JsonbInState *pstate, JsonbValue *scalarVal, bool needCopy)
 {
-	JsonbValue *object = &pstate->contVal;
+	JsonbValue *object = &pstate->parseState->contVal;
+	JsonbPair  *pair;
 
 	Assert(object->type == jbvObject);
 
-	object->val.object.pairs[object->val.object.nPairs++].value = *scalarVal;
+	pair = &object->val.object.pairs[object->val.object.nPairs];
+	pair->value = *scalarVal;
+	object->val.object.nPairs++;
+
+	if (needCopy)
+		copyScalarSubstructure(&pair->value, pstate->outcontext);
 }
 
 /*
- * pushJsonbValue() worker:  Append an element to state when generating a Jsonb
+ * pushJsonbValue() worker:  Append an array element to pstate
  */
 static void
-appendElement(JsonbParseState *pstate, JsonbValue *scalarVal)
+appendElement(JsonbInState *pstate, JsonbValue *scalarVal, bool needCopy)
 {
-	JsonbValue *array = &pstate->contVal;
+	JsonbParseState *ppstate = pstate->parseState;
+	JsonbValue *array = &ppstate->contVal;
+	JsonbValue *elem;
 
 	Assert(array->type == jbvArray);
 
-	if (array->val.array.nElems >= JSONB_MAX_ELEMS)
+	if (array->val.array.nElems >= ppstate->size)
+	{
+		if (unlikely(array->val.array.nElems >= JSONB_MAX_ELEMS))
 			ereport(ERROR,
 					(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
 					 errmsg("number of jsonb array elements exceeds the maximum allowed (%zu)",
 							JSONB_MAX_ELEMS)));
-
-	if (array->val.array.nElems >= pstate->size)
-	{
-		pstate->size *= 2;
+		ppstate->size = Min(ppstate->size * 2, JSONB_MAX_ELEMS);
 		array->val.array.elems = repalloc(array->val.array.elems,
-										  sizeof(JsonbValue) * pstate->size);
+										  sizeof(JsonbValue) * ppstate->size);
 	}
 
-	array->val.array.elems[array->val.array.nElems++] = *scalarVal;
+	elem = &array->val.array.elems[array->val.array.nElems];
+	*elem = *scalarVal;
+	array->val.array.nElems++;
+
+	if (needCopy)
+		copyScalarSubstructure(elem, pstate->outcontext);
+}
+
+/*
+ * Copy any infrastructure of a scalar JsonbValue into the outcontext,
+ * adjusting the pointer(s) in *v.
+ *
+ * We need not deal with containers here, as the routines above ensure
+ * that they are built fresh.
+ */
+static void
+copyScalarSubstructure(JsonbValue *v, MemoryContext outcontext)
+{
+	MemoryContext oldcontext;
+
+	/* Nothing to do if caller did not specify an outcontext */
+	if (outcontext == NULL)
+		return;
+	switch (v->type)
+	{
+		case jbvNull:
+		case jbvBool:
+			/* pass-by-value, nothing to do */
+			break;
+		case jbvString:
+			{
+				char	   *buf = MemoryContextAlloc(outcontext,
+													 v->val.string.len);
+
+				memcpy(buf, v->val.string.val, v->val.string.len);
+				v->val.string.val = buf;
+			}
+			break;
+		case jbvNumeric:
+			oldcontext = MemoryContextSwitchTo(outcontext);
+			v->val.numeric =
+				DatumGetNumeric(datumCopy(NumericGetDatum(v->val.numeric),
+										  false, -1));
+			MemoryContextSwitchTo(oldcontext);
+			break;
+		case jbvDatetime:
+			switch (v->val.datetime.typid)
+			{
+				case DATEOID:
+				case TIMEOID:
+				case TIMESTAMPOID:
+				case TIMESTAMPTZOID:
+					/* pass-by-value, nothing to do */
+					break;
+				case TIMETZOID:
+					/* pass-by-reference */
+					oldcontext = MemoryContextSwitchTo(outcontext);
+					v->val.datetime.value = datumCopy(v->val.datetime.value,
+													  false, TIMETZ_TYPLEN);
+					MemoryContextSwitchTo(oldcontext);
+					break;
+				default:
+					elog(ERROR, "unexpected jsonb datetime type oid %u",
+						 v->val.datetime.typid);
+			}
+			break;
+		default:
+			elog(ERROR, "invalid jsonb scalar type");
+	}
 }
 
 /*

src/backend/utils/adt/jsonfuncs.c

@@ -475,18 +475,18 @@ static Datum populate_domain(DomainIOData *io, Oid typid, const char *colname,
 							 Node *escontext, bool omit_quotes);
 
 /* functions supporting jsonb_delete, jsonb_set and jsonb_concat */
-static JsonbValue *IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
-								  JsonbParseState **state);
-static JsonbValue *setPath(JsonbIterator **it, const Datum *path_elems,
+static void IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
+						   JsonbInState *state);
+static void setPath(JsonbIterator **it, const Datum *path_elems,
 					const bool *path_nulls, int path_len,
-						   JsonbParseState **st, int level, JsonbValue *newval,
+					JsonbInState *st, int level, JsonbValue *newval,
 					int op_type);
 static void setPathObject(JsonbIterator **it, const Datum *path_elems,
-						  const bool *path_nulls, int path_len, JsonbParseState **st,
+						  const bool *path_nulls, int path_len, JsonbInState *st,
 						  int level,
 						  JsonbValue *newval, uint32 npairs, int op_type);
 static void setPathArray(JsonbIterator **it, const Datum *path_elems,
-						 const bool *path_nulls, int path_len, JsonbParseState **st,
+						 const bool *path_nulls, int path_len, JsonbInState *st,
 						 int level,
 						 JsonbValue *newval, uint32 nelems, int op_type);
 
@@ -1679,8 +1679,7 @@ Datum
 jsonb_set_element(Jsonb *jb, const Datum *path, int path_len,
 				  JsonbValue *newval)
 {
-	JsonbValue *res;
-	JsonbParseState *state = NULL;
+	JsonbInState state = {0};
 	JsonbIterator *it;
 	bool	   *path_nulls = palloc0(path_len * sizeof(bool));
 
@@ -1689,17 +1688,17 @@ jsonb_set_element(Jsonb *jb, const Datum *path, int path_len,
 
 	it = JsonbIteratorInit(&jb->root);
 
-	res = setPath(&it, path, path_nulls, path_len, &state, 0, newval,
+	setPath(&it, path, path_nulls, path_len, &state, 0, newval,
 			JB_PATH_CREATE | JB_PATH_FILL_GAPS |
 			JB_PATH_CONSISTENT_POSITION);
 
 	pfree(path_nulls);
 
-	PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
+	PG_RETURN_JSONB_P(JsonbValueToJsonb(state.result));
 }
 
 static void
-push_null_elements(JsonbParseState **ps, int num)
+push_null_elements(JsonbInState *ps, int num)
 {
 	JsonbValue	null;
 
@@ -1718,7 +1717,7 @@ push_null_elements(JsonbParseState **ps, int num)
  * Caller is responsible to make sure such path does not exist yet.
  */
 static void
-push_path(JsonbParseState **st, int level, const Datum *path_elems,
+push_path(JsonbInState *st, int level, const Datum *path_elems,
 		  const bool *path_nulls, int path_len, JsonbValue *newval)
 {
 	/*
@@ -1758,15 +1757,15 @@ push_path(JsonbParseState **st, int level, const Datum *path_elems,
 			newkey.val.string.val = c;
 			newkey.val.string.len = strlen(c);
 
-			(void) pushJsonbValue(st, WJB_BEGIN_OBJECT, NULL);
-			(void) pushJsonbValue(st, WJB_KEY, &newkey);
+			pushJsonbValue(st, WJB_BEGIN_OBJECT, NULL);
+			pushJsonbValue(st, WJB_KEY, &newkey);
 
 			tpath[i - level] = jbvObject;
 		}
 		else
 		{
 			/* integer, an array is expected */
-			(void) pushJsonbValue(st, WJB_BEGIN_ARRAY, NULL);
+			pushJsonbValue(st, WJB_BEGIN_ARRAY, NULL);
 
 			push_null_elements(st, lindex);
 
@@ -1776,11 +1775,9 @@ push_path(JsonbParseState **st, int level, const Datum *path_elems,
 
 	/* Insert an actual value for either an object or array */
 	if (tpath[(path_len - level) - 1] == jbvArray)
-	{
-		(void) pushJsonbValue(st, WJB_ELEM, newval);
-	}
+		pushJsonbValue(st, WJB_ELEM, newval);
 	else
-		(void) pushJsonbValue(st, WJB_VALUE, newval);
+		pushJsonbValue(st, WJB_VALUE, newval);
 
 	/*
 	 * Close everything up to the last but one level. The last one will be
@@ -1792,9 +1789,9 @@ push_path(JsonbParseState **st, int level, const Datum *path_elems,
 			break;
 
 		if (tpath[i - level] == jbvObject)
-			(void) pushJsonbValue(st, WJB_END_OBJECT, NULL);
+			pushJsonbValue(st, WJB_END_OBJECT, NULL);
 		else
-			(void) pushJsonbValue(st, WJB_END_ARRAY, NULL);
+			pushJsonbValue(st, WJB_END_ARRAY, NULL);
 	}
 }
 
@@ -4544,8 +4541,7 @@ jsonb_strip_nulls(PG_FUNCTION_ARGS)
 	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
 	bool		strip_in_arrays = false;
 	JsonbIterator *it;
-	JsonbParseState *parseState = NULL;
-	JsonbValue *res = NULL;
+	JsonbInState parseState = {0};
 	JsonbValue	v,
 				k;
 	JsonbIteratorToken type;
@@ -4581,7 +4577,7 @@ jsonb_strip_nulls(PG_FUNCTION_ARGS)
 				continue;
 
 			/* otherwise, do a delayed push of the key */
-			(void) pushJsonbValue(&parseState, WJB_KEY, &k);
+			pushJsonbValue(&parseState, WJB_KEY, &k);
 		}
 
 		/* if strip_in_arrays is set, also skip null array elements */
@@ -4590,14 +4586,12 @@ jsonb_strip_nulls(PG_FUNCTION_ARGS)
 				continue;
 
 		if (type == WJB_VALUE || type == WJB_ELEM)
-			res = pushJsonbValue(&parseState, type, &v);
+			pushJsonbValue(&parseState, type, &v);
 		else
-			res = pushJsonbValue(&parseState, type, NULL);
+			pushJsonbValue(&parseState, type, NULL);
 	}
 
-	Assert(res != NULL);
-
-	PG_RETURN_POINTER(JsonbValueToJsonb(res));
+	PG_RETURN_POINTER(JsonbValueToJsonb(parseState.result));
 }
 
 /*
@@ -4627,8 +4621,7 @@ jsonb_concat(PG_FUNCTION_ARGS)
 {
 	Jsonb	   *jb1 = PG_GETARG_JSONB_P(0);
 	Jsonb	   *jb2 = PG_GETARG_JSONB_P(1);
-	JsonbParseState *state = NULL;
-	JsonbValue *res;
+	JsonbInState state = {0};
 	JsonbIterator *it1,
 			   *it2;
 
@@ -4649,11 +4642,9 @@ jsonb_concat(PG_FUNCTION_ARGS)
 	it1 = JsonbIteratorInit(&jb1->root);
 	it2 = JsonbIteratorInit(&jb2->root);
 
-	res = IteratorConcat(&it1, &it2, &state);
+	IteratorConcat(&it1, &it2, &state);
 
-	Assert(res != NULL);
-
-	PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
+	PG_RETURN_JSONB_P(JsonbValueToJsonb(state.result));
 }
 
 
@@ -4670,10 +4661,9 @@ jsonb_delete(PG_FUNCTION_ARGS)
 	text	   *key = PG_GETARG_TEXT_PP(1);
 	char	   *keyptr = VARDATA_ANY(key);
 	int			keylen = VARSIZE_ANY_EXHDR(key);
-	JsonbParseState *state = NULL;
+	JsonbInState pstate = {0};
 	JsonbIterator *it;
-	JsonbValue	v,
-			   *res = NULL;
+	JsonbValue	v;
 	bool		skipNested = false;
 	JsonbIteratorToken r;
 
@@ -4702,12 +4692,10 @@ jsonb_delete(PG_FUNCTION_ARGS)
 			continue;
 		}
 
-		res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
+		pushJsonbValue(&pstate, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
 	}
 
-	Assert(res != NULL);
-
-	PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
+	PG_RETURN_JSONB_P(JsonbValueToJsonb(pstate.result));
 }
 
 /*
@@ -4724,10 +4712,9 @@ jsonb_delete_array(PG_FUNCTION_ARGS)
 	Datum	   *keys_elems;
 	bool	   *keys_nulls;
 	int			keys_len;
-	JsonbParseState *state = NULL;
+	JsonbInState pstate = {0};
 	JsonbIterator *it;
-	JsonbValue	v,
-			   *res = NULL;
+	JsonbValue	v;
 	bool		skipNested = false;
 	JsonbIteratorToken r;
 
@@ -4788,12 +4775,10 @@ jsonb_delete_array(PG_FUNCTION_ARGS)
 			}
 		}
 
-		res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
+		pushJsonbValue(&pstate, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
 	}
 
-	Assert(res != NULL);
-
-	PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
+	PG_RETURN_JSONB_P(JsonbValueToJsonb(pstate.result));
 }
 
 /*
@@ -4808,12 +4793,11 @@ jsonb_delete_idx(PG_FUNCTION_ARGS)
 {
 	Jsonb	   *in = PG_GETARG_JSONB_P(0);
 	int			idx = PG_GETARG_INT32(1);
-	JsonbParseState *state = NULL;
+	JsonbInState pstate = {0};
 	JsonbIterator *it;
 	uint32		i = 0,
 				n;
-	JsonbValue	v,
-			   *res = NULL;
+	JsonbValue	v;
 	JsonbIteratorToken r;
 
 	if (JB_ROOT_IS_SCALAR(in))
@@ -4846,7 +4830,7 @@ jsonb_delete_idx(PG_FUNCTION_ARGS)
 	if (idx >= n)
 		PG_RETURN_JSONB_P(in);
 
-	pushJsonbValue(&state, r, NULL);
+	pushJsonbValue(&pstate, r, NULL);
 
 	while ((r = JsonbIteratorNext(&it, &v, true)) != WJB_DONE)
 	{
@@ -4856,12 +4840,10 @@ jsonb_delete_idx(PG_FUNCTION_ARGS)
 				continue;
 		}
 
-		res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
+		pushJsonbValue(&pstate, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
 	}
 
-	Assert(res != NULL);
-
-	PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
+	PG_RETURN_JSONB_P(JsonbValueToJsonb(pstate.result));
 }
 
 /*
@@ -4875,12 +4857,11 @@ jsonb_set(PG_FUNCTION_ARGS)
 	Jsonb	   *newjsonb = PG_GETARG_JSONB_P(2);
 	JsonbValue	newval;
 	bool		create = PG_GETARG_BOOL(3);
-	JsonbValue *res = NULL;
 	Datum	   *path_elems;
 	bool	   *path_nulls;
 	int			path_len;
 	JsonbIterator *it;
-	JsonbParseState *st = NULL;
+	JsonbInState st = {0};
 
 	JsonbToJsonbValue(newjsonb, &newval);
 
@@ -4904,12 +4885,10 @@ jsonb_set(PG_FUNCTION_ARGS)
 
 	it = JsonbIteratorInit(&in->root);
 
-	res = setPath(&it, path_elems, path_nulls, path_len, &st,
+	setPath(&it, path_elems, path_nulls, path_len, &st,
 			0, &newval, create ? JB_PATH_CREATE : JB_PATH_REPLACE);
 
-	Assert(res != NULL);
-
-	PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
+	PG_RETURN_JSONB_P(JsonbValueToJsonb(st.result));
 }
 
 
@@ -4988,12 +4967,11 @@ jsonb_delete_path(PG_FUNCTION_ARGS)
 {
 	Jsonb	   *in = PG_GETARG_JSONB_P(0);
 	ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
-	JsonbValue *res = NULL;
 	Datum	   *path_elems;
 	bool	   *path_nulls;
 	int			path_len;
 	JsonbIterator *it;
-	JsonbParseState *st = NULL;
+	JsonbInState st = {0};
 
 	if (ARR_NDIM(path) > 1)
 		ereport(ERROR,
@@ -5015,12 +4993,10 @@ jsonb_delete_path(PG_FUNCTION_ARGS)
 
 	it = JsonbIteratorInit(&in->root);
 
-	res = setPath(&it, path_elems, path_nulls, path_len, &st,
+	setPath(&it, path_elems, path_nulls, path_len, &st,
 			0, NULL, JB_PATH_DELETE);
 
-	Assert(res != NULL);
-
-	PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
+	PG_RETURN_JSONB_P(JsonbValueToJsonb(st.result));
 }
 
 /*
@@ -5034,12 +5010,11 @@ jsonb_insert(PG_FUNCTION_ARGS)
 	Jsonb	   *newjsonb = PG_GETARG_JSONB_P(2);
 	JsonbValue	newval;
 	bool		after = PG_GETARG_BOOL(3);
-	JsonbValue *res = NULL;
 	Datum	   *path_elems;
 	bool	   *path_nulls;
 	int			path_len;
 	JsonbIterator *it;
-	JsonbParseState *st = NULL;
+	JsonbInState st = {0};
 
 	JsonbToJsonbValue(newjsonb, &newval);
 
@@ -5060,12 +5035,10 @@ jsonb_insert(PG_FUNCTION_ARGS)
 
 	it = JsonbIteratorInit(&in->root);
 
-	res = setPath(&it, path_elems, path_nulls, path_len, &st, 0, &newval,
+	setPath(&it, path_elems, path_nulls, path_len, &st, 0, &newval,
 			after ? JB_PATH_INSERT_AFTER : JB_PATH_INSERT_BEFORE);
 
-	Assert(res != NULL);
-
-	PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
+	PG_RETURN_JSONB_P(JsonbValueToJsonb(st.result));
 }
 
 /*
@@ -5075,13 +5048,12 @@ jsonb_insert(PG_FUNCTION_ARGS)
  * In that case we just append the content of it2 to it1 without any
  * verifications.
  */
-static JsonbValue *
+static void
 IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
-			   JsonbParseState **state)
+			   JsonbInState *state)
 {
 	JsonbValue	v1,
-				v2,
-			   *res = NULL;
+				v2;
 	JsonbIteratorToken r1,
 				r2,
 				rk1,
@@ -5112,7 +5084,7 @@ IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
 		 * automatically override the value from the first object.
 		 */
 		while ((r2 = JsonbIteratorNext(it2, &v2, true)) != WJB_DONE)
-			res = pushJsonbValue(state, r2, r2 != WJB_END_OBJECT ? &v2 : NULL);
+			pushJsonbValue(state, r2, r2 != WJB_END_OBJECT ? &v2 : NULL);
 	}
 	else if (rk1 == WJB_BEGIN_ARRAY && rk2 == WJB_BEGIN_ARRAY)
 	{
@@ -5133,7 +5105,7 @@ IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
 			pushJsonbValue(state, WJB_ELEM, &v2);
 		}
 
-		res = pushJsonbValue(state, WJB_END_ARRAY, NULL /* signal to sort */ );
+		pushJsonbValue(state, WJB_END_ARRAY, NULL /* signal to sort */ );
 	}
 	else if (rk1 == WJB_BEGIN_OBJECT)
 	{
@@ -5149,7 +5121,7 @@ IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
 			pushJsonbValue(state, r1, r1 != WJB_END_OBJECT ? &v1 : NULL);
 
 		while ((r2 = JsonbIteratorNext(it2, &v2, true)) != WJB_DONE)
-			res = pushJsonbValue(state, r2, r2 != WJB_END_ARRAY ? &v2 : NULL);
+			pushJsonbValue(state, r2, r2 != WJB_END_ARRAY ? &v2 : NULL);
 	}
 	else
 	{
@@ -5168,10 +5140,8 @@ IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
 		while ((r2 = JsonbIteratorNext(it2, &v2, true)) != WJB_DONE)
 			pushJsonbValue(state, r2, r2 != WJB_END_OBJECT ? &v2 : NULL);
 
-		res = pushJsonbValue(state, WJB_END_ARRAY, NULL);
+		pushJsonbValue(state, WJB_END_ARRAY, NULL);
 	}
-
-	return res;
 }
 
 /*
@@ -5203,14 +5173,13 @@ IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
  * All path elements before the last must already exist
  * whatever bits in op_type are set, or nothing is done.
  */
-static JsonbValue *
+static void
 setPath(JsonbIterator **it, const Datum *path_elems,
 		const bool *path_nulls, int path_len,
-		JsonbParseState **st, int level, JsonbValue *newval, int op_type)
+		JsonbInState *st, int level, JsonbValue *newval, int op_type)
 {
 	JsonbValue	v;
 	JsonbIteratorToken r;
-	JsonbValue *res;
 
 	check_stack_depth();
 
@@ -5240,20 +5209,20 @@ setPath(JsonbIterator **it, const Datum *path_elems,
 						 errdetail("The path assumes key is a composite object, "
 								   "but it is a scalar value.")));
 
-			(void) pushJsonbValue(st, r, NULL);
+			pushJsonbValue(st, r, NULL);
 			setPathArray(it, path_elems, path_nulls, path_len, st, level,
 						 newval, v.val.array.nElems, op_type);
 			r = JsonbIteratorNext(it, &v, false);
 			Assert(r == WJB_END_ARRAY);
-			res = pushJsonbValue(st, r, NULL);
+			pushJsonbValue(st, r, NULL);
 			break;
 		case WJB_BEGIN_OBJECT:
-			(void) pushJsonbValue(st, r, NULL);
+			pushJsonbValue(st, r, NULL);
 			setPathObject(it, path_elems, path_nulls, path_len, st, level,
 						  newval, v.val.object.nPairs, op_type);
 			r = JsonbIteratorNext(it, &v, true);
 			Assert(r == WJB_END_OBJECT);
-			res = pushJsonbValue(st, r, NULL);
+			pushJsonbValue(st, r, NULL);
 			break;
 		case WJB_ELEM:
 		case WJB_VALUE:
@@ -5271,15 +5240,12 @@ setPath(JsonbIterator **it, const Datum *path_elems,
 						 errdetail("The path assumes key is a composite object, "
 								   "but it is a scalar value.")));
 
-			res = pushJsonbValue(st, r, &v);
+			pushJsonbValue(st, r, &v);
 			break;
 		default:
 			elog(ERROR, "unrecognized iterator result: %d", (int) r);
-			res = NULL;			/* keep compiler quiet */
 			break;
 	}
-
-	return res;
 }
 
 /*
@@ -5287,7 +5253,7 @@ setPath(JsonbIterator **it, const Datum *path_elems,
  */
 static void
 setPathObject(JsonbIterator **it, const Datum *path_elems, const bool *path_nulls,
-			  int path_len, JsonbParseState **st, int level,
+			  int path_len, JsonbInState *st, int level,
 			  JsonbValue *newval, uint32 npairs, int op_type)
 {
 	text	   *pathelem = NULL;
@@ -5314,8 +5280,8 @@ setPathObject(JsonbIterator **it, const Datum *path_elems, const bool *path_null
 		newkey.val.string.val = VARDATA_ANY(pathelem);
 		newkey.val.string.len = VARSIZE_ANY_EXHDR(pathelem);
 
-		(void) pushJsonbValue(st, WJB_KEY, &newkey);
-		(void) pushJsonbValue(st, WJB_VALUE, newval);
+		pushJsonbValue(st, WJB_KEY, &newkey);
+		pushJsonbValue(st, WJB_VALUE, newval);
 	}
 
 	for (i = 0; i < npairs; i++)
@@ -5347,13 +5313,13 @@ setPathObject(JsonbIterator **it, const Datum *path_elems, const bool *path_null
 				r = JsonbIteratorNext(it, &v, true);	/* skip value */
 				if (!(op_type & JB_PATH_DELETE))
 				{
-					(void) pushJsonbValue(st, WJB_KEY, &k);
-					(void) pushJsonbValue(st, WJB_VALUE, newval);
+					pushJsonbValue(st, WJB_KEY, &k);
+					pushJsonbValue(st, WJB_VALUE, newval);
 				}
 			}
 			else
 			{
-				(void) pushJsonbValue(st, r, &k);
+				pushJsonbValue(st, r, &k);
 				setPath(it, path_elems, path_nulls, path_len,
 						st, level + 1, newval, op_type);
 			}
@@ -5369,13 +5335,13 @@ setPathObject(JsonbIterator **it, const Datum *path_elems, const bool *path_null
 				newkey.val.string.val = VARDATA_ANY(pathelem);
 				newkey.val.string.len = VARSIZE_ANY_EXHDR(pathelem);
 
-				(void) pushJsonbValue(st, WJB_KEY, &newkey);
-				(void) pushJsonbValue(st, WJB_VALUE, newval);
+				pushJsonbValue(st, WJB_KEY, &newkey);
+				pushJsonbValue(st, WJB_VALUE, newval);
 			}
 
-			(void) pushJsonbValue(st, r, &k);
+			pushJsonbValue(st, r, &k);
 			r = JsonbIteratorNext(it, &v, false);
-			(void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
+			pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
 			if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
 			{
 				int			walking_level = 1;
@@ -5389,7 +5355,7 @@ setPathObject(JsonbIterator **it, const Datum *path_elems, const bool *path_null
 					if (r == WJB_END_ARRAY || r == WJB_END_OBJECT)
 						--walking_level;
 
-					(void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
+					pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
 				}
 			}
 		}
@@ -5413,7 +5379,7 @@ setPathObject(JsonbIterator **it, const Datum *path_elems, const bool *path_null
 		newkey.val.string.val = VARDATA_ANY(pathelem);
 		newkey.val.string.len = VARSIZE_ANY_EXHDR(pathelem);
 
-		(void) pushJsonbValue(st, WJB_KEY, &newkey);
+		pushJsonbValue(st, WJB_KEY, &newkey);
 		push_path(st, level, path_elems, path_nulls, path_len, newval);
 
 		/* Result is closed with WJB_END_OBJECT outside of this function */
@@ -5425,7 +5391,7 @@ setPathObject(JsonbIterator **it, const Datum *path_elems, const bool *path_null
  */
 static void
 setPathArray(JsonbIterator **it, const Datum *path_elems, const bool *path_nulls,
-			 int path_len, JsonbParseState **st, int level,
+			 int path_len, JsonbInState *st, int level,
 			 JsonbValue *newval, uint32 nelems, int op_type)
 {
 	JsonbValue	v;
@@ -5493,7 +5459,7 @@ setPathArray(JsonbIterator **it, const Datum *path_elems, const bool *path_nulls
 		if (op_type & JB_PATH_FILL_GAPS && nelems == 0 && idx > 0)
 			push_null_elements(st, idx);
 
-		(void) pushJsonbValue(st, WJB_ELEM, newval);
+		pushJsonbValue(st, WJB_ELEM, newval);
 
 		done = true;
 	}
@@ -5512,7 +5478,7 @@ setPathArray(JsonbIterator **it, const Datum *path_elems, const bool *path_nulls
 				r = JsonbIteratorNext(it, &v, true);	/* skip */
 
 				if (op_type & (JB_PATH_INSERT_BEFORE | JB_PATH_CREATE))
-					(void) pushJsonbValue(st, WJB_ELEM, newval);
+					pushJsonbValue(st, WJB_ELEM, newval);
 
 				/*
 				 * We should keep current value only in case of
@@ -5520,20 +5486,20 @@ setPathArray(JsonbIterator **it, const Datum *path_elems, const bool *path_nulls
 				 * otherwise it should be deleted or replaced
 				 */
 				if (op_type & (JB_PATH_INSERT_AFTER | JB_PATH_INSERT_BEFORE))
-					(void) pushJsonbValue(st, r, &v);
+					pushJsonbValue(st, r, &v);
 
 				if (op_type & (JB_PATH_INSERT_AFTER | JB_PATH_REPLACE))
-					(void) pushJsonbValue(st, WJB_ELEM, newval);
+					pushJsonbValue(st, WJB_ELEM, newval);
 			}
 			else
-				(void) setPath(it, path_elems, path_nulls, path_len,
+				setPath(it, path_elems, path_nulls, path_len,
 						st, level + 1, newval, op_type);
 		}
 		else
 		{
 			r = JsonbIteratorNext(it, &v, false);
 
-			(void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
+			pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
 
 			if (r == WJB_BEGIN_ARRAY || r == WJB_BEGIN_OBJECT)
 			{
@@ -5548,7 +5514,7 @@ setPathArray(JsonbIterator **it, const Datum *path_elems, const bool *path_nulls
 					if (r == WJB_END_ARRAY || r == WJB_END_OBJECT)
 						--walking_level;
 
-					(void) pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
+					pushJsonbValue(st, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
 				}
 			}
 		}
@@ -5563,7 +5529,7 @@ setPathArray(JsonbIterator **it, const Datum *path_elems, const bool *path_nulls
 		if (op_type & JB_PATH_FILL_GAPS && idx > nelems)
 			push_null_elements(st, idx - nelems);
 
-		(void) pushJsonbValue(st, WJB_ELEM, newval);
+		pushJsonbValue(st, WJB_ELEM, newval);
 		done = true;
 	}
 
@@ -5808,10 +5774,9 @@ transform_jsonb_string_values(Jsonb *jsonb, void *action_state,
 							  JsonTransformStringValuesAction transform_action)
 {
 	JsonbIterator *it;
-	JsonbValue	v,
-			   *res = NULL;
+	JsonbValue	v;
 	JsonbIteratorToken type;
-	JsonbParseState *st = NULL;
+	JsonbInState st = {0};
 	text	   *out;
 	bool		is_scalar = false;
 
@@ -5827,20 +5792,20 @@ transform_jsonb_string_values(Jsonb *jsonb, void *action_state,
 			out = pg_detoast_datum_packed(out);
 			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);
+			pushJsonbValue(&st, type, type < WJB_BEGIN_ARRAY ? &v : NULL);
 		}
 		else
 		{
-			res = pushJsonbValue(&st, type, (type == WJB_KEY ||
+			pushJsonbValue(&st, type, (type == WJB_KEY ||
 									   type == WJB_VALUE ||
 									   type == WJB_ELEM) ? &v : NULL);
 		}
 	}
 
-	if (res->type == jbvArray)
-		res->val.array.rawScalar = is_scalar;
+	if (st.result->type == jbvArray)
+		st.result->val.array.rawScalar = is_scalar;
 
-	return JsonbValueToJsonb(res);
+	return JsonbValueToJsonb(st.result);
 }
 
 /*

src/backend/utils/adt/jsonpath_exec.c

@@ -2874,8 +2874,7 @@ executeKeyValueMethod(JsonPathExecContext *cxt, JsonPathItem *jsp,
 	{
 		JsonBaseObjectInfo baseObject;
 		JsonbValue	obj;
-		JsonbParseState *ps;
-		JsonbValue *keyval;
+		JsonbInState ps;
 		Jsonb	   *jsonb;
 
 		if (tok != WJB_KEY)
@@ -2889,7 +2888,8 @@ executeKeyValueMethod(JsonPathExecContext *cxt, JsonPathItem *jsp,
 		tok = JsonbIteratorNext(&it, &val, true);
 		Assert(tok == WJB_VALUE);
 
-		ps = NULL;
+		memset(&ps, 0, sizeof(ps));
+
 		pushJsonbValue(&ps, WJB_BEGIN_OBJECT, NULL);
 
 		pushJsonbValue(&ps, WJB_KEY, &keystr);
@@ -2901,9 +2901,9 @@ executeKeyValueMethod(JsonPathExecContext *cxt, JsonPathItem *jsp,
 		pushJsonbValue(&ps, WJB_KEY, &idstr);
 		pushJsonbValue(&ps, WJB_VALUE, &idval);
 
-		keyval = pushJsonbValue(&ps, WJB_END_OBJECT, NULL);
+		pushJsonbValue(&ps, WJB_END_OBJECT, NULL);
 
-		jsonb = JsonbValueToJsonb(keyval);
+		jsonb = JsonbValueToJsonb(ps.result);
 
 		JsonbInitBinary(&obj, jsonb);
 
@@ -3649,7 +3649,7 @@ getScalar(JsonbValue *scalar, enum jbvType type)
 static JsonbValue *
 wrapItemsInArray(const JsonValueList *items)
 {
-	JsonbParseState *ps = NULL;
+	JsonbInState ps = {0};
 	JsonValueListIterator it;
 	JsonbValue *jbv;
 
@@ -3659,7 +3659,9 @@ wrapItemsInArray(const JsonValueList *items)
 	while ((jbv = JsonValueListNext(items, &it)))
 		pushJsonbValue(&ps, WJB_ELEM, jbv);
 
-	return pushJsonbValue(&ps, WJB_END_ARRAY, NULL);
+	pushJsonbValue(&ps, WJB_END_ARRAY, NULL);
+
+	return ps.result;
 }
 
 /* Check if the timezone required for casting from type1 to type2 is used */

src/include/utils/jsonb.h

@@ -67,8 +67,10 @@ typedef enum
 #define JGINFLAG_HASHED 0x10	/* OR'd into flag if value was hashed */
 #define JGIN_MAXLENGTH	125		/* max length of text part before hashing */
 
+/* Forward struct references */
 typedef struct JsonbPair JsonbPair;
 typedef struct JsonbValue JsonbValue;
+typedef struct JsonbParseState JsonbParseState;
 
 /*
  * Jsonbs are varlena objects, so must meet the varlena convention that the
@@ -315,15 +317,40 @@ struct JsonbPair
 	uint32		order;			/* Pair's index in original sequence */
 };
 
-/* Conversion state used when parsing Jsonb from text, or for type coercion */
-typedef struct JsonbParseState
+/*
+ * State used while constructing or manipulating a JsonbValue.
+ * For example, when parsing Jsonb from text, we construct a JsonbValue
+ * data structure and then flatten that into the Jsonb on-disk format.
+ * JsonbValues are also useful in aggregation and type coercion.
+ *
+ * Callers providing a JsonbInState must initialize it to zeroes/nulls,
+ * except for optionally setting outcontext (if that's left NULL,
+ * CurrentMemoryContext is used) and escontext (if that's left NULL,
+ * parsing errors are thrown via ereport).
+ */
+typedef struct JsonbInState
 {
-	JsonbValue	contVal;
-	Size		size;
-	struct JsonbParseState *next;
+	JsonbValue *result;			/* The completed value; NULL until complete */
+	MemoryContext outcontext;	/* The context to build it in, or NULL */
+	struct Node *escontext;		/* Optional soft-error-reporting context */
+	/* Remaining fields should be treated as private to jsonb.c/jsonb_util.c */
+	JsonbParseState *parseState;	/* Stack of parsing contexts */
+	bool		unique_keys;	/* Check object key uniqueness */
+} JsonbInState;
+
+/*
+ * Parsing context for one level of Jsonb array or object nesting.
+ * The contVal will be part of the constructed JsonbValue tree,
+ * but the other fields are just transient state.
+ */
+struct JsonbParseState
+{
+	JsonbValue	contVal;		/* An array or object JsonbValue */
+	Size		size;			/* Allocated length of array or object */
+	JsonbParseState *next;		/* Link to next outer level, if any */
 	bool		unique_keys;	/* Check object key uniqueness */
 	bool		skip_nulls;		/* Skip null object fields */
-} JsonbParseState;
+};
 
 /*
  * JsonbIterator holds details of the type for each iteration. It also stores a
@@ -404,7 +431,7 @@ extern JsonbValue *getKeyJsonValueFromContainer(JsonbContainer *container,
 												JsonbValue *res);
 extern JsonbValue *getIthJsonbValueFromContainer(JsonbContainer *container,
 												 uint32 i);
-extern JsonbValue *pushJsonbValue(JsonbParseState **pstate,
+extern void pushJsonbValue(JsonbInState *pstate,
 						   JsonbIteratorToken seq, JsonbValue *jbval);
 extern JsonbIterator *JsonbIteratorInit(JsonbContainer *container);
 extern JsonbIteratorToken JsonbIteratorNext(JsonbIterator **it, JsonbValue *val,