Detect whether plpgsql assignment targets are "local" variables.

Mark whether the target of a potentially optimizable assignment
is "local", in the sense of being declared inside any exception
block that could trap an error thrown from the assignment.
(This implies that we needn't preserve the variable's value
in case of an error.  This patch doesn't do anything with the
knowledge, but the next one will.)

Normally, this requires a post-parsing scan of the function's
parse tree, since we don't know while parsing a BEGIN ...
construct whether we will find EXCEPTION at its end.  However,
if there are no BEGIN ... EXCEPTION blocks in the function at
all, then all assignments are local, even those to variables
representing function arguments.  We optimize that common case
by initializing the target_is_local flags to "true", and fixing
them up with a post-scan only if we found EXCEPTION.

Note that variables' default-value expressions are never interesting
for expanded-variable optimization, since they couldn't contain a
reference to the target variable anyway.  But the code is set up
to compute their target_param and target_is_local correctly anyway,
for consistency and in case someone thinks of a use for that data.

I added a bit of plpgsql_dumptree support to help verify that this
code sets the flags as expected.  I also added a plpgsql_dumptree
call in plpgsql_compile_inline.  It was at best an oversight that
"#option dump" didn't work in a DO block; now it does.

Author: Tom Lane <tgl@sss.pgh.pa.us>
Reviewed-by: Andrey Borodin <x4mmm@yandex-team.ru>
Reviewed-by: Pavel Borisov <pashkin.elfe@gmail.com>
Discussion: https://postgr.es/m/CACxu=vJaKFNsYxooSnW1wEgsAO5u_v1XYBacfVJ14wgJV_PYeg@mail.gmail.com

Author: tglsfdc

src/pl/plpgsql/src/pl_comp.c

@@ -371,6 +371,7 @@ do_compile(FunctionCallInfo fcinfo,
 
 	function->nstatements = 0;
 	function->requires_procedure_resowner = false;
+	function->has_exception_block = false;
 
 	/*
 	 * Initialize the compiler, particularly the namespace stack.  The
@@ -811,6 +812,9 @@ do_compile(FunctionCallInfo fcinfo,
 
 	plpgsql_finish_datums(function);
 
+	if (function->has_exception_block)
+		plpgsql_mark_local_assignment_targets(function);
+
 	/* Debug dump for completed functions */
 	if (plpgsql_DumpExecTree)
 		plpgsql_dumptree(function);
@@ -906,6 +910,7 @@ plpgsql_compile_inline(char *proc_source)
 
 	function->nstatements = 0;
 	function->requires_procedure_resowner = false;
+	function->has_exception_block = false;
 
 	plpgsql_ns_init();
 	plpgsql_ns_push(func_name, PLPGSQL_LABEL_BLOCK);
@@ -962,6 +967,13 @@ plpgsql_compile_inline(char *proc_source)
 
 	plpgsql_finish_datums(function);
 
+	if (function->has_exception_block)
+		plpgsql_mark_local_assignment_targets(function);
+
+	/* Debug dump for completed functions */
+	if (plpgsql_DumpExecTree)
+		plpgsql_dumptree(function);
+
 	/*
 	 * Pop the error context stack
 	 */

src/pl/plpgsql/src/pl_funcs.c

@@ -598,6 +598,91 @@ plpgsql_statement_tree_walker_impl(PLpgSQL_stmt *stmt,
 }
 
 
+/**********************************************************************
+ * Mark assignment source expressions that have local target variables,
+ * that is, the target variable is declared within the exception block
+ * most closely containing the assignment itself.  (Such target variables
+ * need not be preserved if the assignment's source expression raises an
+ * error, since the variable will no longer be accessible afterwards.
+ * Detecting this allows better optimization.)
+ *
+ * This code need not be called if the plpgsql function contains no exception
+ * blocks, because mark_expr_as_assignment_source will have set all the flags
+ * to true already.  Also, we need not reconsider default-value expressions
+ * for variables, because variable declarations are necessarily within the
+ * nearest exception block.  (In DECLARE ... BEGIN ... EXCEPTION ... END, the
+ * variable initializations are done before entering the exception scope.)
+ *
+ * Within the recursion, local_dnos is a Bitmapset of dnos of variables
+ * known to be declared within the current exception level.
+ **********************************************************************/
+static void mark_stmt(PLpgSQL_stmt *stmt, Bitmapset *local_dnos);
+static void mark_expr(PLpgSQL_expr *expr, Bitmapset *local_dnos);
+
+static void
+mark_stmt(PLpgSQL_stmt *stmt, Bitmapset *local_dnos)
+{
+	if (stmt == NULL)
+		return;
+	if (stmt->cmd_type == PLPGSQL_STMT_BLOCK)
+	{
+		PLpgSQL_stmt_block *block = (PLpgSQL_stmt_block *) stmt;
+
+		if (block->exceptions)
+		{
+			/*
+			 * The block creates a new exception scope, so variables declared
+			 * at outer levels are nonlocal.  For that matter, so are any
+			 * variables declared in the block's DECLARE section.  Hence, we
+			 * must pass down empty local_dnos.
+			 */
+			plpgsql_statement_tree_walker(stmt, mark_stmt, mark_expr, NULL);
+		}
+		else
+		{
+			/*
+			 * Otherwise, the block does not create a new exception scope, and
+			 * any variables it declares can also be considered local within
+			 * it.  Note that only initializable datum types (VAR, REC) are
+			 * included in initvarnos; but that's sufficient for our purposes.
+			 */
+			local_dnos = bms_copy(local_dnos);
+			for (int i = 0; i < block->n_initvars; i++)
+				local_dnos = bms_add_member(local_dnos, block->initvarnos[i]);
+			plpgsql_statement_tree_walker(stmt, mark_stmt, mark_expr,
+										  local_dnos);
+			bms_free(local_dnos);
+		}
+	}
+	else
+		plpgsql_statement_tree_walker(stmt, mark_stmt, mark_expr, local_dnos);
+}
+
+static void
+mark_expr(PLpgSQL_expr *expr, Bitmapset *local_dnos)
+{
+	/*
+	 * If this expression has an assignment target, check whether the target
+	 * is local, and mark the expression accordingly.
+	 */
+	if (expr && expr->target_param >= 0)
+		expr->target_is_local = bms_is_member(expr->target_param, local_dnos);
+}
+
+void
+plpgsql_mark_local_assignment_targets(PLpgSQL_function *func)
+{
+	Bitmapset  *local_dnos;
+
+	/* Function parameters can be treated as local targets at outer level */
+	local_dnos = NULL;
+	for (int i = 0; i < func->fn_nargs; i++)
+		local_dnos = bms_add_member(local_dnos, func->fn_argvarnos[i]);
+	mark_stmt((PLpgSQL_stmt *) func->action, local_dnos);
+	bms_free(local_dnos);
+}
+
+
 /**********************************************************************
  * Release memory when a PL/pgSQL function is no longer needed
  *
@@ -1500,6 +1585,9 @@ static void
 dump_expr(PLpgSQL_expr *expr)
 {
 	printf("'%s'", expr->query);
+	if (expr->target_param >= 0)
+		printf(" target %d%s", expr->target_param,
+			   expr->target_is_local ? " (local)" : "");
 }
 
 void

src/pl/plpgsql/src/pl_gram.y

@@ -2328,6 +2328,8 @@ exception_sect	:
 						PLpgSQL_exception_block *new = palloc(sizeof(PLpgSQL_exception_block));
 						PLpgSQL_variable *var;
 
+						plpgsql_curr_compile->has_exception_block = true;
+
 						var = plpgsql_build_variable("sqlstate", lineno,
 													 plpgsql_build_datatype(TEXTOID,
 																			-1,
@@ -2673,6 +2675,7 @@ make_plpgsql_expr(const char *query,
 	expr->ns = plpgsql_ns_top();
 	/* might get changed later during parsing: */
 	expr->target_param = -1;
+	expr->target_is_local = false;
 	/* other fields are left as zeroes until first execution */
 	return expr;
 }
@@ -2687,9 +2690,21 @@ mark_expr_as_assignment_source(PLpgSQL_expr *expr, PLpgSQL_datum *target)
 	 * other DTYPEs, so no need to mark in other cases.
 	 */
 	if (target->dtype == PLPGSQL_DTYPE_VAR)
+	{
 		expr->target_param = target->dno;
+
+		/*
+		 * For now, assume the target is local to the nearest enclosing
+		 * exception block.  That's correct if the function contains no
+		 * exception blocks; otherwise we'll update this later.
+		 */
+		expr->target_is_local = true;
+	}
 	else
+	{
 		expr->target_param = -1;	/* should be that already */
+		expr->target_is_local = false; /* ditto */
+	}
 }
 
 /* Convenience routine to read an expression with one possible terminator */

src/pl/plpgsql/src/plpgsql.h

@@ -225,9 +225,12 @@ typedef struct PLpgSQL_expr
 	/*
 	 * These fields are used to help optimize assignments to expanded-datum
 	 * variables.  If this expression is the source of an assignment to a
-	 * simple variable, target_param holds that variable's dno (else it's -1).
+	 * simple variable, target_param holds that variable's dno (else it's -1),
+	 * and target_is_local indicates whether the target is declared inside the
+	 * closest exception block containing the assignment.
 	 */
 	int			target_param;	/* dno of assign target, or -1 if none */
+	bool		target_is_local;	/* is it within nearest exception block? */
 
 	/*
 	 * Fields above are set during plpgsql parsing.  Remaining fields are left
@@ -1014,6 +1017,7 @@ typedef struct PLpgSQL_function
 	/* data derived while parsing body */
 	unsigned int nstatements;	/* counter for assigning stmtids */
 	bool		requires_procedure_resowner;	/* contains CALL or DO? */
+	bool		has_exception_block;	/* contains BEGIN...EXCEPTION? */
 
 	/* these fields change when the function is used */
 	struct PLpgSQL_execstate *cur_estate;
@@ -1312,6 +1316,7 @@ extern PLpgSQL_nsitem *plpgsql_ns_find_nearest_loop(PLpgSQL_nsitem *ns_cur);
  */
 extern PGDLLEXPORT const char *plpgsql_stmt_typename(PLpgSQL_stmt *stmt);
 extern const char *plpgsql_getdiag_kindname(PLpgSQL_getdiag_kind kind);
+extern void plpgsql_mark_local_assignment_targets(PLpgSQL_function *func);
 extern void plpgsql_free_function_memory(PLpgSQL_function *func);
 extern void plpgsql_dumptree(PLpgSQL_function *func);