From 0a7fed14e036a214daa668004daa97a023e7cc90 Mon Sep 17 00:00:00 2001
From: Dmitrii Dolgov <9erthalion6@gmail.com>
Date: Thu, 8 May 2025 16:41:01 +0200
Subject: [PATCH v3 2/3] Introduce LocationExpr

Add LocationExpr wrapper node to capture location and length of an
expression in a query. Use it in to wrap expr_list in in_expr and
array_expr conveying location information to ArrayExpr.
---
 src/backend/nodes/nodeFuncs.c    | 23 +++++++++++++++
 src/backend/parser/gram.y        | 31 +++++++++++++++++----
 src/backend/parser/parse_expr.c  | 48 ++++++++++++++++++++++++++++++--
 src/include/nodes/parsenodes.h   | 17 ++++++++++-
 src/include/nodes/primnodes.h    |  7 +++++
 src/tools/pgindent/typedefs.list |  1 +
 6 files changed, 118 insertions(+), 9 deletions(-)

diff --git a/src/backend/nodes/nodeFuncs.c b/src/backend/nodes/nodeFuncs.c
index 7bc823507f1..f0b05630fd1 100644
--- a/src/backend/nodes/nodeFuncs.c
+++ b/src/backend/nodes/nodeFuncs.c
@@ -284,6 +284,12 @@ exprType(const Node *expr)
 		case T_PlaceHolderVar:
 			type = exprType((Node *) ((const PlaceHolderVar *) expr)->phexpr);
 			break;
+		case T_LocationExpr:
+			{
+				const LocationExpr *n = (const LocationExpr *) expr;
+				type = exprType((Node *) n->expr);
+			}
+			break;
 		default:
 			elog(ERROR, "unrecognized node type: %d", (int) nodeTag(expr));
 			type = InvalidOid;	/* keep compiler quiet */
@@ -536,6 +542,11 @@ exprTypmod(const Node *expr)
 			return exprTypmod((Node *) ((const ReturningExpr *) expr)->retexpr);
 		case T_PlaceHolderVar:
 			return exprTypmod((Node *) ((const PlaceHolderVar *) expr)->phexpr);
+		case T_LocationExpr:
+			{
+				const LocationExpr *n = (const LocationExpr *) expr;
+				return exprTypmod((Node *) n->expr);
+			}
 		default:
 			break;
 	}
@@ -1058,6 +1069,9 @@ exprCollation(const Node *expr)
 		case T_PlaceHolderVar:
 			coll = exprCollation((Node *) ((const PlaceHolderVar *) expr)->phexpr);
 			break;
+		case T_LocationExpr:
+			coll = exprCollation((Node *) ((const LocationExpr *) expr)->expr);
+			break;
 		default:
 			elog(ERROR, "unrecognized node type: %d", (int) nodeTag(expr));
 			coll = InvalidOid;	/* keep compiler quiet */
@@ -1306,6 +1320,10 @@ exprSetCollation(Node *expr, Oid collation)
 			/* NextValueExpr's result is an integer type ... */
 			Assert(!OidIsValid(collation)); /* ... so never set a collation */
 			break;
+		case T_LocationExpr:
+			exprSetCollation((Node *) ((LocationExpr *) expr)->expr,
+							 collation);
+			break;
 		default:
 			elog(ERROR, "unrecognized node type: %d", (int) nodeTag(expr));
 			break;
@@ -1803,6 +1821,9 @@ exprLocation(const Node *expr)
 		case T_PartitionRangeDatum:
 			loc = ((const PartitionRangeDatum *) expr)->location;
 			break;
+		case T_LocationExpr:
+			loc = ((const LocationExpr *) expr)->location;
+			break;
 		default:
 			/* for any other node type it's just unknown... */
 			loc = -1;
@@ -4705,6 +4726,8 @@ raw_expression_tree_walker_impl(Node *node,
 					return true;
 			}
 			break;
+		case T_LocationExpr:
+			return WALK(((LocationExpr *) node)->expr);
 		default:
 			elog(ERROR, "unrecognized node type: %d",
 				 (int) nodeTag(node));
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 0b5652071d1..293a81b29b7 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -184,7 +184,8 @@ static void doNegateFloat(Float *v);
 static Node *makeAndExpr(Node *lexpr, Node *rexpr, int location);
 static Node *makeOrExpr(Node *lexpr, Node *rexpr, int location);
 static Node *makeNotExpr(Node *expr, int location);
-static Node *makeAArrayExpr(List *elements, int location);
+static Node *makeAArrayExpr(Node *elements, int location);
+static Node *makeLocationExpr(Node *expr, int location, int length);
 static Node *makeSQLValueFunction(SQLValueFunctionOp op, int32 typmod,
 								  int location);
 static Node *makeXmlExpr(XmlExprOp op, char *name, List *named_args,
@@ -16764,15 +16765,18 @@ type_list:	Typename								{ $$ = list_make1($1); }
 
 array_expr: '[' expr_list ']'
 				{
-					$$ = makeAArrayExpr($2, @1);
+					Node *locExpr = makeLocationExpr((Node *) $2, @1, @3);
+					$$ = makeAArrayExpr(locExpr, @1);
 				}
 			| '[' array_expr_list ']'
 				{
-					$$ = makeAArrayExpr($2, @1);
+					Node *locExpr = makeLocationExpr((Node *) $2, @1, @3);
+					$$ = makeAArrayExpr(locExpr, @1);
 				}
 			| '[' ']'
 				{
-					$$ = makeAArrayExpr(NIL, @1);
+					Node *locExpr =  makeLocationExpr((Node *) NIL, @1, @2);
+					$$ = makeAArrayExpr(locExpr, @1);
 				}
 		;
 
@@ -16902,7 +16906,10 @@ in_expr:	select_with_parens
 					/* other fields will be filled later */
 					$$ = (Node *) n;
 				}
-			| '(' expr_list ')'						{ $$ = (Node *) $2; }
+			| '(' expr_list ')'
+				{
+					$$ = (Node *) makeLocationExpr((Node *) $2, @1, @3);
+				}
 		;
 
 /*
@@ -19300,7 +19307,7 @@ makeNotExpr(Node *expr, int location)
 }
 
 static Node *
-makeAArrayExpr(List *elements, int location)
+makeAArrayExpr(Node *elements, int location)
 {
 	A_ArrayExpr *n = makeNode(A_ArrayExpr);
 
@@ -19309,6 +19316,18 @@ makeAArrayExpr(List *elements, int location)
 	return (Node *) n;
 }
 
+static Node *
+makeLocationExpr(Node *expr, int start_location, int end_location)
+{
+	LocationExpr 	*n = makeNode(LocationExpr);
+
+	n->expr = expr;
+	n->location = start_location + 1;
+	n->length = end_location - start_location - 1;
+
+	return (Node *) n;
+}
+
 static Node *
 makeSQLValueFunction(SQLValueFunctionOp op, int32 typmod, int location)
 {
diff --git a/src/backend/parser/parse_expr.c b/src/backend/parser/parse_expr.c
index 1f8e2d54673..b48beff157e 100644
--- a/src/backend/parser/parse_expr.c
+++ b/src/backend/parser/parse_expr.c
@@ -370,6 +370,25 @@ transformExprRecurse(ParseState *pstate, Node *expr)
 			result = transformJsonFuncExpr(pstate, (JsonFuncExpr *) expr);
 			break;
 
+		case T_LocationExpr:
+			{
+				LocationExpr *loc = (LocationExpr *) expr;
+				if (!IsA(loc->expr, List))
+					result = transformExprRecurse(pstate, loc->expr);
+				else
+					result = loc->expr;
+
+				if (IsA(result, ArrayExpr))
+				{
+					ArrayExpr *arr = (ArrayExpr *) result;
+					arr->loc_range = list_make2_int(loc->location,
+												    loc->length);
+
+					result = (Node *) arr;
+				}
+			}
+			break;
+
 		default:
 			/* should not reach here */
 			elog(ERROR, "unrecognized node type: %d", (int) nodeTag(expr));
@@ -1125,6 +1144,7 @@ transformAExprIn(ParseState *pstate, A_Expr *a)
 {
 	Node	   *result = NULL;
 	Node	   *lexpr;
+	LocationExpr *location = NULL;
 	List	   *rexprs;
 	List	   *rvars;
 	List	   *rnonvars;
@@ -1139,6 +1159,9 @@ transformAExprIn(ParseState *pstate, A_Expr *a)
 	else
 		useOr = true;
 
+	if (IsA(a->rexpr, LocationExpr))
+		location = (LocationExpr *) a->rexpr;
+
 	/*
 	 * We try to generate a ScalarArrayOpExpr from IN/NOT IN, but this is only
 	 * possible if there is a suitable array type available.  If not, we fall
@@ -1152,7 +1175,7 @@ transformAExprIn(ParseState *pstate, A_Expr *a)
 	 */
 	lexpr = transformExprRecurse(pstate, a->lexpr);
 	rexprs = rvars = rnonvars = NIL;
-	foreach(l, (List *) a->rexpr)
+	foreach(l, (List *) transformExprRecurse(pstate, a->rexpr))
 	{
 		Node	   *rexpr = transformExprRecurse(pstate, lfirst(l));
 
@@ -1224,6 +1247,12 @@ transformAExprIn(ParseState *pstate, A_Expr *a)
 			newa->elements = aexprs;
 			newa->multidims = false;
 			newa->location = -1;
+			if (location)
+				newa->loc_range = list_make2_int(location->location,
+											     location->length);
+			else
+				newa->loc_range = NIL;
+
 
 			result = (Node *) make_scalar_array_op(pstate,
 												   a->name,
@@ -2014,12 +2043,22 @@ transformArrayExpr(ParseState *pstate, A_ArrayExpr *a,
 				   Oid array_type, Oid element_type, int32 typmod)
 {
 	ArrayExpr  *newa = makeNode(ArrayExpr);
+	List 	   *elements = NIL;
 	List	   *newelems = NIL;
 	List	   *newcoercedelems = NIL;
 	ListCell   *element;
+	LocationExpr *locExpr = NULL;
 	Oid			coerce_type;
 	bool		coerce_hard;
 
+	if (IsA(a->elements, LocationExpr))
+	{
+		locExpr = (LocationExpr *) a->elements;
+		elements = (List *) locExpr->expr;
+	}
+	else
+		elements = (List *) a->elements;
+
 	/*
 	 * Transform the element expressions
 	 *
@@ -2027,7 +2066,7 @@ transformArrayExpr(ParseState *pstate, A_ArrayExpr *a,
 	 * element expression.
 	 */
 	newa->multidims = false;
-	foreach(element, a->elements)
+	foreach(element, elements)
 	{
 		Node	   *e = (Node *) lfirst(element);
 		Node	   *newe;
@@ -2166,6 +2205,11 @@ transformArrayExpr(ParseState *pstate, A_ArrayExpr *a,
 	newa->element_typeid = element_type;
 	newa->elements = newcoercedelems;
 	newa->location = a->location;
+	if (locExpr)
+		newa->loc_range = list_make2_int(locExpr->location,
+										 locExpr->length);
+	else
+		newa->loc_range = NIL;
 
 	return (Node *) newa;
 }
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 4610fc61293..f3e4ba47af1 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -500,10 +500,25 @@ typedef struct A_Indirection
 typedef struct A_ArrayExpr
 {
 	NodeTag		type;
-	List	   *elements;		/* array element expressions */
+	Node	   *elements;		/* array element expressions */
 	ParseLoc	location;		/* token location, or -1 if unknown */
 } A_ArrayExpr;
 
+/*
+ * A wrapper expression to record start and end location
+ */
+typedef struct LocationExpr
+{
+	NodeTag		type;
+
+	/* the node to be wrapped */
+	Node	   *expr;
+	/* token location, or -1 if unknown */
+	ParseLoc	location;
+	/* token length, or -1 if unknown */
+	ParseLoc	length;
+} LocationExpr;
+
 /*
  * ResTarget -
  *	  result target (used in target list of pre-transformed parse trees)
diff --git a/src/include/nodes/primnodes.h b/src/include/nodes/primnodes.h
index 7d3b4198f26..60dec576908 100644
--- a/src/include/nodes/primnodes.h
+++ b/src/include/nodes/primnodes.h
@@ -1399,6 +1399,13 @@ typedef struct ArrayExpr
 	bool		multidims pg_node_attr(query_jumble_ignore);
 	/* token location, or -1 if unknown */
 	ParseLoc	location;
+
+	/*
+	 * Pair (location, length) for list of elements. Note, that location field
+	 * cannot always be used here instead, since it could be unknown, e.g. if
+	 * the node was created in transformAExprIn.
+	 */
+	List 		*loc_range pg_node_attr(query_jumble_ignore);
 } ArrayExpr;
 
 /*
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 9ea573fae21..1f911b24edd 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -1502,6 +1502,7 @@ LOCALLOCK
 LOCALLOCKOWNER
 LOCALLOCKTAG
 LOCALPREDICATELOCK
+LocationExpr
 LOCK
 LOCKMASK
 LOCKMETHODID
-- 
2.45.1

