From df8b71428481f718ac6bb60a3ca5969a6876da7e Mon Sep 17 00:00:00 2001
From: Alvaro Herrera <alvherre@alvh.no-ip.org>
Date: Thu, 3 Aug 2023 11:44:15 +0200
Subject: [PATCH] JsonLexContext allocation/free

---
 src/backend/utils/adt/json.c             |  38 ++++----
 src/backend/utils/adt/jsonb.c            |  13 +--
 src/backend/utils/adt/jsonfuncs.c        | 106 +++++++++++++----------
 src/bin/pg_verifybackup/parse_manifest.c |   4 +-
 src/common/jsonapi.c                     |  43 +++++++--
 src/include/common/jsonapi.h             |  23 +++--
 src/include/utils/jsonfuncs.h            |   2 +-
 7 files changed, 146 insertions(+), 83 deletions(-)

diff --git a/src/backend/utils/adt/json.c b/src/backend/utils/adt/json.c
index e405791f5d..27f9a51228 100644
--- a/src/backend/utils/adt/json.c
+++ b/src/backend/utils/adt/json.c
@@ -106,11 +106,11 @@ json_in(PG_FUNCTION_ARGS)
 {
 	char	   *json = PG_GETARG_CSTRING(0);
 	text	   *result = cstring_to_text(json);
-	JsonLexContext *lex;
+	JsonLexContext lex;
 
 	/* validate it */
-	lex = makeJsonLexContext(result, false);
-	if (!pg_parse_json_or_errsave(lex, &nullSemAction, fcinfo->context))
+	makeJsonLexContext(&lex, result, false);
+	if (!pg_parse_json_or_errsave(&lex, &nullSemAction, fcinfo->context))
 		PG_RETURN_NULL();
 
 	/* Internal representation is the same as text */
@@ -152,13 +152,13 @@ json_recv(PG_FUNCTION_ARGS)
 	StringInfo	buf = (StringInfo) PG_GETARG_POINTER(0);
 	char	   *str;
 	int			nbytes;
-	JsonLexContext *lex;
+	JsonLexContext lex;
 
 	str = pq_getmsgtext(buf, buf->len - buf->cursor, &nbytes);
 
 	/* Validate it. */
-	lex = makeJsonLexContextCstringLen(str, nbytes, GetDatabaseEncoding(), false);
-	pg_parse_json_or_ereport(lex, &nullSemAction);
+	makeJsonLexContextCstringLen(&lex, str, nbytes, GetDatabaseEncoding(), false);
+	pg_parse_json_or_ereport(&lex, &nullSemAction);
 
 	PG_RETURN_TEXT_P(cstring_to_text_with_len(str, nbytes));
 }
@@ -1625,14 +1625,16 @@ json_unique_object_field_start(void *_state, char *field, bool isnull)
 bool
 json_validate(text *json, bool check_unique_keys, bool throw_error)
 {
-	JsonLexContext *lex = makeJsonLexContext(json, check_unique_keys);
+	JsonLexContext lex;
 	JsonSemAction uniqueSemAction = {0};
 	JsonUniqueParsingState state;
 	JsonParseErrorType result;
 
+	makeJsonLexContext(&lex, json, check_unique_keys);
+
 	if (check_unique_keys)
 	{
-		state.lex = lex;
+		state.lex = &lex;
 		state.stack = NULL;
 		state.id_counter = 0;
 		state.unique = true;
@@ -1644,12 +1646,12 @@ json_validate(text *json, bool check_unique_keys, bool throw_error)
 		uniqueSemAction.object_end = json_unique_object_end;
 	}
 
-	result = pg_parse_json(lex, check_unique_keys ? &uniqueSemAction : &nullSemAction);
+	result = pg_parse_json(&lex, check_unique_keys ? &uniqueSemAction : &nullSemAction);
 
 	if (result != JSON_SUCCESS)
 	{
 		if (throw_error)
-			json_errsave_error(result, lex, NULL);
+			json_errsave_error(result, &lex, NULL);
 
 		return false;			/* invalid json */
 	}
@@ -1664,6 +1666,9 @@ json_validate(text *json, bool check_unique_keys, bool throw_error)
 		return false;			/* not unique keys */
 	}
 
+	if (check_unique_keys)
+		freeJsonLexContext(&lex);
+
 	return true;				/* ok */
 }
 
@@ -1683,18 +1688,17 @@ Datum
 json_typeof(PG_FUNCTION_ARGS)
 {
 	text	   *json = PG_GETARG_TEXT_PP(0);
-	JsonLexContext *lex = makeJsonLexContext(json, false);
+	JsonLexContext lex;
 	char	   *type;
-	JsonTokenType tok;
 	JsonParseErrorType result;
 
 	/* Lex exactly one token from the input and check its type. */
-	result = json_lex(lex);
+	makeJsonLexContext(&lex, json, false);
+	result = json_lex(&lex);
 	if (result != JSON_SUCCESS)
-		json_errsave_error(result, lex, NULL);
-	tok = lex->token_type;
+		json_errsave_error(result, &lex, NULL);
 
-	switch (tok)
+	switch (lex.token_type)
 	{
 		case JSON_TOKEN_OBJECT_START:
 			type = "object";
@@ -1716,7 +1720,7 @@ json_typeof(PG_FUNCTION_ARGS)
 			type = "null";
 			break;
 		default:
-			elog(ERROR, "unexpected json token: %d", tok);
+			elog(ERROR, "unexpected json token: %d", lex.token_type);
 	}
 
 	PG_RETURN_TEXT_P(cstring_to_text(type));
diff --git a/src/backend/utils/adt/jsonb.c b/src/backend/utils/adt/jsonb.c
index 9781852b0c..b10a60ac66 100644
--- a/src/backend/utils/adt/jsonb.c
+++ b/src/backend/utils/adt/jsonb.c
@@ -252,13 +252,13 @@ jsonb_typeof(PG_FUNCTION_ARGS)
 static inline Datum
 jsonb_from_cstring(char *json, int len, bool unique_keys, Node *escontext)
 {
-	JsonLexContext *lex;
+	JsonLexContext lex;
 	JsonbInState state;
 	JsonSemAction sem;
 
 	memset(&state, 0, sizeof(state));
 	memset(&sem, 0, sizeof(sem));
-	lex = makeJsonLexContextCstringLen(json, len, GetDatabaseEncoding(), true);
+	makeJsonLexContextCstringLen(&lex, json, len, GetDatabaseEncoding(), true);
 
 	state.unique_keys = unique_keys;
 	state.escontext = escontext;
@@ -271,7 +271,7 @@ jsonb_from_cstring(char *json, int len, bool unique_keys, Node *escontext)
 	sem.scalar = jsonb_in_scalar;
 	sem.object_field_start = jsonb_in_object_field_start;
 
-	if (!pg_parse_json_or_errsave(lex, &sem, escontext))
+	if (!pg_parse_json_or_errsave(&lex, &sem, escontext))
 		return (Datum) 0;
 
 	/* after parsing, the item member has the composed jsonb structure */
@@ -755,11 +755,11 @@ datum_to_jsonb_internal(Datum val, bool is_null, JsonbInState *result,
 			case JSONTYPE_JSON:
 				{
 					/* parse the json right into the existing result object */
-					JsonLexContext *lex;
+					JsonLexContext lex;
 					JsonSemAction sem;
 					text	   *json = DatumGetTextPP(val);
 
-					lex = makeJsonLexContext(json, true);
+					makeJsonLexContext(&lex, json, true);
 
 					memset(&sem, 0, sizeof(sem));
 
@@ -772,7 +772,8 @@ datum_to_jsonb_internal(Datum val, bool is_null, JsonbInState *result,
 					sem.scalar = jsonb_in_scalar;
 					sem.object_field_start = jsonb_in_object_field_start;
 
-					pg_parse_json_or_ereport(lex, &sem);
+					pg_parse_json_or_ereport(&lex, &sem);
+					freeJsonLexContext(&lex);
 				}
 				break;
 			case JSONTYPE_JSONB:
diff --git a/src/backend/utils/adt/jsonfuncs.c b/src/backend/utils/adt/jsonfuncs.c
index a4bfa5e404..3f855d8f2b 100644
--- a/src/backend/utils/adt/jsonfuncs.c
+++ b/src/backend/utils/adt/jsonfuncs.c
@@ -526,7 +526,7 @@ pg_parse_json_or_errsave(JsonLexContext *lex, JsonSemAction *sem,
  * directly.
  */
 JsonLexContext *
-makeJsonLexContext(text *json, bool need_escapes)
+makeJsonLexContext(JsonLexContext *lex, text *json, bool need_escapes)
 {
 	/*
 	 * Most callers pass a detoasted datum, but it's not clear that they all
@@ -534,7 +534,8 @@ makeJsonLexContext(text *json, bool need_escapes)
 	 */
 	json = pg_detoast_datum_packed(json);
 
-	return makeJsonLexContextCstringLen(VARDATA_ANY(json),
+	return makeJsonLexContextCstringLen(lex,
+										VARDATA_ANY(json),
 										VARSIZE_ANY_EXHDR(json),
 										GetDatabaseEncoding(),
 										need_escapes);
@@ -725,17 +726,19 @@ json_object_keys(PG_FUNCTION_ARGS)
 	if (SRF_IS_FIRSTCALL())
 	{
 		text	   *json = PG_GETARG_TEXT_PP(0);
-		JsonLexContext *lex = makeJsonLexContext(json, true);
+		JsonLexContext lex;
 		JsonSemAction *sem;
 		MemoryContext oldcontext;
 
+		makeJsonLexContext(&lex, json, true);
+
 		funcctx = SRF_FIRSTCALL_INIT();
 		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
 
 		state = palloc(sizeof(OkeysState));
 		sem = palloc0(sizeof(JsonSemAction));
 
-		state->lex = lex;
+		state->lex = &lex;
 		state->result_size = 256;
 		state->result_count = 0;
 		state->sent_count = 0;
@@ -747,12 +750,10 @@ json_object_keys(PG_FUNCTION_ARGS)
 		sem->object_field_start = okeys_object_field_start;
 		/* remainder are all NULL, courtesy of palloc0 above */
 
-		pg_parse_json_or_ereport(lex, sem);
+		pg_parse_json_or_ereport(&lex, sem);
 		/* keys are now in state->result */
 
-		pfree(lex->strval->data);
-		pfree(lex->strval);
-		pfree(lex);
+		freeJsonLexContext(&lex);
 		pfree(sem);
 
 		MemoryContextSwitchTo(oldcontext);
@@ -1096,13 +1097,13 @@ get_worker(text *json,
 		   int npath,
 		   bool normalize_results)
 {
-	JsonLexContext *lex = makeJsonLexContext(json, true);
 	JsonSemAction *sem = palloc0(sizeof(JsonSemAction));
 	GetState   *state = palloc0(sizeof(GetState));
 
 	Assert(npath >= 0);
 
-	state->lex = lex;
+	state->lex = makeJsonLexContext(NULL, json, true);
+
 	/* is it "_as_text" variant? */
 	state->normalize_results = normalize_results;
 	state->npath = npath;
@@ -1140,7 +1141,7 @@ get_worker(text *json,
 		sem->array_element_end = get_array_element_end;
 	}
 
-	pg_parse_json_or_ereport(lex, sem);
+	pg_parse_json_or_ereport(state->lex, sem);
 
 	return state->tresult;
 }
@@ -1842,25 +1843,24 @@ json_array_length(PG_FUNCTION_ARGS)
 {
 	text	   *json = PG_GETARG_TEXT_PP(0);
 	AlenState  *state;
-	JsonLexContext *lex;
 	JsonSemAction *sem;
+	JsonLexContext lex;
 
-	lex = makeJsonLexContext(json, false);
 	state = palloc0(sizeof(AlenState));
 	sem = palloc0(sizeof(JsonSemAction));
 
+	state->lex = makeJsonLexContext(&lex, json, false);
 	/* palloc0 does this for us */
 #if 0
 	state->count = 0;
 #endif
-	state->lex = lex;
 
 	sem->semstate = (void *) state;
 	sem->object_start = alen_object_start;
 	sem->scalar = alen_scalar;
 	sem->array_element_start = alen_array_element_start;
 
-	pg_parse_json_or_ereport(lex, sem);
+	pg_parse_json_or_ereport(state->lex, sem);
 
 	PG_RETURN_INT32(state->count);
 }
@@ -2049,12 +2049,12 @@ static Datum
 each_worker(FunctionCallInfo fcinfo, bool as_text)
 {
 	text	   *json = PG_GETARG_TEXT_PP(0);
-	JsonLexContext *lex;
+	JsonLexContext lex;
 	JsonSemAction *sem;
 	ReturnSetInfo *rsi;
 	EachState  *state;
 
-	lex = makeJsonLexContext(json, true);
+	makeJsonLexContext(&lex, json, true);
 	state = palloc0(sizeof(EachState));
 	sem = palloc0(sizeof(JsonSemAction));
 
@@ -2072,12 +2072,12 @@ each_worker(FunctionCallInfo fcinfo, bool as_text)
 
 	state->normalize_results = as_text;
 	state->next_scalar = false;
-	state->lex = lex;
+	state->lex = &lex;
 	state->tmp_cxt = AllocSetContextCreate(CurrentMemoryContext,
 										   "json_each temporary cxt",
 										   ALLOCSET_DEFAULT_SIZES);
 
-	pg_parse_json_or_ereport(lex, sem);
+	pg_parse_json_or_ereport(&lex, sem);
 
 	MemoryContextDelete(state->tmp_cxt);
 
@@ -2299,13 +2299,14 @@ static Datum
 elements_worker(FunctionCallInfo fcinfo, const char *funcname, bool as_text)
 {
 	text	   *json = PG_GETARG_TEXT_PP(0);
-
-	/* elements only needs escaped strings when as_text */
-	JsonLexContext *lex = makeJsonLexContext(json, as_text);
+	JsonLexContext lex;
 	JsonSemAction *sem;
 	ReturnSetInfo *rsi;
 	ElementsState *state;
 
+	/* elements only needs escaped strings when as_text */
+	makeJsonLexContext(&lex, json, as_text);
+
 	state = palloc0(sizeof(ElementsState));
 	sem = palloc0(sizeof(JsonSemAction));
 
@@ -2323,12 +2324,12 @@ elements_worker(FunctionCallInfo fcinfo, const char *funcname, bool as_text)
 	state->function_name = funcname;
 	state->normalize_results = as_text;
 	state->next_scalar = false;
-	state->lex = lex;
+	state->lex = &lex;
 	state->tmp_cxt = AllocSetContextCreate(CurrentMemoryContext,
 										   "json_array_elements temporary cxt",
 										   ALLOCSET_DEFAULT_SIZES);
 
-	pg_parse_json_or_ereport(lex, sem);
+	pg_parse_json_or_ereport(&lex, sem);
 
 	MemoryContextDelete(state->tmp_cxt);
 
@@ -2704,7 +2705,8 @@ populate_array_json(PopulateArrayContext *ctx, char *json, int len)
 	PopulateArrayState state;
 	JsonSemAction sem;
 
-	state.lex = makeJsonLexContextCstringLen(json, len, GetDatabaseEncoding(), true);
+	state.lex = makeJsonLexContextCstringLen(NULL, json, len,
+											 GetDatabaseEncoding(), true);
 	state.ctx = ctx;
 
 	memset(&sem, 0, sizeof(sem));
@@ -2720,7 +2722,7 @@ populate_array_json(PopulateArrayContext *ctx, char *json, int len)
 	/* number of dimensions should be already known */
 	Assert(ctx->ndims > 0 && ctx->dims);
 
-	pfree(state.lex);
+	freeJsonLexContext(state.lex);
 }
 
 /*
@@ -3547,7 +3549,6 @@ get_json_object_as_hash(char *json, int len, const char *funcname)
 	HASHCTL		ctl;
 	HTAB	   *tab;
 	JHashState *state;
-	JsonLexContext *lex = makeJsonLexContextCstringLen(json, len, GetDatabaseEncoding(), true);
 	JsonSemAction *sem;
 
 	ctl.keysize = NAMEDATALEN;
@@ -3563,7 +3564,8 @@ get_json_object_as_hash(char *json, int len, const char *funcname)
 
 	state->function_name = funcname;
 	state->hash = tab;
-	state->lex = lex;
+	state->lex = makeJsonLexContextCstringLen(NULL, json, len,
+											  GetDatabaseEncoding(), true);
 
 	sem->semstate = (void *) state;
 	sem->array_start = hash_array_start;
@@ -3571,7 +3573,9 @@ get_json_object_as_hash(char *json, int len, const char *funcname)
 	sem->object_field_start = hash_object_field_start;
 	sem->object_field_end = hash_object_field_end;
 
-	pg_parse_json_or_ereport(lex, sem);
+	pg_parse_json_or_ereport(state->lex, sem);
+
+	freeJsonLexContext(state->lex);
 
 	return tab;
 }
@@ -3863,12 +3867,12 @@ populate_recordset_worker(FunctionCallInfo fcinfo, const char *funcname,
 	if (is_json)
 	{
 		text	   *json = PG_GETARG_TEXT_PP(json_arg_num);
-		JsonLexContext *lex;
+		JsonLexContext lex;
 		JsonSemAction *sem;
 
 		sem = palloc0(sizeof(JsonSemAction));
 
-		lex = makeJsonLexContext(json, true);
+		makeJsonLexContext(&lex, json, true);
 
 		sem->semstate = (void *) state;
 		sem->array_start = populate_recordset_array_start;
@@ -3879,9 +3883,12 @@ populate_recordset_worker(FunctionCallInfo fcinfo, const char *funcname,
 		sem->object_start = populate_recordset_object_start;
 		sem->object_end = populate_recordset_object_end;
 
-		state->lex = lex;
+		state->lex = &lex;
 
-		pg_parse_json_or_ereport(lex, sem);
+		pg_parse_json_or_ereport(&lex, sem);
+
+		freeJsonLexContext(&lex);
+		state->lex = NULL;
 	}
 	else
 	{
@@ -4217,16 +4224,16 @@ json_strip_nulls(PG_FUNCTION_ARGS)
 {
 	text	   *json = PG_GETARG_TEXT_PP(0);
 	StripnullState *state;
-	JsonLexContext *lex;
+	JsonLexContext lex;
 	JsonSemAction *sem;
 
-	lex = makeJsonLexContext(json, true);
+	makeJsonLexContext(&lex, json, true);
 	state = palloc0(sizeof(StripnullState));
 	sem = palloc0(sizeof(JsonSemAction));
 
 	state->strval = makeStringInfo();
 	state->skip_next_null = false;
-	state->lex = lex;
+	state->lex = &lex;
 
 	sem->semstate = (void *) state;
 	sem->object_start = sn_object_start;
@@ -4237,7 +4244,7 @@ json_strip_nulls(PG_FUNCTION_ARGS)
 	sem->array_element_start = sn_array_element_start;
 	sem->object_field_start = sn_object_field_start;
 
-	pg_parse_json_or_ereport(lex, sem);
+	pg_parse_json_or_ereport(&lex, sem);
 
 	PG_RETURN_TEXT_P(cstring_to_text_with_len(state->strval->data,
 											  state->strval->len));
@@ -5433,11 +5440,13 @@ void
 iterate_json_values(text *json, uint32 flags, void *action_state,
 					JsonIterateStringValuesAction action)
 {
-	JsonLexContext *lex = makeJsonLexContext(json, true);
+	JsonLexContext lex;
 	JsonSemAction *sem = palloc0(sizeof(JsonSemAction));
 	IterateJsonStringValuesState *state = palloc0(sizeof(IterateJsonStringValuesState));
 
-	state->lex = lex;
+	makeJsonLexContext(&lex, json, true);
+
+	state->lex = &lex;
 	state->action = action;
 	state->action_state = action_state;
 	state->flags = flags;
@@ -5446,7 +5455,7 @@ iterate_json_values(text *json, uint32 flags, void *action_state,
 	sem->scalar = iterate_values_scalar;
 	sem->object_field_start = iterate_values_object_field_start;
 
-	pg_parse_json_or_ereport(lex, sem);
+	pg_parse_json_or_ereport(&lex, sem);
 }
 
 /*
@@ -5553,11 +5562,12 @@ text *
 transform_json_string_values(text *json, void *action_state,
 							 JsonTransformStringValuesAction transform_action)
 {
-	JsonLexContext *lex = makeJsonLexContext(json, true);
+	JsonLexContext lex;
 	JsonSemAction *sem = palloc0(sizeof(JsonSemAction));
 	TransformJsonStringValuesState *state = palloc0(sizeof(TransformJsonStringValuesState));
 
-	state->lex = lex;
+	makeJsonLexContext(&lex, json, true);
+	state->lex = &lex;
 	state->strval = makeStringInfo();
 	state->action = transform_action;
 	state->action_state = action_state;
@@ -5571,7 +5581,7 @@ transform_json_string_values(text *json, void *action_state,
 	sem->array_element_start = transform_string_values_array_element_start;
 	sem->object_field_start = transform_string_values_object_field_start;
 
-	pg_parse_json_or_ereport(lex, sem);
+	pg_parse_json_or_ereport(&lex, sem);
 
 	return cstring_to_text_with_len(state->strval->data, state->strval->len);
 }
@@ -5670,19 +5680,19 @@ transform_string_values_scalar(void *state, char *token, JsonTokenType tokentype
 JsonTokenType
 json_get_first_token(text *json, bool throw_error)
 {
-	JsonLexContext *lex;
+	JsonLexContext lex;
 	JsonParseErrorType result;
 
-	lex = makeJsonLexContext(json, false);
+	makeJsonLexContext(&lex, json, false);
 
 	/* Lex exactly one token from the input and check its type. */
-	result = json_lex(lex);
+	result = json_lex(&lex);
 
 	if (result == JSON_SUCCESS)
-		return lex->token_type;
+		return lex.token_type;
 
 	if (throw_error)
-		json_errsave_error(result, lex, NULL);
+		json_errsave_error(result, &lex, NULL);
 
 	return JSON_TOKEN_INVALID;	/* invalid json */
 }
diff --git a/src/bin/pg_verifybackup/parse_manifest.c b/src/bin/pg_verifybackup/parse_manifest.c
index 2379f7be7b..f0acd9f1e7 100644
--- a/src/bin/pg_verifybackup/parse_manifest.c
+++ b/src/bin/pg_verifybackup/parse_manifest.c
@@ -130,7 +130,7 @@ json_parse_manifest(JsonManifestParseContext *context, char *buffer,
 	parse.saw_version_field = false;
 
 	/* Create a JSON lexing context. */
-	lex = makeJsonLexContextCstringLen(buffer, size, PG_UTF8, true);
+	lex = makeJsonLexContextCstringLen(NULL, buffer, size, PG_UTF8, true);
 
 	/* Set up semantic actions. */
 	sem.semstate = &parse;
@@ -153,6 +153,8 @@ json_parse_manifest(JsonManifestParseContext *context, char *buffer,
 
 	/* Verify the manifest checksum. */
 	verify_manifest_checksum(&parse, buffer, size);
+
+	freeJsonLexContext(lex);
 }
 
 /*
diff --git a/src/common/jsonapi.c b/src/common/jsonapi.c
index 2e86589cfd..e30d8491c9 100644
--- a/src/common/jsonapi.c
+++ b/src/common/jsonapi.c
@@ -135,26 +135,59 @@ IsValidJsonNumber(const char *str, int len)
 
 /*
  * makeJsonLexContextCstringLen
+ *		Initialize the given JsonLexContext object, or create one
  *
- * lex constructor, with or without StringInfo object for de-escaped lexemes.
+ * If a valid 'lex' pointer is given, it is initialized.  This can
+ * be used for stack-allocated structs, saving overhead.  Otherwise,
+ * one is allocated.
  *
- * Without is better as it makes the processing faster, so only make one
- * if really required.
+ * If need_escapes is true, ->strval stores the unescaped lexemes.
+ * Unescaping is expensive, so only request it when necessary.
+ *
+ * If either need_escapes or lex was given as NULL, then caller
+ * is responsible for freeing the object, either by calling
+ * freeJsonLexContext() or via memory context cleanup.
  */
 JsonLexContext *
-makeJsonLexContextCstringLen(char *json, int len, int encoding, bool need_escapes)
+makeJsonLexContextCstringLen(JsonLexContext *lex, char *json,
+							 int len, int encoding, bool need_escapes)
 {
-	JsonLexContext *lex = palloc0(sizeof(JsonLexContext));
+	if (lex == NULL)
+	{
+		lex = palloc0(sizeof(JsonLexContext));
+		lex->flags |= JSONLEX_FREE_STRUCT;
+	}
+	else
+		memset(lex, 0, sizeof(JsonLexContext));
 
 	lex->input = lex->token_terminator = lex->line_start = json;
 	lex->line_number = 1;
 	lex->input_length = len;
 	lex->input_encoding = encoding;
 	if (need_escapes)
+	{
 		lex->strval = makeStringInfo();
+		lex->flags |= JSONLEX_FREE_STRVAL;
+	}
+
 	return lex;
 }
 
+/*
+ * Free memory in a JsonLexContext
+ */
+void
+freeJsonLexContext(JsonLexContext *lex)
+{
+	if (lex->flags & JSONLEX_FREE_STRVAL)
+	{
+		pfree(lex->strval->data);
+		pfree(lex->strval);
+	}
+	if (lex->flags & JSONLEX_FREE_STRUCT)
+		pfree(lex);
+}
+
 /*
  * pg_parse_json
  *
diff --git a/src/include/common/jsonapi.h b/src/include/common/jsonapi.h
index 4310084b2b..a03d8310d4 100644
--- a/src/include/common/jsonapi.h
+++ b/src/include/common/jsonapi.h
@@ -71,6 +71,8 @@ typedef enum JsonParseErrorType
  * AFTER the end of the token, i.e. where there would be a nul byte
  * if we were using nul-terminated strings.
  */
+#define JSONLEX_FREE_STRUCT			(1 << 0)
+#define JSONLEX_FREE_STRVAL			(1 << 1)
 typedef struct JsonLexContext
 {
 	char	   *input;
@@ -84,6 +86,7 @@ typedef struct JsonLexContext
 	int			line_number;	/* line number, starting from 1 */
 	char	   *line_start;		/* where that line starts within input */
 	StringInfo	strval;
+	bits32		flags;
 } JsonLexContext;
 
 typedef JsonParseErrorType (*json_struct_action) (void *state);
@@ -151,15 +154,25 @@ extern JsonParseErrorType json_count_array_elements(JsonLexContext *lex,
 													int *elements);
 
 /*
- * constructor for JsonLexContext, with or without strval element.
- * If supplied, the strval element will contain a de-escaped version of
- * the lexeme. However, doing this imposes a performance penalty, so
- * it should be avoided if the de-escaped lexeme is not required.
+ * initializer for JsonLexContext.
+ *
+ * If a valid 'lex' pointer is given, it is initialized.  This can be used
+ * for stack-allocated structs, saving overhead.  If NULL is given, a new
+ * struct is allocated.
+ *
+ * If need_escapes is true, ->strval stores the unescaped lexemes.
+ * Unescaping is expensive, so only request it when necessary.
+ *
+ * If either need_escapes or lex was given as NULL, then the caller is
+ * responsible for freeing the returned struct, either by calling
+ * freeJsonLexContext() or via memory context cleanup.
  */
-extern JsonLexContext *makeJsonLexContextCstringLen(char *json,
+extern JsonLexContext *makeJsonLexContextCstringLen(JsonLexContext *lex,
+													char *json,
 													int len,
 													int encoding,
 													bool need_escapes);
+extern void freeJsonLexContext(JsonLexContext *lex);
 
 /* lex one token */
 extern JsonParseErrorType json_lex(JsonLexContext *lex);
diff --git a/src/include/utils/jsonfuncs.h b/src/include/utils/jsonfuncs.h
index c677ac8ff7..8d77aa9de0 100644
--- a/src/include/utils/jsonfuncs.h
+++ b/src/include/utils/jsonfuncs.h
@@ -37,7 +37,7 @@ typedef void (*JsonIterateStringValuesAction) (void *state, char *elem_value, in
 typedef text *(*JsonTransformStringValuesAction) (void *state, char *elem_value, int elem_len);
 
 /* build a JsonLexContext from a text datum */
-extern JsonLexContext *makeJsonLexContext(text *json, bool need_escapes);
+extern JsonLexContext *makeJsonLexContext(JsonLexContext *lex, text *json, bool need_escapes);
 
 /* try to parse json, and errsave(escontext) on failure */
 extern bool pg_parse_json_or_errsave(JsonLexContext *lex, JsonSemAction *sem,
-- 
2.39.2

