Further fixes for MULTIEXPR_SUBLINK fix.

Some more things I didn't think about in commits 3f7323cbb et al:

MULTIEXPR_SUBLINK subplans might have been converted to initplans
instead of regular subplans, in which case they won't show up in
the modified targetlist.  Fortunately, this would only happen if
they have no input parameters, which means that the problem we
originally needed to fix can't happen with them.  Therefore, there's
no need to clone their output parameters, and thus it doesn't hurt
that we'll fail to see them in the first pass over the targetlist.
Nonetheless, this complicates matters greatly, because now we have
to distinguish output Params of initplans (which shouldn't get
renumbered) from those of regular subplans (which should).
This also breaks the simplistic scheme I used of assuming that the
subplans found in the targetlist have consecutive subLinkIds.
We really can't avoid the need to know the subplans' subLinkIds in
this code.  To fix that, add subLinkId as the last field of SubPlan.
We can get away with that change in back branches because SubPlan
nodes will never be stored in the catalogs, and there's no ABI
break for external code that might be looking at the existing
fields of SubPlan.

Secondly, rewriteTargetListIU might have rolled up multiple
FieldStores or SubscriptingRefs into one targetlist entry,
breaking the assumption that there's at most one Param to fix
per targetlist entry.  (That assumption is OK I think in the
ruleutils.c code I stole the logic from in 18f51083c, because
that only deals with pre-rewrite query trees.  But it's
definitely not OK here.)  Abandon that shortcut and just do a
full tree walk on the targetlist to ensure we find all the
Params we have to change.

Per bug #17606 from Andre Lin.  As before, only v10-v13 need the
patch.

Discussion: https://postgr.es/m/17606-e5c8ad18d31db96a@postgresql.org

8 files changed, 119 insertions, 83 deletions

diff --git a/src/backend/nodes/copyfuncs.c b/src/backend/nodes/copyfuncs.c
index 8db09fda63..692b6c1559 100644
--- a/src/backend/nodes/copyfuncs.c
+++ b/src/backend/nodes/copyfuncs.c
@@ -1719,6 +1719,7 @@ _copySubPlan(const SubPlan *from)
 	COPY_NODE_FIELD(args);
 	COPY_SCALAR_FIELD(startup_cost);
 	COPY_SCALAR_FIELD(per_call_cost);
+	COPY_SCALAR_FIELD(subLinkId);
 
 	return newnode;
 }
diff --git a/src/backend/nodes/equalfuncs.c b/src/backend/nodes/equalfuncs.c
index 092fb37923..337be2f896 100644
--- a/src/backend/nodes/equalfuncs.c
+++ b/src/backend/nodes/equalfuncs.c
@@ -455,6 +455,7 @@ _equalSubPlan(const SubPlan *a, const SubPlan *b)
 	COMPARE_NODE_FIELD(args);
 	COMPARE_SCALAR_FIELD(startup_cost);
 	COMPARE_SCALAR_FIELD(per_call_cost);
+	COMPARE_SCALAR_FIELD(subLinkId);
 
 	return true;
 }
diff --git a/src/backend/nodes/outfuncs.c b/src/backend/nodes/outfuncs.c
index 4547eed43b..21ececf0c2 100644
--- a/src/backend/nodes/outfuncs.c
+++ b/src/backend/nodes/outfuncs.c
@@ -1349,6 +1349,7 @@ _outSubPlan(StringInfo str, const SubPlan *node)
 	WRITE_NODE_FIELD(args);
 	WRITE_FLOAT_FIELD(startup_cost, "%.2f");
 	WRITE_FLOAT_FIELD(per_call_cost, "%.2f");
+	WRITE_INT_FIELD(subLinkId);
 }
 
 static void
diff --git a/src/backend/nodes/readfuncs.c b/src/backend/nodes/readfuncs.c
index 62c945b6c5..7976b369ba 100644
--- a/src/backend/nodes/readfuncs.c
+++ b/src/backend/nodes/readfuncs.c
@@ -2536,6 +2536,7 @@ _readSubPlan(void)
 	READ_NODE_FIELD(args);
 	READ_FLOAT_FIELD(startup_cost);
 	READ_FLOAT_FIELD(per_call_cost);
+	READ_INT_FIELD(subLinkId);
 
 	READ_DONE();
 }
diff --git a/src/backend/optimizer/plan/subselect.c b/src/backend/optimizer/plan/subselect.c
index a873d72b3a..3e65af2dce 100644
--- a/src/backend/optimizer/plan/subselect.c
+++ b/src/backend/optimizer/plan/subselect.c
@@ -86,6 +86,7 @@ static bool subplan_is_hashable(Plan *plan);
 static bool testexpr_is_hashable(Node *testexpr, List *param_ids);
 static bool test_opexpr_is_hashable(OpExpr *testexpr, List *param_ids);
 static bool hash_ok_operator(OpExpr *expr);
+static bool SS_make_multiexprs_unique_walker(Node *node, void *context);
 static bool contain_dml(Node *node);
 static bool contain_dml_walker(Node *node, void *context);
 static bool contain_outer_selfref(Node *node);
@@ -335,6 +336,7 @@ build_subplan(PlannerInfo *root, Plan *plan, PlannerInfo *subroot,
 	 */
 	splan = makeNode(SubPlan);
 	splan->subLinkType = subLinkType;
+	splan->subLinkId = subLinkId;
 	splan->testexpr = NULL;
 	splan->paramIds = NIL;
 	get_first_col_type(plan, &splan->firstColType, &splan->firstColTypmod,
@@ -858,12 +860,17 @@ hash_ok_operator(OpExpr *expr)
  * SubPlans, inheritance_planner() must call this to assign new, unique Param
  * IDs to the cloned MULTIEXPR_SUBLINKs' output parameters.  See notes in
  * ExecScanSubPlan.
+ *
+ * We do not need to renumber Param IDs for MULTIEXPR_SUBLINK plans that are
+ * initplans, because those don't have input parameters that could cause
+ * confusion.  Such initplans will not appear in the targetlist anyway, but
+ * they still complicate matters because the surviving regular subplans might
+ * not have consecutive subLinkIds.
  */
 void
 SS_make_multiexprs_unique(PlannerInfo *root, PlannerInfo *subroot)
 {
-	List	   *new_multiexpr_params = NIL;
-	int			offset;
+	List	   *param_mapping = NIL;
 	ListCell   *lc;
 
 	/*
@@ -876,6 +883,9 @@ SS_make_multiexprs_unique(PlannerInfo *root, PlannerInfo *subroot)
 		SubPlan    *splan;
 		Plan	   *plan;
 		List	   *params;
+		int			oldId,
+					newId;
+		ListCell   *lc2;
 
 		if (!IsA(tent->expr, SubPlan))
 			continue;
@@ -898,86 +908,77 @@ SS_make_multiexprs_unique(PlannerInfo *root, PlannerInfo *subroot)
 										  &splan->setParam);
 
 		/*
-		 * We will append the replacement-Params lists to
-		 * root->multiexpr_params, but for the moment just make a local list.
-		 * Since we lack easy access here to the original subLinkId, we have
-		 * to fall back on the slightly shaky assumption that the MULTIEXPR
-		 * SubPlans appear in the targetlist in subLinkId order.  This should
-		 * be safe enough given the way that the parser builds the targetlist
-		 * today.  I wouldn't want to rely on it going forward, but since this
-		 * code has a limited lifespan it should be fine.  We can partially
-		 * protect against problems with assertions below.
+		 * Append the new replacement-Params list to root->multiexpr_params.
+		 * Then its index in that list becomes the new subLinkId of the
+		 * SubPlan.
 		 */
-		new_multiexpr_params = lappend(new_multiexpr_params, params);
+		root->multiexpr_params = lappend(root->multiexpr_params, params);
+		oldId = splan->subLinkId;
+		newId = list_length(root->multiexpr_params);
+		Assert(newId > oldId);
+		splan->subLinkId = newId;
+
+		/*
+		 * Add a mapping entry to param_mapping so that we can update the
+		 * associated Params below.  Leave zeroes in the list for any
+		 * subLinkIds we don't encounter; those must have been converted to
+		 * initplans.
+		 */
+		while (list_length(param_mapping) < oldId)
+			param_mapping = lappend_int(param_mapping, 0);
+		lc2 = list_nth_cell(param_mapping, oldId - 1);
+		lfirst_int(lc2) = newId;
 	}
 
 	/*
-	 * Now we must find the Param nodes that reference the MULTIEXPR outputs
-	 * and update their sublink IDs so they'll reference the new outputs.
-	 * Fortunately, those too must be in the cloned targetlist, but they could
-	 * be buried under FieldStores and SubscriptingRefs and CoerceToDomains
-	 * (cf processIndirection()), and underneath those there could be an
-	 * implicit type coercion.
+	 * Unless all the MULTIEXPRs were converted to initplans, we must now find
+	 * the Param nodes that reference the MULTIEXPR outputs and update their
+	 * sublink IDs so they'll reference the new outputs.  While such Params
+	 * must be in the cloned targetlist, they could be buried under stuff such
+	 * as FieldStores and SubscriptingRefs and type coercions.
 	 */
-	offset = list_length(root->multiexpr_params);
+	if (param_mapping != NIL)
+		SS_make_multiexprs_unique_walker((Node *) subroot->parse->targetList,
+										 (void *) param_mapping);
+}
 
-	foreach(lc, subroot->parse->targetList)
+/*
+ * Locate PARAM_MULTIEXPR Params in an expression tree, and update as needed.
+ * (We can update-in-place because the tree was already copied.)
+ */
+static bool
+SS_make_multiexprs_unique_walker(Node *node, void *context)
+{
+	if (node == NULL)
+		return false;
+	if (IsA(node, Param))
 	{
-		TargetEntry *tent = (TargetEntry *) lfirst(lc);
-		Node	   *expr;
-		Param	   *p;
+		Param	   *p = (Param *) node;
+		List	   *param_mapping = (List *) context;
 		int			subqueryid;
 		int			colno;
+		int			newId;
 
-		expr = (Node *) tent->expr;
-		while (expr)
-		{
-			if (IsA(expr, FieldStore))
-			{
-				FieldStore *fstore = (FieldStore *) expr;
-
-				expr = (Node *) linitial(fstore->newvals);
-			}
-			else if (IsA(expr, SubscriptingRef))
-			{
-				SubscriptingRef *sbsref = (SubscriptingRef *) expr;
-
-				if (sbsref->refassgnexpr == NULL)
-					break;
-
-				expr = (Node *) sbsref->refassgnexpr;
-			}
-			else if (IsA(expr, CoerceToDomain))
-			{
-				CoerceToDomain *cdomain = (CoerceToDomain *) expr;
-
-				if (cdomain->coercionformat != COERCE_IMPLICIT_CAST)
-					break;
-				expr = (Node *) cdomain->arg;
-			}
-			else
-				break;
-		}
-		expr = strip_implicit_coercions(expr);
-		if (expr == NULL || !IsA(expr, Param))
-			continue;
-		p = (Param *) expr;
 		if (p->paramkind != PARAM_MULTIEXPR)
-			continue;
+			return false;
 		subqueryid = p->paramid >> 16;
 		colno = p->paramid & 0xFFFF;
-		Assert(subqueryid > 0 &&
-			   subqueryid <= list_length(new_multiexpr_params));
-		Assert(colno > 0 &&
-			   colno <= list_length((List *) list_nth(new_multiexpr_params,
-													  subqueryid - 1)));
-		subqueryid += offset;
-		p->paramid = (subqueryid << 16) + colno;
-	}
 
-	/* Finally, attach new replacement lists to the global list */
-	root->multiexpr_params = list_concat(root->multiexpr_params,
-										 new_multiexpr_params);
+		/*
+		 * If subqueryid doesn't have a mapping entry, it must refer to an
+		 * initplan, so don't change the Param.
+		 */
+		Assert(subqueryid > 0);
+		if (subqueryid > list_length(param_mapping))
+			return false;
+		newId = list_nth_int(param_mapping, subqueryid - 1);
+		if (newId == 0)
+			return false;
+		p->paramid = (newId << 16) + colno;
+		return false;
+	}
+	return expression_tree_walker(node, SS_make_multiexprs_unique_walker,
+								  context);
 }
 
 
@@ -1110,6 +1111,7 @@ SS_process_ctes(PlannerInfo *root)
 		 */
 		splan = makeNode(SubPlan);
 		splan->subLinkType = CTE_SUBLINK;
+		splan->subLinkId = 0;
 		splan->testexpr = NULL;
 		splan->paramIds = NIL;
 		get_first_col_type(plan, &splan->firstColType, &splan->firstColTypmod,
@@ -3075,6 +3077,7 @@ SS_make_initplan_from_plan(PlannerInfo *root,
 	 */
 	node = makeNode(SubPlan);
 	node->subLinkType = EXPR_SUBLINK;
+	node->subLinkId = 0;
 	node->plan_id = list_length(root->glob->subplans);
 	node->plan_name = psprintf("InitPlan %d (returns $%d)",
 							   node->plan_id, prm->paramid);
diff --git a/src/include/nodes/primnodes.h b/src/include/nodes/primnodes.h
index d73be2ad46..cfc4f5a9a7 100644
--- a/src/include/nodes/primnodes.h
+++ b/src/include/nodes/primnodes.h
@@ -731,6 +731,8 @@ typedef struct SubPlan
 	/* Estimated execution costs: */
 	Cost		startup_cost;	/* one-time setup cost */
 	Cost		per_call_cost;	/* cost for each subplan evaluation */
+	/* Copied from original SubLink, but placed at end for ABI stability */
+	int			subLinkId;		/* ID (1..n); 0 if not MULTIEXPR */
 } SubPlan;
 
 /*
diff --git a/src/test/regress/expected/inherit.out b/src/test/regress/expected/inherit.out
index ed99bdde7f..658fed7922 100644
--- a/src/test/regress/expected/inherit.out
+++ b/src/test/regress/expected/inherit.out
@@ -1717,23 +1717,36 @@ reset enable_bitmapscan;
 --
 -- Check handling of MULTIEXPR SubPlans in inherited updates
 --
-create table inhpar(f1 int, f2 text[]);
+create table inhpar(f1 int, f2 text[], f3 int);
 insert into inhpar select generate_series(1,10);
 create table inhcld() inherits(inhpar);
 insert into inhcld select generate_series(11,10000);
 vacuum analyze inhcld;
 vacuum analyze inhpar;
 explain (verbose, costs off)
-update inhpar set (f1, f2[1]) = (select p2.unique2, p2.stringu1
-                                 from int4_tbl limit 1)
+update inhpar set
+  (f1, f2[1])    = (select p2.unique2, p2.stringu1 from int4_tbl limit 1),
+  (f2[2], f2[3]) = (select 'x', 'y' from int4_tbl limit 1),
+  (f3, f2[4])    = (select p2.unique2, p2.stringu1 from int4_tbl limit 1),
+  (f2[5], f2[6]) = (select 'x', 'y' from int4_tbl limit 1)
 from onek p2 where inhpar.f1 = p2.unique1;
-                                          QUERY PLAN                                           
------------------------------------------------------------------------------------------------
+                                                                                                       QUERY PLAN                                                                                                        
+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
  Update on public.inhpar
    Update on public.inhpar
    Update on public.inhcld inhpar_1
+   InitPlan 2 (returns $4,$5)
+     ->  Limit
+           Output: 'x'::text, 'y'::text
+           ->  Seq Scan on public.int4_tbl int4_tbl_1
+                 Output: 'x'::text, 'y'::text
+   InitPlan 4 (returns $10,$11)
+     ->  Limit
+           Output: 'x'::text, 'y'::text
+           ->  Seq Scan on public.int4_tbl int4_tbl_3
+                 Output: 'x'::text, 'y'::text
    ->  Merge Join
-         Output: $4, inhpar.f2[1] := $5, (SubPlan 1 (returns $2,$3)), inhpar.ctid, p2.ctid
+         Output: $12, (((((inhpar.f2[1] := $13)[2] := $4)[3] := $5)[4] := $15)[5] := $10)[6] := $11, $14, (SubPlan 1 (returns $2,$3)), NULL::record, (SubPlan 3 (returns $8,$9)), NULL::record, inhpar.ctid, p2.ctid
          Merge Cond: (p2.unique1 = inhpar.f1)
          ->  Index Scan using onek_unique1 on public.onek p2
                Output: p2.unique2, p2.stringu1, p2.ctid, p2.unique1
@@ -1747,8 +1760,13 @@ from onek p2 where inhpar.f1 = p2.unique1;
                  Output: (p2.unique2), (p2.stringu1)
                  ->  Seq Scan on public.int4_tbl
                        Output: p2.unique2, p2.stringu1
+         SubPlan 3 (returns $8,$9)
+           ->  Limit
+                 Output: (p2.unique2), (p2.stringu1)
+                 ->  Seq Scan on public.int4_tbl int4_tbl_2
+                       Output: p2.unique2, p2.stringu1
    ->  Hash Join
-         Output: $6, inhpar_1.f2[1] := $7, (SubPlan 1 (returns $2,$3)), inhpar_1.ctid, p2.ctid
+         Output: $16, (((((inhpar_1.f2[1] := $17)[2] := $4)[3] := $5)[4] := $19)[5] := $10)[6] := $11, $18, (SubPlan 1 (returns $2,$3)), NULL::record, (SubPlan 3 (returns $8,$9)), NULL::record, inhpar_1.ctid, p2.ctid
          Hash Cond: (inhpar_1.f1 = p2.unique1)
          ->  Seq Scan on public.inhcld inhpar_1
                Output: inhpar_1.f2, inhpar_1.ctid, inhpar_1.f1
@@ -1756,10 +1774,13 @@ from onek p2 where inhpar.f1 = p2.unique1;
                Output: p2.unique2, p2.stringu1, p2.ctid, p2.unique1
                ->  Seq Scan on public.onek p2
                      Output: p2.unique2, p2.stringu1, p2.ctid, p2.unique1
-(27 rows)
+(42 rows)
 
-update inhpar set (f1, f2[1]) = (select p2.unique2, p2.stringu1
-                                 from int4_tbl limit 1)
+update inhpar set
+  (f1, f2[1])    = (select p2.unique2, p2.stringu1 from int4_tbl limit 1),
+  (f2[2], f2[3]) = (select 'x', 'y' from int4_tbl limit 1),
+  (f3, f2[4])    = (select p2.unique2, p2.stringu1 from int4_tbl limit 1),
+  (f2[5], f2[6]) = (select 'x', 'y' from int4_tbl limit 1)
 from onek p2 where inhpar.f1 = p2.unique1;
 drop table inhpar cascade;
 NOTICE:  drop cascades to table inhcld
diff --git a/src/test/regress/sql/inherit.sql b/src/test/regress/sql/inherit.sql
index 7d3813704e..a4a33a3da8 100644
--- a/src/test/regress/sql/inherit.sql
+++ b/src/test/regress/sql/inherit.sql
@@ -632,7 +632,7 @@ reset enable_bitmapscan;
 --
 -- Check handling of MULTIEXPR SubPlans in inherited updates
 --
-create table inhpar(f1 int, f2 text[]);
+create table inhpar(f1 int, f2 text[], f3 int);
 insert into inhpar select generate_series(1,10);
 create table inhcld() inherits(inhpar);
 insert into inhcld select generate_series(11,10000);
@@ -640,11 +640,17 @@ vacuum analyze inhcld;
 vacuum analyze inhpar;
 
 explain (verbose, costs off)
-update inhpar set (f1, f2[1]) = (select p2.unique2, p2.stringu1
-                                 from int4_tbl limit 1)
+update inhpar set
+  (f1, f2[1])    = (select p2.unique2, p2.stringu1 from int4_tbl limit 1),
+  (f2[2], f2[3]) = (select 'x', 'y' from int4_tbl limit 1),
+  (f3, f2[4])    = (select p2.unique2, p2.stringu1 from int4_tbl limit 1),
+  (f2[5], f2[6]) = (select 'x', 'y' from int4_tbl limit 1)
 from onek p2 where inhpar.f1 = p2.unique1;
-update inhpar set (f1, f2[1]) = (select p2.unique2, p2.stringu1
-                                 from int4_tbl limit 1)
+update inhpar set
+  (f1, f2[1])    = (select p2.unique2, p2.stringu1 from int4_tbl limit 1),
+  (f2[2], f2[3]) = (select 'x', 'y' from int4_tbl limit 1),
+  (f3, f2[4])    = (select p2.unique2, p2.stringu1 from int4_tbl limit 1),
+  (f2[5], f2[6]) = (select 'x', 'y' from int4_tbl limit 1)
 from onek p2 where inhpar.f1 = p2.unique1;
 
 drop table inhpar cascade;