From fd2b9dbab24622e741108ca55a3f1fcae3c1719b Mon Sep 17 00:00:00 2001
From: amitlan <amitlangote09@gmail.com>
Date: Fri, 15 Jul 2022 12:57:50 +0900
Subject: [PATCH v2 3/3] Use one ExprState to implement JsonItemCoercions

Current implementation assigns one ExprState for every JsonCoercion
member of JsonItemCoercions, which seems wasteful.

This commit instead assigns the ExprState to JsonItemCoercions itself
and arranges to compute the member JsonCoercion's expressions using
ExprEvalSteps instead.
---
 src/backend/executor/execExpr.c       | 113 ++++++++++++++++++----
 src/backend/executor/execExprInterp.c | 133 ++++++++++++++++++++------
 src/backend/jit/llvm/llvmjit_expr.c   |   3 +
 src/backend/optimizer/util/clauses.c  |   6 +-
 src/include/executor/execExpr.h       |  75 +++++++++------
 src/include/utils/jsonb.h             |   2 +
 6 files changed, 256 insertions(+), 76 deletions(-)

diff --git a/src/backend/executor/execExpr.c b/src/backend/executor/execExpr.c
index 4dbb24aaee..be98544e05 100644
--- a/src/backend/executor/execExpr.c
+++ b/src/backend/executor/execExpr.c
@@ -2786,7 +2786,8 @@ ExecInitExprRec(Expr *node, ExprState *state,
 				 * query.
 				 */
 				jsestate->needSubtrans =
-					ExecEvalJsonNeedsSubTransaction(jexpr, &jsestate->coercions);
+					ExecEvalJsonNeedsSubTransaction(jexpr,
+													true /* coerce */);
 				jsestate->coercionNeedSubtrans =
 					(!jsestate->needSubtrans && !throwErrors);
 
@@ -2810,26 +2811,102 @@ ExecInitExprRec(Expr *node, ExprState *state,
 
 				/* Set up coercion related data structures. */
 				if (jexpr->coercions)
+					jsestate->coercions =
+						ExecInitExprWithCaseValue((Expr *) jexpr->coercions,
+												  state->parent,
+												  resv, resnull);
+
+				break;
+			}
+
+		case T_JsonItemCoercions:
+			{
+				JsonItemCoercions *coercions = castNode(JsonItemCoercions,
+														node);
+				JsonItemCoercionsState *jcstate =
+					palloc0(sizeof(JsonItemCoercionsState));
+				JsonCoercion **coercion;
+				JsonItemCoercionState *cstate;
+				ExprEvalStep *as;
+				int			json_item_coercion_step_id;
+				List	   *adjust_jumps = NIL;
+				ListCell   *lc;
+
+				/*
+				 * First push a step to read the value provided by the parent
+				 * JsonExpr via a CaseTestExpr.
+				 */
+				scratch.opcode = EEOP_CASE_TESTVAL;
+				scratch.d.casetest.value = state->innermost_caseval;
+				scratch.d.casetest.isnull = state->innermost_casenull;
+				ExprEvalPushStep(state, &scratch);
+
+				/* Push the control step. */
+				scratch.opcode = EEOP_JSON_ITEM_COERCION;
+				scratch.d.json_item_coercion.jcstate = jcstate;
+				json_item_coercion_step_id = state->steps_len;
+				ExprEvalPushStep(state, &scratch);
+				/* Will set jump_skip_coercion target address below. */
+
+				/*
+				 * Now push the steps of individual coercion's expression, if
+				 * needed.
+				 */
+				for (cstate = &jcstate->null,
+					 coercion = &coercions->null;
+					 coercion <= &coercions->composite;
+					 coercion++, cstate++)
 				{
-					JsonCoercion **coercion;
-					struct JsonCoercionState *cstate;
-					Datum	   *caseval;
-					bool	   *casenull;
-
-					caseval = resv;
-					casenull = resnull;
-
-					for (cstate = &jsestate->coercions.null,
-						 coercion = &jexpr->coercions->null;
-						 coercion <= &jexpr->coercions->composite;
-						 coercion++, cstate++)
+					cstate->coercion = *coercion;
+					if (cstate->coercion && cstate->coercion->expr)
 					{
-						cstate->coercion = *coercion;
-						cstate->estate = *coercion ?
-							ExecInitExprWithCaseValue((Expr *) (*coercion)->expr,
-													  state->parent,
-													  caseval, casenull) : NULL;
+						Datum	   *save_innermost_caseval;
+						bool	   *save_innermost_isnull;
+
+						cstate->jump_eval_expr = state->steps_len;
+
+						/* Push step(s) to compute (*coercion)->expr. */
+						save_innermost_caseval = state->innermost_caseval;
+						save_innermost_isnull = state->innermost_casenull;
+
+						state->innermost_caseval = resv;
+						state->innermost_casenull = resnull;
+
+						ExecInitExprRec((Expr *) cstate->coercion->expr,
+										state, resv, resnull);
+
+						state->innermost_caseval = save_innermost_caseval;
+						state->innermost_casenull = save_innermost_isnull;
 					}
+					else
+						cstate->jump_eval_expr = -1;
+
+					/* Emit JUMP step to jump to end of coercions code */
+					scratch.opcode = EEOP_JUMP;
+
+					/*
+					 * Remember JUMP step address to set the actual jump
+					 * target address below.
+					 */
+					adjust_jumps = lappend_int(adjust_jumps,
+											   state->steps_len);
+					ExprEvalPushStep(state, &scratch);
+				}
+
+				/*
+				 * Adjust the jump target address of the control step and all
+				 * the jumps we added in the above loop to make them point to
+				 * the step after the last step that would have been added
+				 * above.
+				 */
+				as = &state->steps[json_item_coercion_step_id];
+				as->d.json_item_coercion.jump_skip_item_coercion = state->steps_len;
+				foreach(lc, adjust_jumps)
+				{
+					int		jump_step_id = lfirst_int(lc);
+
+					as = &state->steps[jump_step_id];
+					as->d.jump.jumpdone = state->steps_len;
 				}
 
 				break;
diff --git a/src/backend/executor/execExprInterp.c b/src/backend/executor/execExprInterp.c
index 5309017027..144536465a 100644
--- a/src/backend/executor/execExprInterp.c
+++ b/src/backend/executor/execExprInterp.c
@@ -506,6 +506,7 @@ ExecInterpExpr(ExprState *state, ExprContext *econtext, bool *isnull)
 		&&CASE_EEOP_JSONEXPR_COERCION,
 		&&CASE_EEOP_JSONEXPR_COERCION_ERROR,
 		&&CASE_EEOP_JSONEXPR_RESULT,
+		&&CASE_EEOP_JSON_ITEM_COERCION,
 		&&CASE_EEOP_AGG_STRICT_DESERIALIZE,
 		&&CASE_EEOP_AGG_DESERIALIZE,
 		&&CASE_EEOP_AGG_STRICT_INPUT_CHECK_ARGS,
@@ -1844,7 +1845,7 @@ ExecInterpExpr(ExprState *state, ExprContext *econtext, bool *isnull)
 			 */
 			if (pre_eval->formatted_expr.isnull || pre_eval->pathspec.isnull)
 			{
-				post_eval->coercion = NULL;
+				post_eval->coercions = NULL;
 				/* A signal to coercion step to not use a sub-trancaction. */
 				post_eval->coercion_error = true;
 				*op->resvalue = 0;
@@ -1948,7 +1949,7 @@ ExecInterpExpr(ExprState *state, ExprContext *econtext, bool *isnull)
 
 				/* Coercion the ON ERROR expression. */
 				post_eval->coercion_done = false;
-				post_eval->coercion = NULL;
+				post_eval->coercions = NULL;
 				EEO_JUMP(op->d.jsonexpr_coercion_error.jump_coercion);
 			}
 
@@ -2008,6 +2009,23 @@ ExecInterpExpr(ExprState *state, ExprContext *econtext, bool *isnull)
 			EEO_NEXT();
 		}
 
+		EEO_CASE(EEOP_JSON_ITEM_COERCION)
+		{
+			JsonItemCoercionsState *jcstate = op->d.json_item_coercion.jcstate;
+			JsonItemCoercionState *cstate = NULL;
+
+			*op->resvalue =
+				ExecPrepareJsonItemCoercion(*op->resvalue, jcstate, &cstate);
+			if (cstate->coercion && cstate->coercion->expr)
+			{
+				Assert(cstate->jump_eval_expr >= 0);
+				EEO_JUMP(cstate->jump_eval_expr);
+			}
+
+			/* Skip over all of the steps added for this JsonItemCoercions. */
+			EEO_JUMP(op->d.json_item_coercion.jump_skip_item_coercion);
+		}
+
 		EEO_CASE(EEOP_LAST)
 		{
 			/* unreachable */
@@ -4831,8 +4849,7 @@ ExecEvalJsonExprCoercion(ExprEvalStep *op, ExprContext *econtext,
 	JsonExprState *jsestate = op->d.jsonexpr_coercion.jsestate;
 	JsonExprPostEvalState *post_eval = &jsestate->post_eval;
 	JsonExpr   *jexpr = jsestate->jsexpr;
-	ExprState  *estate = post_eval->coercion ?
-						post_eval->coercion->estate : NULL;
+	ExprState  *estate = post_eval->coercions;
 
 	/*
 	 * Initially assume we will indeed perform the coercion either
@@ -4944,19 +4961,84 @@ EvalJsonPathVar(void *cxt, char *varName, int varNameLen,
 	return id;
 }
 
+/*
+ * Return a coercion among those given in 'coercions' for given
+ * JSON item.
+ */
+JsonCoercion *
+ExecGetJsonItemCoercion(JsonbValue *item, JsonItemCoercions *coercions)
+{
+	if (item->type == jbvBinary &&
+		JsonContainerIsScalar(item->val.binary.data))
+	{
+		JsonbValue	buf;
+		bool		res PG_USED_FOR_ASSERTS_ONLY;
+
+		res = JsonbExtractScalar(item->val.binary.data, &buf);
+		item = &buf;
+		Assert(res);
+	}
+
+	/* get coercion state reference and datum of the corresponding SQL type */
+	switch (item->type)
+	{
+		case jbvNull:
+			return coercions->null;
+
+		case jbvString:
+			return coercions->string;
+
+		case jbvNumeric:
+			return coercions->numeric;
+
+		case jbvBool:
+			return coercions->boolean;
+
+		case jbvDatetime:
+			switch (item->val.datetime.typid)
+			{
+				case DATEOID:
+					return coercions->date;
+				case TIMEOID:
+					return coercions->time;
+				case TIMETZOID:
+					return coercions->timetz;
+				case TIMESTAMPOID:
+					return coercions->timestamp;
+				case TIMESTAMPTZOID:
+					return coercions->timestamptz;
+				default:
+					elog(ERROR, "unexpected jsonb datetime type oid %u",
+						 item->val.datetime.typid);
+			}
+			break;
+
+		case jbvArray:
+		case jbvObject:
+		case jbvBinary:
+			return coercions->composite;
+
+		default:
+			elog(ERROR, "unexpected jsonb value type %d", item->type);
+	}
+
+	Assert(false);
+	return NULL;
+}
+
 /*
  * Prepare SQL/JSON item coercion to the output type. Returned a datum of the
  * corresponding SQL type and a pointer to the coercion state.
  */
 Datum
-ExecPrepareJsonItemCoercion(JsonbValue *item,
-							JsonReturning *returning,
-							struct JsonCoercionsState *coercions,
-							struct JsonCoercionState **pcoercion)
+ExecPrepareJsonItemCoercion(Datum itemval,
+							JsonItemCoercionsState *coercions,
+							JsonItemCoercionState **pcoercion)
 {
-	struct JsonCoercionState *coercion;
+	JsonbValue *item = DatumGetJsonbValueP(itemval);
+	JsonItemCoercionState *coercion;
 	Datum		res;
-	JsonbValue	buf;
+	JsonbValue 	buf;
 
 	if (item->type == jbvBinary &&
 		JsonContainerIsScalar(item->val.binary.data))
@@ -5136,7 +5218,7 @@ ExecEvalJsonExpr(ExprEvalStep *op, ExprContext *econtext,
 
 		case JSON_VALUE_OP:
 			{
-				struct JsonCoercionState *jcstate;
+				JsonCoercion *coercion;
 				JsonbValue *jbv = JsonPathValue(item, path, empty, error,
 												pre_eval->args);
 
@@ -5162,14 +5244,12 @@ ExecEvalJsonExpr(ExprEvalStep *op, ExprContext *econtext,
 					break;
 				}
 
-				/* Use coercion from SQL/JSON item type to the output type */
-				res = ExecPrepareJsonItemCoercion(jbv,
-												  jsestate->jsexpr->returning,
-												  &jsestate->coercions,
-												  &jcstate);
-				if (jcstate->coercion &&
-					(jcstate->coercion->via_io ||
-					 jcstate->coercion->via_populate))
+				/*
+				 * Error out no cast exists to coerce SQL/JSON item to the
+				 * the output type
+				 */
+				coercion = ExecGetJsonItemCoercion(jbv, jsestate->jsexpr->coercions);
+				if (coercion && (coercion->via_io || coercion->via_populate))
 				{
 					if (error)
 					{
@@ -5186,15 +5266,11 @@ ExecEvalJsonExpr(ExprEvalStep *op, ExprContext *econtext,
 							(errcode(ERRCODE_SQL_JSON_ITEM_CANNOT_BE_CAST_TO_TARGET_TYPE),
 							 errmsg("SQL/JSON item cannot be cast to target type")));
 				}
-				else if (jcstate->estate == NULL)
-				{
-					/* No coercion needed */
-					post_eval->coercion_done = true;
-				}
 				else
 				{
-					/* Coerce using specific expression */
-					post_eval->coercion = jcstate;
+					/* Coerce using specific expression. */
+					res = JsonbValuePGetDatum(jbv);
+					post_eval->coercions = jsestate->coercions;
 				}
 				break;
 			}
@@ -5252,8 +5328,7 @@ ExecEvalJsonExpr(ExprEvalStep *op, ExprContext *econtext,
 }
 
 bool
-ExecEvalJsonNeedsSubTransaction(JsonExpr *jsexpr,
-								struct JsonCoercionsState *coercions)
+ExecEvalJsonNeedsSubTransaction(JsonExpr *jsexpr, bool coerce)
 {
 	if (jsexpr->on_error->btype == JSON_BEHAVIOR_ERROR)
 		return false;
@@ -5261,7 +5336,7 @@ ExecEvalJsonNeedsSubTransaction(JsonExpr *jsexpr,
 	if (jsexpr->op == JSON_EXISTS_OP && !jsexpr->result_coercion)
 		return false;
 
-	if (!coercions)
+	if (!coerce)
 		return true;
 
 	return false;
diff --git a/src/backend/jit/llvm/llvmjit_expr.c b/src/backend/jit/llvm/llvmjit_expr.c
index c133600005..5f06a87373 100644
--- a/src/backend/jit/llvm/llvmjit_expr.c
+++ b/src/backend/jit/llvm/llvmjit_expr.c
@@ -2397,6 +2397,9 @@ llvm_compile_expr(ExprState *state)
 									opblocks[opno + 1]);
 				}
 				break;
+			case EEOP_JSON_ITEM_COERCION:
+				/* XXX - handle! */
+				break;
 
 			case EEOP_LAST:
 				Assert(false);
diff --git a/src/backend/optimizer/util/clauses.c b/src/backend/optimizer/util/clauses.c
index 533df86ff7..9c02218155 100644
--- a/src/backend/optimizer/util/clauses.c
+++ b/src/backend/optimizer/util/clauses.c
@@ -901,7 +901,11 @@ max_parallel_hazard_walker(Node *node, max_parallel_hazard_context *context)
 	{
 		JsonExpr   *jsexpr = (JsonExpr *) node;
 
-		if (ExecEvalJsonNeedsSubTransaction(jsexpr, NULL))
+		/*
+		 * XXX - don't really know why it makes sense to ignore the coercion
+		 * part here.
+		 */
+		if (ExecEvalJsonNeedsSubTransaction(jsexpr, false /* coerce */))
 		{
 			context->max_hazard = PROPARALLEL_UNSAFE;
 			return true;
diff --git a/src/include/executor/execExpr.h b/src/include/executor/execExpr.h
index 358b831f7a..2639d7c73d 100644
--- a/src/include/executor/execExpr.h
+++ b/src/include/executor/execExpr.h
@@ -251,6 +251,7 @@ typedef enum ExprEvalOp
 	EEOP_JSONEXPR_COERCION,
 	EEOP_JSONEXPR_COERCION_ERROR,
 	EEOP_JSONEXPR_RESULT,
+	EEOP_JSON_ITEM_COERCION,
 
 	/* aggregation related nodes */
 	EEOP_AGG_STRICT_DESERIALIZE,
@@ -753,6 +754,13 @@ typedef struct ExprEvalStep
 			JsonExpr *jexpr;
 			int		jump_skip_result;
 		}		jsonexpr_result;
+
+		/* for EEOP_JSON_ITEM_COERCION */
+		struct
+		{
+			struct JsonItemCoercionsState *jcstate;
+			int		jump_skip_item_coercion;
+		}		json_item_coercion;
 	}			d;
 } ExprEvalStep;
 
@@ -833,10 +841,8 @@ typedef struct JsonExprPostEvalState
 	/* Cache for json_populate_type() called for coercion in some cases */
 	void	   *cache;
 
-	/*
-	 * Coercion to be applied to JSON item returned by JsonPathValue(),
-	 * set by ExecPrepareJsonItemCoercion(). */
-	struct JsonCoercionState *coercion;
+	/* Coercion state; same as parent JsonExprState.coercions when not NULL */
+	ExprState   *coercions;
 }	JsonExprPostEvalState;
 
 /*
@@ -875,26 +881,40 @@ typedef struct JsonExprState
 		Oid			typioparam;
 	}			input;			/* I/O info for output type */
 
-	struct JsonCoercionsState
-	{
-		struct JsonCoercionState
-		{
-			JsonCoercion *coercion; /* coercion expression */
-			ExprState  *estate; /* coercion expression state */
-		}			null,
-					string,
-					numeric,
-					boolean,
-					date,
-					time,
-					timetz,
-					timestamp,
-					timestamptz,
-					composite;
-	}			coercions;		/* states for coercion from SQL/JSON item
-								 * types directly to the output type */
+	/* To apply coercion to the final result if needed. */
+	ExprState			   *coercions;
 } JsonExprState;
 
+/*
+ * State for a given member of JsonItemCoercions.
+ */
+typedef struct JsonItemCoercionState
+{
+	/* Expression used to evaluate the coercion */
+	JsonCoercion   *coercion;
+
+	/* ExprEvalStep to compute this coercion's expression */
+	int				jump_eval_expr;
+} JsonItemCoercionState;
+
+/*
+ * State for evaluating the coercion for a given JSON item using one of
+ * the following coercions.
+ */
+typedef struct JsonItemCoercionsState
+{
+	JsonItemCoercionState	null;
+	JsonItemCoercionState	string;
+	JsonItemCoercionState	numeric;
+	JsonItemCoercionState	boolean;
+	JsonItemCoercionState	date;
+	JsonItemCoercionState	time;
+	JsonItemCoercionState	timetz;
+	JsonItemCoercionState	timestamp;
+	JsonItemCoercionState	timestamptz;
+	JsonItemCoercionState	composite;
+} JsonItemCoercionsState;
+
 /* functions in execExpr.c */
 extern void ExprEvalPushStep(ExprState *es, const ExprEvalStep *s);
 
@@ -956,12 +976,11 @@ extern void ExecEvalJsonConstructor(ExprState *state, ExprEvalStep *op,
 									ExprContext *econtext);
 extern void ExecEvalJson(ExprState *state, ExprEvalStep *op,
 						 ExprContext *econtext);
-extern Datum ExecPrepareJsonItemCoercion(struct JsonbValue *item,
-										 JsonReturning *returning,
-										 struct JsonCoercionsState *coercions,
-										 struct JsonCoercionState **pjcstate);
-extern bool ExecEvalJsonNeedsSubTransaction(JsonExpr *jsexpr,
-											struct JsonCoercionsState *);
+JsonCoercion *ExecGetJsonItemCoercion(struct JsonbValue *item, JsonItemCoercions *coercions);
+extern Datum ExecPrepareJsonItemCoercion(Datum itemval,
+										 JsonItemCoercionsState *coercions,
+										 JsonItemCoercionState **pjcstate);
+extern bool ExecEvalJsonNeedsSubTransaction(JsonExpr *jsexpr, bool coerce);
 extern Datum ExecEvalExprPassingCaseValue(ExprState *estate,
 										  ExprContext *econtext, bool *isnull,
 										  Datum caseval_datum,
diff --git a/src/include/utils/jsonb.h b/src/include/utils/jsonb.h
index bae466b523..6bdd9f5121 100644
--- a/src/include/utils/jsonb.h
+++ b/src/include/utils/jsonb.h
@@ -69,8 +69,10 @@ typedef enum
 
 /* Convenience macros */
 #define DatumGetJsonbP(d)	((Jsonb *) PG_DETOAST_DATUM(d))
+#define DatumGetJsonbValueP(d)	((JsonbValue *) DatumGetPointer(d))
 #define DatumGetJsonbPCopy(d)	((Jsonb *) PG_DETOAST_DATUM_COPY(d))
 #define JsonbPGetDatum(p)	PointerGetDatum(p)
+#define JsonbValuePGetDatum(p)	PointerGetDatum(p)
 #define PG_GETARG_JSONB_P(x)	DatumGetJsonbP(PG_GETARG_DATUM(x))
 #define PG_GETARG_JSONB_P_COPY(x)	DatumGetJsonbPCopy(PG_GETARG_DATUM(x))
 #define PG_RETURN_JSONB_P(x)	PG_RETURN_POINTER(x)
-- 
2.35.3

