Allow extension functions to participate in in-place updates.

Commit 1dc5ebc90 allowed PL/pgSQL to perform in-place updates
of expanded-object variables that are being updated with
assignments like "x := f(x, ...)".  However this was allowed
only for a hard-wired list of functions f(), since we need to
be sure that f() will not modify the variable if it fails.
It was always envisioned that we should make that extensible,
but at the time we didn't have a good way to do so.  Since
then we've invented the idea of "support functions" to allow
attaching specialized optimization knowledge to functions,
and that is a perfect mechanism for doing this.

Hence, adjust PL/pgSQL to use a support function request instead
of hard-wired logic to decide if in-place update is safe.
Preserve the previous optimizations by creating support functions
for the three functions that were previously hard-wired.

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

3 files changed, 40 insertions, 50 deletions

diff --git a/src/pl/plpgsql/src/expected/plpgsql_array.out b/src/pl/plpgsql/src/expected/plpgsql_array.out
index e5db6d60876..4c6b3ce998a 100644
--- a/src/pl/plpgsql/src/expected/plpgsql_array.out
+++ b/src/pl/plpgsql/src/expected/plpgsql_array.out
@@ -57,10 +57,11 @@ begin
   -- test scenarios for optimization of updates of R/W expanded objects
   a := array_append(a, 42);  -- optimizable using "transfer" method
   a := a || a[3];  -- optimizable using "inplace" method
+  a := a[1] || a;  -- ditto, but let's test array_prepend
   a := a || a;     -- not optimizable
   raise notice 'a = %', a;
 end$$;
-NOTICE:  a = {1,2,3,42,3,1,2,3,42,3}
+NOTICE:  a = {1,1,2,3,42,3,1,1,2,3,42,3}
 create temp table onecol as select array[1,2] as f1;
 do $$ declare a int[];
 begin a := f1 from onecol; raise notice 'a = %', a; end$$;
diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index 28b6c85d8d2..d4377ceecbf 100644
--- a/src/pl/plpgsql/src/pl_exec.c
+++ b/src/pl/plpgsql/src/pl_exec.c
@@ -29,6 +29,7 @@
 #include "mb/stringinfo_mb.h"
 #include "miscadmin.h"
 #include "nodes/nodeFuncs.h"
+#include "nodes/supportnodes.h"
 #include "optimizer/optimizer.h"
 #include "parser/parse_coerce.h"
 #include "parser/parse_type.h"
@@ -8411,7 +8412,7 @@ exec_check_rw_parameter(PLpgSQL_expr *expr, int paramid)
 	Expr	   *sexpr = expr->expr_simple_expr;
 	Oid			funcid;
 	List	   *fargs;
-	ListCell   *lc;
+	Oid			prosupport;
 
 	/* Assume unsafe */
 	expr->expr_rwopt = PLPGSQL_RWOPT_NOPE;
@@ -8480,64 +8481,51 @@ exec_check_rw_parameter(PLpgSQL_expr *expr, int paramid)
 	{
 		SubscriptingRef *sbsref = (SubscriptingRef *) sexpr;
 
-		/* We only trust standard varlena arrays to be safe */
-		/* TODO: install some extensibility here */
-		if (get_typsubscript(sbsref->refcontainertype, NULL) !=
-			F_ARRAY_SUBSCRIPT_HANDLER)
-			return;
-
-		/* We can optimize the refexpr if it's the target, otherwise not */
-		if (sbsref->refexpr && IsA(sbsref->refexpr, Param))
-		{
-			Param	   *param = (Param *) sbsref->refexpr;
+		funcid = get_typsubscript(sbsref->refcontainertype, NULL);
 
-			if (param->paramkind == PARAM_EXTERN &&
-				param->paramid == paramid)
-			{
-				/* Found the Param we want to pass as read/write */
-				expr->expr_rwopt = PLPGSQL_RWOPT_INPLACE;
-				expr->expr_rw_param = param;
-				return;
-			}
-		}
-
-		return;
+		/*
+		 * We assume that only the refexpr and refassgnexpr (if any) are
+		 * relevant to the support function's decision.  If that turns out to
+		 * be a bad idea, we could incorporate the subscript expressions into
+		 * the fargs list somehow.
+		 */
+		fargs = list_make2(sbsref->refexpr, sbsref->refassgnexpr);
 	}
 	else
 		return;
 
 	/*
-	 * The top-level function must be one that we trust to be "safe".
-	 * Currently we hard-wire the list, but it would be very desirable to
-	 * allow extensions to mark their functions as safe ...
+	 * The top-level function must be one that can handle in-place update
+	 * safely.  We allow functions to declare their ability to do that via a
+	 * support function request.
 	 */
-	if (!(funcid == F_ARRAY_APPEND ||
-		  funcid == F_ARRAY_PREPEND))
-		return;
-
-	/*
-	 * The target variable (in the form of a Param) must appear as a direct
-	 * argument of the top-level function.  References further down in the
-	 * tree can't be optimized; but on the other hand, they don't invalidate
-	 * optimizing the top-level call, since that will be executed last.
-	 */
-	foreach(lc, fargs)
+	prosupport = get_func_support(funcid);
+	if (OidIsValid(prosupport))
 	{
-		Node	   *arg = (Node *) lfirst(lc);
+		SupportRequestModifyInPlace req;
+		Param	   *param;
 
-		if (arg && IsA(arg, Param))
-		{
-			Param	   *param = (Param *) arg;
+		req.type = T_SupportRequestModifyInPlace;
+		req.funcid = funcid;
+		req.args = fargs;
+		req.paramid = paramid;
 
-			if (param->paramkind == PARAM_EXTERN &&
-				param->paramid == paramid)
-			{
-				/* Found the Param we want to pass as read/write */
-				expr->expr_rwopt = PLPGSQL_RWOPT_INPLACE;
-				expr->expr_rw_param = param;
-				return;
-			}
-		}
+		param = (Param *)
+			DatumGetPointer(OidFunctionCall1(prosupport,
+											 PointerGetDatum(&req)));
+
+		if (param == NULL)
+			return;				/* support function fails */
+
+		/* Verify support function followed the API */
+		Assert(IsA(param, Param));
+		Assert(param->paramkind == PARAM_EXTERN);
+		Assert(param->paramid == paramid);
+
+		/* Found the Param we want to pass as read/write */
+		expr->expr_rwopt = PLPGSQL_RWOPT_INPLACE;
+		expr->expr_rw_param = param;
+		return;
 	}
 }
 
diff --git a/src/pl/plpgsql/src/sql/plpgsql_array.sql b/src/pl/plpgsql/src/sql/plpgsql_array.sql
index 4a346203dc2..da984a99414 100644
--- a/src/pl/plpgsql/src/sql/plpgsql_array.sql
+++ b/src/pl/plpgsql/src/sql/plpgsql_array.sql
@@ -53,6 +53,7 @@ begin
   -- test scenarios for optimization of updates of R/W expanded objects
   a := array_append(a, 42);  -- optimizable using "transfer" method
   a := a || a[3];  -- optimizable using "inplace" method
+  a := a[1] || a;  -- ditto, but let's test array_prepend
   a := a || a;     -- not optimizable
   raise notice 'a = %', a;
 end$$;