Eliminate code duplication in replace_rte_variables callbacks

The callback functions ReplaceVarsFromTargetList_callback and
pullup_replace_vars_callback are both used to replace Vars in an
expression tree that reference a particular RTE with items from a
targetlist, and they both need to expand whole-tuple references and
deal with OLD/NEW RETURNING list Vars.  As a result, currently there
is significant code duplication between these two functions.

This patch introduces a new function, ReplaceVarFromTargetList, to
perform the replacement and calls it from both callback functions,
thereby eliminating code duplication.

Author: Dean Rasheed <dean.a.rasheed@gmail.com>
Author: Richard Guo <guofenglinux@gmail.com>
Reviewed-by: Jian He <jian.universality@gmail.com>
Discussion: https://postgr.es/m/CAEZATCWhr=FM4X5kCPvVs-g2XEk+ceLsNtBK_zZMkqFn9vUjsw@mail.gmail.com

3 files changed, 96 insertions, 142 deletions

diff --git a/src/backend/optimizer/prep/prepjointree.c b/src/backend/optimizer/prep/prepjointree.c
index 8cdacb6aa63..bcc40dd5a84 100644
--- a/src/backend/optimizer/prep/prepjointree.c
+++ b/src/backend/optimizer/prep/prepjointree.c
@@ -2666,131 +2666,38 @@ pullup_replace_vars_callback(Var *var,
 		/* Just copy the entry and fall through to adjust phlevelsup etc */
 		newnode = copyObject(rcon->rv_cache[varattno]);
 	}
-	else if (varattno == InvalidAttrNumber)
+	else
 	{
-		/* Must expand whole-tuple reference into RowExpr */
-		RowExpr    *rowexpr;
-		List	   *colnames;
-		List	   *fields;
-		bool		save_wrap_non_vars = rcon->wrap_non_vars;
-		int			save_sublevelsup = context->sublevels_up;
-
 		/*
-		 * If generating an expansion for a var of a named rowtype (ie, this
-		 * is a plain relation RTE), then we must include dummy items for
-		 * dropped columns.  If the var is RECORD (ie, this is a JOIN), then
-		 * omit dropped columns.  In the latter case, attach column names to
-		 * the RowExpr for use of the executor and ruleutils.c.
-		 *
-		 * In order to be able to cache the results, we always generate the
-		 * expansion with varlevelsup = 0, and then adjust below if needed.
+		 * Generate the replacement expression.  This takes care of expanding
+		 * wholerow references and dealing with non-default varreturningtype.
 		 */
-		expandRTE(rcon->target_rte,
-				  var->varno, 0 /* not varlevelsup */ ,
-				  var->varreturningtype, var->location,
-				  (var->vartype != RECORDOID),
-				  &colnames, &fields);
-		/* Expand the generated per-field Vars, but don't insert PHVs there */
-		rcon->wrap_non_vars = false;
-		context->sublevels_up = 0;	/* to match the expandRTE output */
-		fields = (List *) replace_rte_variables_mutator((Node *) fields,
-														context);
-		rcon->wrap_non_vars = save_wrap_non_vars;
-		context->sublevels_up = save_sublevelsup;
-
-		rowexpr = makeNode(RowExpr);
-		rowexpr->args = fields;
-		rowexpr->row_typeid = var->vartype;
-		rowexpr->row_format = COERCE_IMPLICIT_CAST;
-		rowexpr->colnames = (var->vartype == RECORDOID) ? colnames : NIL;
-		rowexpr->location = var->location;
-		newnode = (Node *) rowexpr;
-
-		/* Handle any OLD/NEW RETURNING list Vars */
-		if (var->varreturningtype != VAR_RETURNING_DEFAULT)
-		{
-			/*
-			 * Wrap the RowExpr in a ReturningExpr node, so that the executor
-			 * returns NULL if the OLD/NEW row does not exist.
-			 */
-			ReturningExpr *rexpr = makeNode(ReturningExpr);
-
-			rexpr->retlevelsup = 0;
-			rexpr->retold = (var->varreturningtype == VAR_RETURNING_OLD);
-			rexpr->retexpr = (Expr *) newnode;
-
-			newnode = (Node *) rexpr;
-		}
-
-		/*
-		 * Insert PlaceHolderVar if needed.  Notice that we are wrapping one
-		 * PlaceHolderVar around the whole RowExpr, rather than putting one
-		 * around each element of the row.  This is because we need the
-		 * expression to yield NULL, not ROW(NULL,NULL,...) when it is forced
-		 * to null by an outer join.
-		 */
-		if (need_phv)
-		{
-			newnode = (Node *)
-				make_placeholder_expr(rcon->root,
-									  (Expr *) newnode,
-									  bms_make_singleton(rcon->varno));
-			/* cache it with the PHV, and with phlevelsup etc not set yet */
-			rcon->rv_cache[InvalidAttrNumber] = copyObject(newnode);
-		}
-	}
-	else
-	{
-		/* Normal case referencing one targetlist element */
-		TargetEntry *tle = get_tle_by_resno(rcon->targetlist, varattno);
-
-		if (tle == NULL)		/* shouldn't happen */
-			elog(ERROR, "could not find attribute %d in subquery targetlist",
-				 varattno);
-
-		/* Make a copy of the tlist item to return */
-		newnode = (Node *) copyObject(tle->expr);
-
-		/* Handle any OLD/NEW RETURNING list Vars */
-		if (var->varreturningtype != VAR_RETURNING_DEFAULT)
-		{
-			/*
-			 * Copy varreturningtype onto any Vars in the tlist item that
-			 * refer to result_relation (which had better be non-zero).
-			 */
-			if (rcon->result_relation == 0)
-				elog(ERROR, "variable returning old/new found outside RETURNING list");
-
-			SetVarReturningType((Node *) newnode, rcon->result_relation,
-								0, var->varreturningtype);
-
-			/*
-			 * If the replacement expression in the targetlist is not simply a
-			 * Var referencing result_relation, wrap it in a ReturningExpr
-			 * node, so that the executor returns NULL if the OLD/NEW row does
-			 * not exist.
-			 */
-			if (!IsA(newnode, Var) ||
-				((Var *) newnode)->varno != rcon->result_relation ||
-				((Var *) newnode)->varlevelsup != 0)
-			{
-				ReturningExpr *rexpr = makeNode(ReturningExpr);
-
-				rexpr->retlevelsup = 0;
-				rexpr->retold = (var->varreturningtype == VAR_RETURNING_OLD);
-				rexpr->retexpr = (Expr *) newnode;
-
-				newnode = (Node *) rexpr;
-			}
-		}
+		newnode = ReplaceVarFromTargetList(var,
+										   rcon->target_rte,
+										   rcon->targetlist,
+										   rcon->result_relation,
+										   REPLACEVARS_REPORT_ERROR,
+										   0);
 
 		/* Insert PlaceHolderVar if needed */
 		if (need_phv)
 		{
 			bool		wrap;
 
-			if (newnode && IsA(newnode, Var) &&
-				((Var *) newnode)->varlevelsup == 0)
+			if (varattno == InvalidAttrNumber)
+			{
+				/*
+				 * Insert PlaceHolderVar for whole-tuple reference.  Notice
+				 * that we are wrapping one PlaceHolderVar around the whole
+				 * RowExpr, rather than putting one around each element of the
+				 * row.  This is because we need the expression to yield NULL,
+				 * not ROW(NULL,NULL,...) when it is forced to null by an
+				 * outer join.
+				 */
+				wrap = true;
+			}
+			else if (newnode && IsA(newnode, Var) &&
+					 ((Var *) newnode)->varlevelsup == 0)
 			{
 				/*
 				 * Simple Vars always escape being wrapped, unless they are
@@ -2936,7 +2843,7 @@ pullup_replace_vars_callback(Var *var,
 				 * Cache it if possible (ie, if the attno is in range, which
 				 * it probably always should be).
 				 */
-				if (varattno > InvalidAttrNumber &&
+				if (varattno >= InvalidAttrNumber &&
 					varattno <= list_length(rcon->targetlist))
 					rcon->rv_cache[varattno] = copyObject(newnode);
 			}
diff --git a/src/backend/rewrite/rewriteManip.c b/src/backend/rewrite/rewriteManip.c
index 6994b8c5425..98a265cd3d5 100644
--- a/src/backend/rewrite/rewriteManip.c
+++ b/src/backend/rewrite/rewriteManip.c
@@ -1010,7 +1010,7 @@ SetVarReturningType_walker(Node *node, SetVarReturningType_context *context)
 	return expression_tree_walker(node, SetVarReturningType_walker, context);
 }
 
-void
+static void
 SetVarReturningType(Node *node, int result_relation, int sublevels_up,
 					VarReturningType returning_type)
 {
@@ -1797,6 +1797,11 @@ map_variable_attnos(Node *node,
  * referencing result_relation, it is wrapped in a ReturningExpr node (causing
  * the executor to return NULL if the OLD/NEW row doesn't exist).
  *
+ * Note that ReplaceVarFromTargetList always generates the replacement
+ * expression with varlevelsup = 0.  The caller is responsible for adjusting
+ * the varlevelsup if needed.  This simplifies the caller's life if it wants to
+ * cache the replacement expressions.
+ *
  * outer_hasSubLinks works the same as for replace_rte_variables().
  */
 
@@ -1814,6 +1819,30 @@ ReplaceVarsFromTargetList_callback(Var *var,
 								   replace_rte_variables_context *context)
 {
 	ReplaceVarsFromTargetList_context *rcon = (ReplaceVarsFromTargetList_context *) context->callback_arg;
+	Node	   *newnode;
+
+	newnode = ReplaceVarFromTargetList(var,
+									   rcon->target_rte,
+									   rcon->targetlist,
+									   rcon->result_relation,
+									   rcon->nomatch_option,
+									   rcon->nomatch_varno);
+
+	/* Must adjust varlevelsup if replaced Var is within a subquery */
+	if (var->varlevelsup > 0)
+		IncrementVarSublevelsUp(newnode, var->varlevelsup, 0);
+
+	return newnode;
+}
+
+Node *
+ReplaceVarFromTargetList(Var *var,
+						 RangeTblEntry *target_rte,
+						 List *targetlist,
+						 int result_relation,
+						 ReplaceVarsNoMatchOption nomatch_option,
+						 int nomatch_varno)
+{
 	TargetEntry *tle;
 
 	if (var->varattno == InvalidAttrNumber)
@@ -1822,6 +1851,7 @@ ReplaceVarsFromTargetList_callback(Var *var,
 		RowExpr    *rowexpr;
 		List	   *colnames;
 		List	   *fields;
+		ListCell   *lc;
 
 		/*
 		 * If generating an expansion for a var of a named rowtype (ie, this
@@ -1830,29 +1860,46 @@ ReplaceVarsFromTargetList_callback(Var *var,
 		 * omit dropped columns.  In the latter case, attach column names to
 		 * the RowExpr for use of the executor and ruleutils.c.
 		 *
+		 * In order to be able to cache the results, we always generate the
+		 * expansion with varlevelsup = 0.  The caller is responsible for
+		 * adjusting it if needed.
+		 *
 		 * The varreturningtype is copied onto each individual field Var, so
 		 * that it is handled correctly when we recurse.
 		 */
-		expandRTE(rcon->target_rte,
-				  var->varno, var->varlevelsup, var->varreturningtype,
-				  var->location, (var->vartype != RECORDOID),
+		expandRTE(target_rte,
+				  var->varno, 0 /* not varlevelsup */ ,
+				  var->varreturningtype, var->location,
+				  (var->vartype != RECORDOID),
 				  &colnames, &fields);
-		/* Adjust the generated per-field Vars... */
-		fields = (List *) replace_rte_variables_mutator((Node *) fields,
-														context);
 		rowexpr = makeNode(RowExpr);
-		rowexpr->args = fields;
+		/* the fields will be set below */
+		rowexpr->args = NIL;
 		rowexpr->row_typeid = var->vartype;
 		rowexpr->row_format = COERCE_IMPLICIT_CAST;
 		rowexpr->colnames = (var->vartype == RECORDOID) ? colnames : NIL;
 		rowexpr->location = var->location;
+		/* Adjust the generated per-field Vars... */
+		foreach(lc, fields)
+		{
+			Node	   *field = lfirst(lc);
+
+			if (field && IsA(field, Var))
+				field = ReplaceVarFromTargetList((Var *) field,
+												 target_rte,
+												 targetlist,
+												 result_relation,
+												 nomatch_option,
+												 nomatch_varno);
+			rowexpr->args = lappend(rowexpr->args, field);
+		}
 
 		/* Wrap it in a ReturningExpr, if needed, per comments above */
 		if (var->varreturningtype != VAR_RETURNING_DEFAULT)
 		{
 			ReturningExpr *rexpr = makeNode(ReturningExpr);
 
-			rexpr->retlevelsup = var->varlevelsup;
+			rexpr->retlevelsup = 0;
 			rexpr->retold = (var->varreturningtype == VAR_RETURNING_OLD);
 			rexpr->retexpr = (Expr *) rowexpr;
 
@@ -1863,12 +1910,12 @@ ReplaceVarsFromTargetList_callback(Var *var,
 	}
 
 	/* Normal case referencing one targetlist element */
-	tle = get_tle_by_resno(rcon->targetlist, var->varattno);
+	tle = get_tle_by_resno(targetlist, var->varattno);
 
 	if (tle == NULL || tle->resjunk)
 	{
 		/* Failed to find column in targetlist */
-		switch (rcon->nomatch_option)
+		switch (nomatch_option)
 		{
 			case REPLACEVARS_REPORT_ERROR:
 				/* fall through, throw error below */
@@ -1876,7 +1923,8 @@ ReplaceVarsFromTargetList_callback(Var *var,
 
 			case REPLACEVARS_CHANGE_VARNO:
 				var = copyObject(var);
-				var->varno = rcon->nomatch_varno;
+				var->varno = nomatch_varno;
+				var->varlevelsup = 0;
 				/* we leave the syntactic referent alone */
 				return (Node *) var;
 
@@ -1906,10 +1954,6 @@ ReplaceVarsFromTargetList_callback(Var *var,
 		/* Make a copy of the tlist item to return */
 		Expr	   *newnode = copyObject(tle->expr);
 
-		/* Must adjust varlevelsup if tlist item is from higher query */
-		if (var->varlevelsup > 0)
-			IncrementVarSublevelsUp((Node *) newnode, var->varlevelsup, 0);
-
 		/*
 		 * Check to see if the tlist item contains a PARAM_MULTIEXPR Param,
 		 * and throw error if so.  This case could only happen when expanding
@@ -1932,20 +1976,20 @@ ReplaceVarsFromTargetList_callback(Var *var,
 			 * Copy varreturningtype onto any Vars in the tlist item that
 			 * refer to result_relation (which had better be non-zero).
 			 */
-			if (rcon->result_relation == 0)
+			if (result_relation == 0)
 				elog(ERROR, "variable returning old/new found outside RETURNING list");
 
-			SetVarReturningType((Node *) newnode, rcon->result_relation,
-								var->varlevelsup, var->varreturningtype);
+			SetVarReturningType((Node *) newnode, result_relation,
+								0, var->varreturningtype);
 
 			/* Wrap it in a ReturningExpr, if needed, per comments above */
 			if (!IsA(newnode, Var) ||
-				((Var *) newnode)->varno != rcon->result_relation ||
-				((Var *) newnode)->varlevelsup != var->varlevelsup)
+				((Var *) newnode)->varno != result_relation ||
+				((Var *) newnode)->varlevelsup != 0)
 			{
 				ReturningExpr *rexpr = makeNode(ReturningExpr);
 
-				rexpr->retlevelsup = var->varlevelsup;
+				rexpr->retlevelsup = 0;
 				rexpr->retold = (var->varreturningtype == VAR_RETURNING_OLD);
 				rexpr->retexpr = newnode;
 
diff --git a/src/include/rewrite/rewriteManip.h b/src/include/rewrite/rewriteManip.h
index 466edd7c1c2..ea3908739c6 100644
--- a/src/include/rewrite/rewriteManip.h
+++ b/src/include/rewrite/rewriteManip.h
@@ -55,9 +55,6 @@ extern void IncrementVarSublevelsUp(Node *node, int delta_sublevels_up,
 extern void IncrementVarSublevelsUp_rtable(List *rtable,
 										   int delta_sublevels_up, int min_sublevels_up);
 
-extern void SetVarReturningType(Node *node, int result_relation, int sublevels_up,
-								VarReturningType returning_type);
-
 extern bool rangeTableEntry_used(Node *node, int rt_index,
 								 int sublevels_up);
 
@@ -92,6 +89,12 @@ extern Node *map_variable_attnos(Node *node,
 								 const struct AttrMap *attno_map,
 								 Oid to_rowtype, bool *found_whole_row);
 
+extern Node *ReplaceVarFromTargetList(Var *var,
+									  RangeTblEntry *target_rte,
+									  List *targetlist,
+									  int result_relation,
+									  ReplaceVarsNoMatchOption nomatch_option,
+									  int nomatch_varno);
 extern Node *ReplaceVarsFromTargetList(Node *node,
 									   int target_varno, int sublevels_up,
 									   RangeTblEntry *target_rte,