1 files changed, 76 insertions, 477 deletions

diff --git a/src/backend/utils/adt/json.c b/src/backend/utils/adt/json.c
index 5fdb7e32ce4..fee2ffb55c8 100644
--- a/src/backend/utils/adt/json.c
+++ b/src/backend/utils/adt/json.c
@@ -13,10 +13,7 @@
  */
 #include "postgres.h"
 
-#include "access/hash.h"
-#include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
-#include "common/hashfn.h"
 #include "funcapi.h"
 #include "libpq/pqformat.h"
 #include "miscadmin.h"
@@ -30,41 +27,20 @@
 #include "utils/lsyscache.h"
 #include "utils/typcache.h"
 
-/* Common context for key uniqueness check */
-typedef struct HTAB *JsonUniqueCheckState;	/* hash table for key names */
-
-/* Hash entry for JsonUniqueCheckState */
-typedef struct JsonUniqueHashEntry
-{
-	const char *key;
-	int			key_len;
-	int			object_id;
-} JsonUniqueHashEntry;
-
-/* Context for key uniqueness check in builder functions */
-typedef struct JsonUniqueBuilderState
-{
-	JsonUniqueCheckState check; /* unique check */
-	StringInfoData skipped_keys;	/* skipped keys with NULL values */
-	MemoryContext mcxt;			/* context for saving skipped keys */
-} JsonUniqueBuilderState;
-
-/* Element of object stack for key uniqueness check during json parsing */
-typedef struct JsonUniqueStackEntry
+typedef enum					/* type categories for datum_to_json */
 {
-	struct JsonUniqueStackEntry *parent;
-	int			object_id;
-} JsonUniqueStackEntry;
-
-/* State for key uniqueness check during json parsing */
-typedef struct JsonUniqueParsingState
-{
-	JsonLexContext *lex;
-	JsonUniqueCheckState check;
-	JsonUniqueStackEntry *stack;
-	int			id_counter;
-	bool		unique;
-} JsonUniqueParsingState;
+	JSONTYPE_NULL,				/* null, so we didn't bother to identify */
+	JSONTYPE_BOOL,				/* boolean (built-in types only) */
+	JSONTYPE_NUMERIC,			/* numeric (ditto) */
+	JSONTYPE_DATE,				/* we use special formatting for datetimes */
+	JSONTYPE_TIMESTAMP,
+	JSONTYPE_TIMESTAMPTZ,
+	JSONTYPE_JSON,				/* JSON itself (and JSONB) */
+	JSONTYPE_ARRAY,				/* array */
+	JSONTYPE_COMPOSITE,			/* composite */
+	JSONTYPE_CAST,				/* something with an explicit cast to JSON */
+	JSONTYPE_OTHER				/* all else */
+} JsonTypeCategory;
 
 typedef struct JsonAggState
 {
@@ -73,7 +49,6 @@ typedef struct JsonAggState
 	Oid			key_output_func;
 	JsonTypeCategory val_category;
 	Oid			val_output_func;
-	JsonUniqueBuilderState unique_check;
 } JsonAggState;
 
 static void composite_to_json(Datum composite, StringInfo result,
@@ -84,6 +59,9 @@ static void array_dim_to_json(StringInfo result, int dim, int ndims, int *dims,
 							  bool use_line_feeds);
 static void array_to_json_internal(Datum array, StringInfo result,
 								   bool use_line_feeds);
+static void json_categorize_type(Oid typoid,
+								 JsonTypeCategory *tcategory,
+								 Oid *outfuncoid);
 static void datum_to_json(Datum val, bool is_null, StringInfo result,
 						  JsonTypeCategory tcategory, Oid outfuncoid,
 						  bool key_scalar);
@@ -162,7 +140,7 @@ json_recv(PG_FUNCTION_ARGS)
  * output function OID.  If the returned category is JSONTYPE_CAST, we
  * return the OID of the type->JSON cast function instead.
  */
-void
+static void
 json_categorize_type(Oid typoid,
 					 JsonTypeCategory *tcategory,
 					 Oid *outfuncoid)
@@ -744,48 +722,6 @@ row_to_json_pretty(PG_FUNCTION_ARGS)
 	PG_RETURN_TEXT_P(cstring_to_text_with_len(result->data, result->len));
 }
 
-Datum
-to_json_worker(Datum val, JsonTypeCategory tcategory, Oid outfuncoid)
-{
-	StringInfo	result = makeStringInfo();
-
-	datum_to_json(val, false, result, tcategory, outfuncoid, false);
-
-	return PointerGetDatum(cstring_to_text_with_len(result->data, result->len));
-}
-
-bool
-to_json_is_immutable(Oid typoid)
-{
-	JsonTypeCategory tcategory;
-	Oid			outfuncoid;
-
-	json_categorize_type(typoid, &tcategory, &outfuncoid);
-
-	switch (tcategory)
-	{
-		case JSONTYPE_BOOL:
-		case JSONTYPE_JSON:
-			return true;
-
-		case JSONTYPE_DATE:
-		case JSONTYPE_TIMESTAMP:
-		case JSONTYPE_TIMESTAMPTZ:
-			return false;
-
-		case JSONTYPE_ARRAY:
-			return false;		/* TODO recurse into elements */
-
-		case JSONTYPE_COMPOSITE:
-			return false;		/* TODO recurse into fields */
-
-		case JSONTYPE_NUMERIC:
-		case JSONTYPE_CAST:
-		default:
-			return func_volatile(outfuncoid) == PROVOLATILE_IMMUTABLE;
-	}
-}
-
 /*
  * SQL function to_json(anyvalue)
  */
@@ -794,6 +730,7 @@ to_json(PG_FUNCTION_ARGS)
 {
 	Datum		val = PG_GETARG_DATUM(0);
 	Oid			val_type = get_fn_expr_argtype(fcinfo->flinfo, 0);
+	StringInfo	result;
 	JsonTypeCategory tcategory;
 	Oid			outfuncoid;
 
@@ -805,7 +742,11 @@ to_json(PG_FUNCTION_ARGS)
 	json_categorize_type(val_type,
 						 &tcategory, &outfuncoid);
 
-	PG_RETURN_DATUM(to_json_worker(val, tcategory, outfuncoid));
+	result = makeStringInfo();
+
+	datum_to_json(val, false, result, tcategory, outfuncoid, false);
+
+	PG_RETURN_TEXT_P(cstring_to_text_with_len(result->data, result->len));
 }
 
 /*
@@ -813,8 +754,8 @@ to_json(PG_FUNCTION_ARGS)
  *
  * aggregate input column as a json array value.
  */
-static Datum
-json_agg_transfn_worker(FunctionCallInfo fcinfo, bool absent_on_null)
+Datum
+json_agg_transfn(PG_FUNCTION_ARGS)
 {
 	MemoryContext aggcontext,
 				oldcontext;
@@ -854,13 +795,8 @@ json_agg_transfn_worker(FunctionCallInfo fcinfo, bool absent_on_null)
 	else
 	{
 		state = (JsonAggState *) PG_GETARG_POINTER(0);
-	}
-
-	if (absent_on_null && PG_ARGISNULL(1))
-		PG_RETURN_POINTER(state);
-
-	if (state->str->len > 1)
 		appendStringInfoString(state->str, ", ");
+	}
 
 	/* fast path for NULLs */
 	if (PG_ARGISNULL(1))
@@ -873,7 +809,7 @@ json_agg_transfn_worker(FunctionCallInfo fcinfo, bool absent_on_null)
 	val = PG_GETARG_DATUM(1);
 
 	/* add some whitespace if structured type and not first item */
-	if (!PG_ARGISNULL(0) && state->str->len > 1 &&
+	if (!PG_ARGISNULL(0) &&
 		(state->val_category == JSONTYPE_ARRAY ||
 		 state->val_category == JSONTYPE_COMPOSITE))
 	{
@@ -891,25 +827,6 @@ json_agg_transfn_worker(FunctionCallInfo fcinfo, bool absent_on_null)
 	PG_RETURN_POINTER(state);
 }
 
-
-/*
- * json_agg aggregate function
- */
-Datum
-json_agg_transfn(PG_FUNCTION_ARGS)
-{
-	return json_agg_transfn_worker(fcinfo, false);
-}
-
-/*
- * json_agg_strict aggregate function
- */
-Datum
-json_agg_strict_transfn(PG_FUNCTION_ARGS)
-{
-	return json_agg_transfn_worker(fcinfo, true);
-}
-
 /*
  * json_agg final function
  */
@@ -933,108 +850,18 @@ json_agg_finalfn(PG_FUNCTION_ARGS)
 	PG_RETURN_TEXT_P(catenate_stringinfo_string(state->str, "]"));
 }
 
-/* Functions implementing hash table for key uniqueness check */
-static uint32
-json_unique_hash(const void *key, Size keysize)
-{
-	const JsonUniqueHashEntry *entry = (JsonUniqueHashEntry *) key;
-	uint32		hash = hash_bytes_uint32(entry->object_id);
-
-	hash ^= hash_bytes((const unsigned char *) entry->key, entry->key_len);
-
-	return DatumGetUInt32(hash);
-}
-
-static int
-json_unique_hash_match(const void *key1, const void *key2, Size keysize)
-{
-	const JsonUniqueHashEntry *entry1 = (const JsonUniqueHashEntry *) key1;
-	const JsonUniqueHashEntry *entry2 = (const JsonUniqueHashEntry *) key2;
-
-	if (entry1->object_id != entry2->object_id)
-		return entry1->object_id > entry2->object_id ? 1 : -1;
-
-	if (entry1->key_len != entry2->key_len)
-		return entry1->key_len > entry2->key_len ? 1 : -1;
-
-	return strncmp(entry1->key, entry2->key, entry1->key_len);
-}
-
-/* Functions implementing object key uniqueness check */
-static void
-json_unique_check_init(JsonUniqueCheckState *cxt)
-{
-	HASHCTL		ctl;
-
-	memset(&ctl, 0, sizeof(ctl));
-	ctl.keysize = sizeof(JsonUniqueHashEntry);
-	ctl.entrysize = sizeof(JsonUniqueHashEntry);
-	ctl.hcxt = CurrentMemoryContext;
-	ctl.hash = json_unique_hash;
-	ctl.match = json_unique_hash_match;
-
-	*cxt = hash_create("json object hashtable",
-					   32,
-					   &ctl,
-					   HASH_ELEM | HASH_CONTEXT | HASH_FUNCTION | HASH_COMPARE);
-}
-
-static bool
-json_unique_check_key(JsonUniqueCheckState *cxt, const char *key, int object_id)
-{
-	JsonUniqueHashEntry entry;
-	bool		found;
-
-	entry.key = key;
-	entry.key_len = strlen(key);
-	entry.object_id = object_id;
-
-	(void) hash_search(*cxt, &entry, HASH_ENTER, &found);
-
-	return !found;
-}
-
-static void
-json_unique_builder_init(JsonUniqueBuilderState *cxt)
-{
-	json_unique_check_init(&cxt->check);
-	cxt->mcxt = CurrentMemoryContext;
-	cxt->skipped_keys.data = NULL;
-}
-
-/* On-demand initialization of skipped_keys StringInfo structure */
-static StringInfo
-json_unique_builder_get_skipped_keys(JsonUniqueBuilderState *cxt)
-{
-	StringInfo	out = &cxt->skipped_keys;
-
-	if (!out->data)
-	{
-		MemoryContext oldcxt = MemoryContextSwitchTo(cxt->mcxt);
-
-		initStringInfo(out);
-		MemoryContextSwitchTo(oldcxt);
-	}
-
-	return out;
-}
-
 /*
  * json_object_agg transition function.
  *
  * aggregate two input columns as a single json object value.
  */
-static Datum
-json_object_agg_transfn_worker(FunctionCallInfo fcinfo,
-							   bool absent_on_null, bool unique_keys)
+Datum
+json_object_agg_transfn(PG_FUNCTION_ARGS)
 {
 	MemoryContext aggcontext,
 				oldcontext;
 	JsonAggState *state;
-	StringInfo	out;
 	Datum		arg;
-	bool		skip;
-	int			key_offset;
 
 	if (!AggCheckCallContext(fcinfo, &aggcontext))
 	{
@@ -1055,10 +882,6 @@ json_object_agg_transfn_worker(FunctionCallInfo fcinfo,
 		oldcontext = MemoryContextSwitchTo(aggcontext);
 		state = (JsonAggState *) palloc(sizeof(JsonAggState));
 		state->str = makeStringInfo();
-		if (unique_keys)
-			json_unique_builder_init(&state->unique_check);
-		else
-			memset(&state->unique_check, 0, sizeof(state->unique_check));
 		MemoryContextSwitchTo(oldcontext);
 
 		arg_type = get_fn_expr_argtype(fcinfo->flinfo, 1);
@@ -1086,6 +909,7 @@ json_object_agg_transfn_worker(FunctionCallInfo fcinfo,
 	else
 	{
 		state = (JsonAggState *) PG_GETARG_POINTER(0);
+		appendStringInfoString(state->str, ", ");
 	}
 
 	/*
@@ -1101,49 +925,11 @@ json_object_agg_transfn_worker(FunctionCallInfo fcinfo,
 				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 				 errmsg("field name must not be null")));
 
-	/* Skip null values if absent_on_null */
-	skip = absent_on_null && PG_ARGISNULL(2);
-
-	if (skip)
-	{
-		/* If key uniqueness check is needed we must save skipped keys */
-		if (!unique_keys)
-			PG_RETURN_POINTER(state);
-
-		out = json_unique_builder_get_skipped_keys(&state->unique_check);
-	}
-	else
-	{
-		out = state->str;
-
-		/*
-		 * Append comma delimiter only if we have already outputted some
-		 * fields after the initial string "{ ".
-		 */
-		if (out->len > 2)
-			appendStringInfoString(out, ", ");
-	}
-
 	arg = PG_GETARG_DATUM(1);
 
-	key_offset = out->len;
-
-	datum_to_json(arg, false, out, state->key_category,
+	datum_to_json(arg, false, state->str, state->key_category,
 				  state->key_output_func, true);
 
-	if (unique_keys)
-	{
-		const char *key = &out->data[key_offset];
-
-		if (!json_unique_check_key(&state->unique_check.check, key, 0))
-			ereport(ERROR,
-					(errcode(ERRCODE_DUPLICATE_JSON_OBJECT_KEY_VALUE),
-					 errmsg("duplicate JSON key %s", key)));
-
-		if (skip)
-			PG_RETURN_POINTER(state);
-	}
-
 	appendStringInfoString(state->str, " : ");
 
 	if (PG_ARGISNULL(2))
@@ -1158,42 +944,6 @@ json_object_agg_transfn_worker(FunctionCallInfo fcinfo,
 }
 
 /*
- * json_object_agg aggregate function
- */
-Datum
-json_object_agg_transfn(PG_FUNCTION_ARGS)
-{
-	return json_object_agg_transfn_worker(fcinfo, false, false);
-}
-
-/*
- * json_object_agg_strict aggregate function
- */
-Datum
-json_object_agg_strict_transfn(PG_FUNCTION_ARGS)
-{
-	return json_object_agg_transfn_worker(fcinfo, true, false);
-}
-
-/*
- * json_object_agg_unique aggregate function
- */
-Datum
-json_object_agg_unique_transfn(PG_FUNCTION_ARGS)
-{
-	return json_object_agg_transfn_worker(fcinfo, false, true);
-}
-
-/*
- * json_object_agg_unique_strict aggregate function
- */
-Datum
-json_object_agg_unique_strict_transfn(PG_FUNCTION_ARGS)
-{
-	return json_object_agg_transfn_worker(fcinfo, true, true);
-}
-
-/*
  * json_object_agg final function.
  */
 Datum
@@ -1234,14 +984,25 @@ catenate_stringinfo_string(StringInfo buffer, const char *addon)
 	return result;
 }
 
+/*
+ * SQL function json_build_object(variadic "any")
+ */
 Datum
-json_build_object_worker(int nargs, Datum *args, bool *nulls, Oid *types,
-						 bool absent_on_null, bool unique_keys)
+json_build_object(PG_FUNCTION_ARGS)
 {
+	int			nargs;
 	int			i;
 	const char *sep = "";
 	StringInfo	result;
-	JsonUniqueBuilderState unique_check;
+	Datum	   *args;
+	bool	   *nulls;
+	Oid		   *types;
+
+	/* fetch argument values to build the object */
+	nargs = extract_variadic_args(fcinfo, 0, false, &args, &types, &nulls);
+
+	if (nargs < 0)
+		PG_RETURN_NULL();
 
 	if (nargs % 2 != 0)
 		ereport(ERROR,
@@ -1255,32 +1016,10 @@ json_build_object_worker(int nargs, Datum *args, bool *nulls, Oid *types,
 
 	appendStringInfoChar(result, '{');
 
-	if (unique_keys)
-		json_unique_builder_init(&unique_check);
-
 	for (i = 0; i < nargs; i += 2)
 	{
-		StringInfo	out;
-		bool		skip;
-		int			key_offset;
-
-		/* Skip null values if absent_on_null */
-		skip = absent_on_null && nulls[i + 1];
-
-		if (skip)
-		{
-			/* If key uniqueness check is needed we must save skipped keys */
-			if (!unique_keys)
-				continue;
-
-			out = json_unique_builder_get_skipped_keys(&unique_check);
-		}
-		else
-		{
-			appendStringInfoString(result, sep);
-			sep = ", ";
-			out = result;
-		}
+		appendStringInfoString(result, sep);
+		sep = ", ";
 
 		/* process key */
 		if (nulls[i])
@@ -1289,24 +1028,7 @@ json_build_object_worker(int nargs, Datum *args, bool *nulls, Oid *types,
 					 errmsg("argument %d cannot be null", i + 1),
 					 errhint("Object keys should be text.")));
 
-		/* save key offset before key appending */
-		key_offset = out->len;
-
-		add_json(args[i], false, out, types[i], true);
-
-		if (unique_keys)
-		{
-			/* check key uniqueness after key appending */
-			const char *key = &out->data[key_offset];
-
-			if (!json_unique_check_key(&unique_check.check, key, 0))
-				ereport(ERROR,
-						(errcode(ERRCODE_DUPLICATE_JSON_OBJECT_KEY_VALUE),
-						 errmsg("duplicate JSON key %s", key)));
-
-			if (skip)
-				continue;
-		}
+		add_json(args[i], false, result, types[i], true);
 
 		appendStringInfoString(result, " : ");
 
@@ -1316,27 +1038,7 @@ json_build_object_worker(int nargs, Datum *args, bool *nulls, Oid *types,
 
 	appendStringInfoChar(result, '}');
 
-	return PointerGetDatum(cstring_to_text_with_len(result->data, result->len));
-}
-
-/*
- * SQL function json_build_object(variadic "any")
- */
-Datum
-json_build_object(PG_FUNCTION_ARGS)
-{
-	Datum	   *args;
-	bool	   *nulls;
-	Oid		   *types;
-
-	/* build argument values to build the object */
-	int			nargs = extract_variadic_args(fcinfo, 0, true,
-											  &args, &types, &nulls);
-
-	if (nargs < 0)
-		PG_RETURN_NULL();
-
-	PG_RETURN_DATUM(json_build_object_worker(nargs, args, nulls, types, false, false));
+	PG_RETURN_TEXT_P(cstring_to_text_with_len(result->data, result->len));
 }
 
 /*
@@ -1348,13 +1050,25 @@ json_build_object_noargs(PG_FUNCTION_ARGS)
 	PG_RETURN_TEXT_P(cstring_to_text_with_len("{}", 2));
 }
 
+/*
+ * SQL function json_build_array(variadic "any")
+ */
 Datum
-json_build_array_worker(int nargs, Datum *args, bool *nulls, Oid *types,
-						bool absent_on_null)
+json_build_array(PG_FUNCTION_ARGS)
 {
+	int			nargs;
 	int			i;
 	const char *sep = "";
 	StringInfo	result;
+	Datum	   *args;
+	bool	   *nulls;
+	Oid		   *types;
+
+	/* fetch argument values to build the array */
+	nargs = extract_variadic_args(fcinfo, 0, false, &args, &types, &nulls);
+
+	if (nargs < 0)
+		PG_RETURN_NULL();
 
 	result = makeStringInfo();
 
@@ -1362,9 +1076,6 @@ json_build_array_worker(int nargs, Datum *args, bool *nulls, Oid *types,
 
 	for (i = 0; i < nargs; i++)
 	{
-		if (absent_on_null && nulls[i])
-			continue;
-
 		appendStringInfoString(result, sep);
 		sep = ", ";
 		add_json(args[i], nulls[i], result, types[i], false);
@@ -1372,27 +1083,7 @@ json_build_array_worker(int nargs, Datum *args, bool *nulls, Oid *types,
 
 	appendStringInfoChar(result, ']');
 
-	return PointerGetDatum(cstring_to_text_with_len(result->data, result->len));
-}
-
-/*
- * SQL function json_build_array(variadic "any")
- */
-Datum
-json_build_array(PG_FUNCTION_ARGS)
-{
-	Datum	   *args;
-	bool	   *nulls;
-	Oid		   *types;
-
-	/* build argument values to build the object */
-	int			nargs = extract_variadic_args(fcinfo, 0, true,
-											  &args, &types, &nulls);
-
-	if (nargs < 0)
-		PG_RETURN_NULL();
-
-	PG_RETURN_DATUM(json_build_array_worker(nargs, args, nulls, types, false));
+	PG_RETURN_TEXT_P(cstring_to_text_with_len(result->data, result->len));
 }
 
 /*
@@ -1618,106 +1309,6 @@ escape_json(StringInfo buf, const char *str)
 	appendStringInfoCharMacro(buf, '"');
 }
 
-/* Semantic actions for key uniqueness check */
-static void
-json_unique_object_start(void *_state)
-{
-	JsonUniqueParsingState *state = _state;
-	JsonUniqueStackEntry *entry;
-
-	if (!state->unique)
-		return;
-
-	/* push object entry to stack */
-	entry = palloc(sizeof(*entry));
-	entry->object_id = state->id_counter++;
-	entry->parent = state->stack;
-	state->stack = entry;
-}
-
-static void
-json_unique_object_end(void *_state)
-{
-	JsonUniqueParsingState *state = _state;
-	JsonUniqueStackEntry *entry;
-
-	if (!state->unique)
-		return;
-
-	entry = state->stack;
-	state->stack = entry->parent;	/* pop object from stack */
-	pfree(entry);
-}
-
-static void
-json_unique_object_field_start(void *_state, char *field, bool isnull)
-{
-	JsonUniqueParsingState *state = _state;
-	JsonUniqueStackEntry *entry;
-
-	if (!state->unique)
-		return;
-
-	/* find key collision in the current object */
-	if (json_unique_check_key(&state->check, field, state->stack->object_id))
-		return;
-
-	state->unique = false;
-
-	/* pop all objects entries */
-	while ((entry = state->stack))
-	{
-		state->stack = entry->parent;
-		pfree(entry);
-	}
-}
-
-/* Validate JSON text and additionally check key uniqueness */
-bool
-json_validate(text *json, bool check_unique_keys, bool throw_error)
-{
-	JsonLexContext *lex = makeJsonLexContext(json, check_unique_keys);
-	JsonSemAction uniqueSemAction = {0};
-	JsonUniqueParsingState state;
-	JsonParseErrorType result;
-
-	if (check_unique_keys)
-	{
-		state.lex = lex;
-		state.stack = NULL;
-		state.id_counter = 0;
-		state.unique = true;
-		json_unique_check_init(&state.check);
-
-		uniqueSemAction.semstate = &state;
-		uniqueSemAction.object_start = json_unique_object_start;
-		uniqueSemAction.object_field_start = json_unique_object_field_start;
-		uniqueSemAction.object_end = json_unique_object_end;
-	}
-
-	result = pg_parse_json(lex, check_unique_keys ? &uniqueSemAction : &nullSemAction);
-
-	if (result != JSON_SUCCESS)
-	{
-		if (throw_error)
-			json_ereport_error(result, lex);
-
-		return false;			/* invalid json */
-	}
-
-	if (check_unique_keys && !state.unique)
-	{
-		if (throw_error)
-			ereport(ERROR,
-					(errcode(ERRCODE_DUPLICATE_JSON_OBJECT_KEY_VALUE),
-					 errmsg("duplicate JSON object key value")));
-
-		return false;			/* not unique keys */
-	}
-
-	return true;				/* ok */
-}
-
 /*
  * SQL function json_typeof(json) -> text
  *
@@ -1733,13 +1324,21 @@ json_validate(text *json, bool check_unique_keys, bool throw_error)
 Datum
 json_typeof(PG_FUNCTION_ARGS)
 {
-	text	   *json = PG_GETARG_TEXT_PP(0);
-	char	   *type;
+	text	   *json;
+
+	JsonLexContext *lex;
 	JsonTokenType tok;
+	char	   *type;
+	JsonParseErrorType result;
 
-	/* Lex exactly one token from the input and check its type. */
-	tok = json_get_first_token(json, true);
+	json = PG_GETARG_TEXT_PP(0);
+	lex = makeJsonLexContext(json, false);
 
+	/* Lex exactly one token from the input and check its type. */
+	result = json_lex(lex);
+	if (result != JSON_SUCCESS)
+		json_ereport_error(result, lex);
+	tok = lex->token_type;
 	switch (tok)
 	{
 		case JSON_TOKEN_OBJECT_START: