1 files changed, 113 insertions, 0 deletions
diff --git a/src/backend/rewrite/rewriteManip.c b/src/backend/rewrite/rewriteManip.c
index ec6b5bf5a92..9c778efd1c3 100644
--- a/src/backend/rewrite/rewriteManip.c
+++ b/src/backend/rewrite/rewriteManip.c
@@ -1218,6 +1218,119 @@ replace_rte_variables_mutator(Node *node,
 
 
 /*
+ * map_variable_attnos() finds all user-column Vars in an expression tree
+ * that reference a particular RTE, and adjusts their varattnos according
+ * to the given mapping array (varattno n is replaced by attno_map[n-1]).
+ * Vars for system columns are not modified.
+ *
+ * A zero in the mapping array represents a dropped column, which should not
+ * appear in the expression.
+ *
+ * If the expression tree contains a whole-row Var for the target RTE,
+ * the Var is not changed but *found_whole_row is returned as TRUE.
+ * For most callers this is an error condition, but we leave it to the caller
+ * to report the error so that useful context can be provided.  (In some
+ * usages it would be appropriate to modify the Var's vartype and insert a
+ * ConvertRowtypeExpr node to map back to the original vartype.  We might
+ * someday extend this function's API to support that.  For now, the only
+ * concession to that future need is that this function is a tree mutator
+ * not just a walker.)
+ *
+ * This could be built using replace_rte_variables and a callback function,
+ * but since we don't ever need to insert sublinks, replace_rte_variables is
+ * overly complicated.
+ */
+
+typedef struct
+{
+	int			target_varno;		/* RTE index to search for */
+	int			sublevels_up;		/* (current) nesting depth */
+	const AttrNumber *attno_map;	/* map array for user attnos */
+	int			map_length;			/* number of entries in attno_map[] */
+	bool	   *found_whole_row;	/* output flag */
+} map_variable_attnos_context;
+
+static Node *
+map_variable_attnos_mutator(Node *node,
+							map_variable_attnos_context *context)
+{
+	if (node == NULL)
+		return NULL;
+	if (IsA(node, Var))
+	{
+		Var		   *var = (Var *) node;
+
+		if (var->varno == context->target_varno &&
+			var->varlevelsup == context->sublevels_up)
+		{
+			/* Found a matching variable, make the substitution */
+			Var	   *newvar = (Var *) palloc(sizeof(Var));
+			int		attno = var->varattno;
+
+			*newvar = *var;
+			if (attno > 0)
+			{
+				/* user-defined column, replace attno */
+				if (attno > context->map_length ||
+					context->attno_map[attno - 1] == 0)
+					elog(ERROR, "unexpected varattno %d in expression to be mapped",
+						 attno);
+				newvar->varattno = newvar->varoattno = context->attno_map[attno - 1];
+			}
+			else if (attno == 0)
+			{
+				/* whole-row variable, warn caller */
+				*(context->found_whole_row) = true;
+			}
+			return (Node *) newvar;
+		}
+		/* otherwise fall through to copy the var normally */
+	}
+	else if (IsA(node, Query))
+	{
+		/* Recurse into RTE subquery or not-yet-planned sublink subquery */
+		Query	   *newnode;
+
+		context->sublevels_up++;
+		newnode = query_tree_mutator((Query *) node,
+									 map_variable_attnos_mutator,
+									 (void *) context,
+									 0);
+		context->sublevels_up--;
+		return (Node *) newnode;
+	}
+	return expression_tree_mutator(node, map_variable_attnos_mutator,
+								   (void *) context);
+}
+
+Node *
+map_variable_attnos(Node *node,
+					int target_varno, int sublevels_up,
+					const AttrNumber *attno_map, int map_length,
+					bool *found_whole_row)
+{
+	map_variable_attnos_context context;
+
+	context.target_varno = target_varno;
+	context.sublevels_up = sublevels_up;
+	context.attno_map = attno_map;
+	context.map_length = map_length;
+	context.found_whole_row = found_whole_row;
+
+	*found_whole_row = false;
+
+	/*
+	 * Must be prepared to start with a Query or a bare expression tree; if
+	 * it's a Query, we don't want to increment sublevels_up.
+	 */
+	return query_or_expression_tree_mutator(node,
+											map_variable_attnos_mutator,
+											(void *) &context,
+											0);
+}
+
+
+/*
  * ResolveNew - replace Vars with corresponding items from a targetlist
  *
  * Vars matching target_varno and sublevels_up are replaced by the