From 29845033cce1de0e80b3c19d4e56b514d094e956 Mon Sep 17 00:00:00 2001
From: Nikita Glukhov <n.gluhov@postgrespro.ru>
Date: Sat, 1 Apr 2023 16:21:20 +0300
Subject: [PATCH v7 2/5] Pass field accessors to generic subscripting

---
 src/backend/executor/execExpr.c    | 24 +++++++---
 src/backend/nodes/nodeFuncs.c      | 73 ++++++++++++++++++++++++++----
 src/backend/parser/parse_collate.c | 22 +++++++--
 src/backend/parser/parse_expr.c    | 58 ++++++++++++++++--------
 src/backend/parser/parse_node.c    | 39 ++++++++++++++--
 src/backend/parser/parse_target.c  |  3 +-
 src/backend/utils/adt/arraysubs.c  | 13 ++++--
 src/backend/utils/adt/jsonbsubs.c  | 11 ++++-
 src/backend/utils/adt/ruleutils.c  | 29 ++++++++----
 src/include/parser/parse_node.h    |  3 +-
 10 files changed, 218 insertions(+), 57 deletions(-)

diff --git a/src/backend/executor/execExpr.c b/src/backend/executor/execExpr.c
index 03566c4d181..be4213455e5 100644
--- a/src/backend/executor/execExpr.c
+++ b/src/backend/executor/execExpr.c
@@ -3328,9 +3328,15 @@ ExecInitSubscriptingRef(ExprEvalStep *scratch, SubscriptingRef *sbsref,
 		{
 			sbsrefstate->upperprovided[i] = true;
 			/* Each subscript is evaluated into appropriate array entry */
-			ExecInitExprRec(e, state,
-							&sbsrefstate->upperindex[i],
-							&sbsrefstate->upperindexnull[i]);
+			if (IsA(e, String))
+			{
+				sbsrefstate->upperindex[i] = CStringGetTextDatum(strVal(e));
+				sbsrefstate->upperindexnull[i] = false;
+			}
+			else
+				ExecInitExprRec(e, state,
+								&sbsrefstate->upperindex[i],
+								&sbsrefstate->upperindexnull[i]);
 		}
 		i++;
 	}
@@ -3351,9 +3357,15 @@ ExecInitSubscriptingRef(ExprEvalStep *scratch, SubscriptingRef *sbsref,
 		{
 			sbsrefstate->lowerprovided[i] = true;
 			/* Each subscript is evaluated into appropriate array entry */
-			ExecInitExprRec(e, state,
-							&sbsrefstate->lowerindex[i],
-							&sbsrefstate->lowerindexnull[i]);
+			if (IsA(e, String))
+			{
+				sbsrefstate->lowerindex[i] = CStringGetTextDatum(strVal(e));
+				sbsrefstate->lowerindexnull[i] = false;
+			}
+			else
+				ExecInitExprRec(e, state,
+								&sbsrefstate->lowerindex[i],
+								&sbsrefstate->lowerindexnull[i]);
 		}
 		i++;
 	}
diff --git a/src/backend/nodes/nodeFuncs.c b/src/backend/nodes/nodeFuncs.c
index 7bc823507f1..a9c29ab8f29 100644
--- a/src/backend/nodes/nodeFuncs.c
+++ b/src/backend/nodes/nodeFuncs.c
@@ -2182,12 +2182,28 @@ expression_tree_walker_impl(Node *node,
 		case T_SubscriptingRef:
 			{
 				SubscriptingRef *sbsref = (SubscriptingRef *) node;
+				ListCell   *lc;
+
+				/*
+				 * Recurse directly for upper/lower container index lists,
+				 * skipping String subscripts used for dot notation.
+				 */
+				foreach(lc, sbsref->refupperindexpr)
+				{
+					Node	   *expr = lfirst(lc);
+
+					if (expr && !IsA(expr, String) && WALK(expr))
+						return true;
+				}
+
+				foreach(lc, sbsref->reflowerindexpr)
+				{
+					Node	   *expr = lfirst(lc);
+
+					if (expr && !IsA(expr, String) && WALK(expr))
+						return true;
+				}
 
-				/* recurse directly for upper/lower container index lists */
-				if (LIST_WALK(sbsref->refupperindexpr))
-					return true;
-				if (LIST_WALK(sbsref->reflowerindexpr))
-					return true;
 				/* walker must see the refexpr and refassgnexpr, however */
 				if (WALK(sbsref->refexpr))
 					return true;
@@ -3082,12 +3098,51 @@ expression_tree_mutator_impl(Node *node,
 			{
 				SubscriptingRef *sbsref = (SubscriptingRef *) node;
 				SubscriptingRef *newnode;
+				ListCell   *lc;
+				List	   *exprs = NIL;
 
 				FLATCOPY(newnode, sbsref, SubscriptingRef);
-				MUTATE(newnode->refupperindexpr, sbsref->refupperindexpr,
-					   List *);
-				MUTATE(newnode->reflowerindexpr, sbsref->reflowerindexpr,
-					   List *);
+
+				foreach(lc, sbsref->refupperindexpr)
+				{
+					Node	   *expr = lfirst(lc);
+
+					if (expr && IsA(expr, String))
+					{
+						String	   *str;
+
+						FLATCOPY(str, expr, String);
+						expr = (Node *) str;
+					}
+					else
+						expr = mutator(expr, context);
+
+					exprs = lappend(exprs, expr);
+				}
+
+				newnode->refupperindexpr = exprs;
+
+				exprs = NIL;
+
+				foreach(lc, sbsref->reflowerindexpr)
+				{
+					Node	   *expr = lfirst(lc);
+
+					if (expr && IsA(expr, String))
+					{
+						String	   *str;
+
+						FLATCOPY(str, expr, String);
+						expr = (Node *) str;
+					}
+					else
+						expr = mutator(expr, context);
+
+					exprs = lappend(exprs, expr);
+				}
+
+				newnode->reflowerindexpr = exprs;
+
 				MUTATE(newnode->refexpr, sbsref->refexpr,
 					   Expr *);
 				MUTATE(newnode->refassgnexpr, sbsref->refassgnexpr,
diff --git a/src/backend/parser/parse_collate.c b/src/backend/parser/parse_collate.c
index d2e218353f3..be6dea6ffd2 100644
--- a/src/backend/parser/parse_collate.c
+++ b/src/backend/parser/parse_collate.c
@@ -680,11 +680,25 @@ assign_collations_walker(Node *node, assign_collations_context *context)
 							 * contribute anything.)
 							 */
 							SubscriptingRef *sbsref = (SubscriptingRef *) node;
+							ListCell   *lc;
+
+							/* skip String subscripts used for dot notation */
+							foreach(lc, sbsref->refupperindexpr)
+							{
+								Node	   *expr = lfirst(lc);
+
+								if (expr && !IsA(expr, String))
+									assign_expr_collations(context->pstate, expr);
+							}
+
+							foreach(lc, sbsref->reflowerindexpr)
+							{
+								Node	   *expr = lfirst(lc);
+
+								if (expr && !IsA(expr, String))
+									assign_expr_collations(context->pstate, expr);
+							}
 
-							assign_expr_collations(context->pstate,
-												   (Node *) sbsref->refupperindexpr);
-							assign_expr_collations(context->pstate,
-												   (Node *) sbsref->reflowerindexpr);
 							(void) assign_collations_walker((Node *) sbsref->refexpr,
 															&loccontext);
 							(void) assign_collations_walker((Node *) sbsref->refassgnexpr,
diff --git a/src/backend/parser/parse_expr.c b/src/backend/parser/parse_expr.c
index 2c0f4a50b21..afe953fdbea 100644
--- a/src/backend/parser/parse_expr.c
+++ b/src/backend/parser/parse_expr.c
@@ -461,19 +461,40 @@ transformIndirection(ParseState *pstate, A_Indirection *ind)
 		}
 		else
 		{
-			Node	   *newresult;
-
 			Assert(IsA(n, String));
+			subscripts = lappend(subscripts, n);
+		}
+	}
+
+	/* process trailing subscripts, if any */
+	while (subscripts)
+	{
+		Node	   *newresult = (Node *)
+			transformContainerSubscripts(pstate,
+										 result,
+										 exprType(result),
+										 exprTypmod(result),
+										 &subscripts,
+										 false,
+										 true);
+
+		if (!newresult)
+		{
+			/* generic subscripting failed */
+			Node	   *n;
+
+			Assert(subscripts);
+
+			n = linitial(subscripts);
 
-			/* process subscripts before this field selection */
-			while (subscripts)
-				result = (Node *) transformContainerSubscripts(pstate,
-															   result,
-															   exprType(result),
-															   exprTypmod(result),
-															   &subscripts,
-															   false);
+			if (!IsA(n, String))
+				ereport(ERROR,
+						(errcode(ERRCODE_DATATYPE_MISMATCH),
+						 errmsg("cannot subscript type %s because it does not support subscripting",
+								format_type_be(exprType(result))),
+						 parser_errposition(pstate, exprLocation(result))));
 
+			/* try to find function for field selection */
 			newresult = ParseFuncOrColumn(pstate,
 										  list_make1(n),
 										  list_make1(result),
@@ -481,19 +502,16 @@ transformIndirection(ParseState *pstate, A_Indirection *ind)
 										  NULL,
 										  false,
 										  location);
-			if (newresult == NULL)
+
+			if (!newresult)
 				unknown_attribute(pstate, result, strVal(n), location);
-			result = newresult;
+
+			/* consume field select */
+			subscripts = list_delete_first(subscripts);
 		}
+
+		result = newresult;
 	}
-	/* process trailing subscripts, if any */
-	while (subscripts)
-		result = (Node *) transformContainerSubscripts(pstate,
-													   result,
-													   exprType(result),
-													   exprTypmod(result),
-													   &subscripts,
-													   false);
 
 	return result;
 }
diff --git a/src/backend/parser/parse_node.c b/src/backend/parser/parse_node.c
index 19a6b678e67..c1f8055564f 100644
--- a/src/backend/parser/parse_node.c
+++ b/src/backend/parser/parse_node.c
@@ -245,13 +245,15 @@ transformContainerSubscripts(ParseState *pstate,
 							 Oid containerType,
 							 int32 containerTypMod,
 							 List **indirection,
-							 bool isAssignment)
+							 bool isAssignment,
+							 bool noError)
 {
 	SubscriptingRef *sbsref;
 	const SubscriptRoutines *sbsroutines;
 	Oid			elementType;
 	bool		isSlice = false;
 	ListCell   *idx;
+	int			indirection_length = list_length(*indirection);
 
 	/*
 	 * Determine the actual container type, smashing any domain.  In the
@@ -267,11 +269,16 @@ transformContainerSubscripts(ParseState *pstate,
 	 */
 	sbsroutines = getSubscriptingRoutines(containerType, &elementType);
 	if (!sbsroutines)
+	{
+		if (noError)
+			return NULL;
+
 		ereport(ERROR,
 				(errcode(ERRCODE_DATATYPE_MISMATCH),
 				 errmsg("cannot subscript type %s because it does not support subscripting",
 						format_type_be(containerType)),
 				 parser_errposition(pstate, exprLocation(containerBase))));
+	}
 
 	/*
 	 * Detect whether any of the indirection items are slice specifiers.
@@ -282,9 +289,9 @@ transformContainerSubscripts(ParseState *pstate,
 	 */
 	foreach(idx, *indirection)
 	{
-		A_Indices  *ai = lfirst_node(A_Indices, idx);
+		Node	   *ai = lfirst(idx);
 
-		if (ai->is_slice)
+		if (IsA(ai, A_Indices) && castNode(A_Indices, ai)->is_slice)
 		{
 			isSlice = true;
 			break;
@@ -312,6 +319,32 @@ transformContainerSubscripts(ParseState *pstate,
 	sbsroutines->transform(sbsref, indirection, pstate,
 						   isSlice, isAssignment);
 
+	/*
+	 * Error out, if datatyoe falied to consume any indirection elements.
+	 */
+	if (list_length(*indirection) == indirection_length)
+	{
+		Node	   *ind = linitial(*indirection);
+
+		if (noError)
+			return NULL;
+
+		if (IsA(ind, String))
+			ereport(ERROR,
+					(errcode(ERRCODE_DATATYPE_MISMATCH),
+					 errmsg("type %s does not support column notation",
+							format_type_be(containerType)),
+					 parser_errposition(pstate, exprLocation(containerBase))));
+		else if (IsA(ind, A_Indices))
+			ereport(ERROR,
+					(errcode(ERRCODE_DATATYPE_MISMATCH),
+					 errmsg("type %s does not support array subscripting",
+							format_type_be(containerType)),
+					 parser_errposition(pstate, exprLocation(containerBase))));
+		else
+			elog(ERROR, "invalid indirection operation: %d", nodeTag(ind));
+	}
+
 	/*
 	 * Verify we got a valid type (this defends, for example, against someone
 	 * using array_subscript_handler as typsubscript without setting typelem).
diff --git a/src/backend/parser/parse_target.c b/src/backend/parser/parse_target.c
index 39fd82f8371..5e126145ea5 100644
--- a/src/backend/parser/parse_target.c
+++ b/src/backend/parser/parse_target.c
@@ -937,7 +937,8 @@ transformAssignmentSubscripts(ParseState *pstate,
 										  containerType,
 										  containerTypMod,
 										  &subscripts,
-										  true);
+										  true,
+										  false);
 
 	if (subscripts)
 		elog(ERROR, "subscripting assignment is not supported");
diff --git a/src/backend/utils/adt/arraysubs.c b/src/backend/utils/adt/arraysubs.c
index edd85f7ba67..fe18df86e45 100644
--- a/src/backend/utils/adt/arraysubs.c
+++ b/src/backend/utils/adt/arraysubs.c
@@ -61,6 +61,7 @@ array_subscript_transform(SubscriptingRef *sbsref,
 	List	   *upperIndexpr = NIL;
 	List	   *lowerIndexpr = NIL;
 	ListCell   *idx;
+	int			ndim;
 
 	/*
 	 * Transform the subscript expressions, and separate upper and lower
@@ -72,9 +73,14 @@ array_subscript_transform(SubscriptingRef *sbsref,
 	 */
 	foreach(idx, *indirection)
 	{
-		A_Indices  *ai = lfirst_node(A_Indices, idx);
+		A_Indices  *ai;
 		Node	   *subexpr;
 
+		if (!IsA(lfirst(idx), A_Indices))
+			break;
+
+		ai = lfirst_node(A_Indices, idx);
+
 		if (isSlice)
 		{
 			if (ai->lidx)
@@ -144,14 +150,15 @@ array_subscript_transform(SubscriptingRef *sbsref,
 	sbsref->reflowerindexpr = lowerIndexpr;
 
 	/* Verify subscript list lengths are within implementation limit */
-	if (list_length(upperIndexpr) > MAXDIM)
+	ndim = list_length(upperIndexpr);
+	if (ndim > MAXDIM)
 		ereport(ERROR,
 				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
 				 errmsg("number of array dimensions (%d) exceeds the maximum allowed (%d)",
 						list_length(upperIndexpr), MAXDIM)));
 	/* We need not check lowerIndexpr separately */
 
-	*indirection = NIL;
+	*indirection = list_delete_first_n(*indirection, ndim);
 
 	/*
 	 * Determine the result type of the subscripting operation.  It's the same
diff --git a/src/backend/utils/adt/jsonbsubs.c b/src/backend/utils/adt/jsonbsubs.c
index 8ad6aa1ad4f..a0d38a0fd80 100644
--- a/src/backend/utils/adt/jsonbsubs.c
+++ b/src/backend/utils/adt/jsonbsubs.c
@@ -55,9 +55,14 @@ jsonb_subscript_transform(SubscriptingRef *sbsref,
 	 */
 	foreach(idx, *indirection)
 	{
-		A_Indices  *ai = lfirst_node(A_Indices, idx);
+		A_Indices  *ai;
 		Node	   *subExpr;
 
+		if (!IsA(lfirst(idx), A_Indices))
+			break;
+
+		ai = lfirst_node(A_Indices, idx);
+
 		if (isSlice)
 		{
 			Node	   *expr = ai->uidx ? ai->uidx : ai->lidx;
@@ -160,7 +165,9 @@ jsonb_subscript_transform(SubscriptingRef *sbsref,
 	sbsref->refrestype = JSONBOID;
 	sbsref->reftypmod = -1;
 
-	*indirection = NIL;
+	/* Remove processed elements */
+	if (upperIndexpr)
+		*indirection = list_delete_first_n(*indirection, list_length(upperIndexpr));
 }
 
 /*
diff --git a/src/backend/utils/adt/ruleutils.c b/src/backend/utils/adt/ruleutils.c
index 54dad975553..af5417d0859 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -47,6 +47,7 @@
 #include "nodes/makefuncs.h"
 #include "nodes/nodeFuncs.h"
 #include "nodes/pathnodes.h"
+#include "nodes/subscripting.h"
 #include "optimizer/optimizer.h"
 #include "parser/parse_agg.h"
 #include "parser/parse_func.h"
@@ -12916,17 +12917,29 @@ printSubscripts(SubscriptingRef *sbsref, deparse_context *context)
 	lowlist_item = list_head(sbsref->reflowerindexpr);	/* could be NULL */
 	foreach(uplist_item, sbsref->refupperindexpr)
 	{
-		appendStringInfoChar(buf, '[');
-		if (lowlist_item)
+		Node	   *up = (Node *) lfirst(uplist_item);
+
+		if (IsA(up, String))
+		{
+			appendStringInfoChar(buf, '.');
+			appendStringInfoString(buf, quote_identifier(strVal(up)));
+		}
+		else
 		{
+			appendStringInfoChar(buf, '[');
+			if (lowlist_item)
+			{
+				/* If subexpression is NULL, get_rule_expr prints nothing */
+				get_rule_expr((Node *) lfirst(lowlist_item), context, false);
+				appendStringInfoChar(buf, ':');
+			}
 			/* If subexpression is NULL, get_rule_expr prints nothing */
-			get_rule_expr((Node *) lfirst(lowlist_item), context, false);
-			appendStringInfoChar(buf, ':');
-			lowlist_item = lnext(sbsref->reflowerindexpr, lowlist_item);
+			get_rule_expr((Node *) lfirst(uplist_item), context, false);
+			appendStringInfoChar(buf, ']');
 		}
-		/* If subexpression is NULL, get_rule_expr prints nothing */
-		get_rule_expr((Node *) lfirst(uplist_item), context, false);
-		appendStringInfoChar(buf, ']');
+
+		if (lowlist_item)
+			lowlist_item = lnext(sbsref->reflowerindexpr, lowlist_item);
 	}
 }
 
diff --git a/src/include/parser/parse_node.h b/src/include/parser/parse_node.h
index 5ae11ccec33..71b04bd503c 100644
--- a/src/include/parser/parse_node.h
+++ b/src/include/parser/parse_node.h
@@ -378,7 +378,8 @@ extern SubscriptingRef *transformContainerSubscripts(ParseState *pstate,
 													 Oid containerType,
 													 int32 containerTypMod,
 													 List **indirection,
-													 bool isAssignment);
+													 bool isAssignment,
+													 bool noError);
 extern Const *make_const(ParseState *pstate, A_Const *aconst);
 
 #endif							/* PARSE_NODE_H */
-- 
2.39.5 (Apple Git-154)

