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

src/backend/utils/adt/array_userfuncs.c

@@ -16,6 +16,7 @@
 #include "common/int.h"
 #include "common/pg_prng.h"
 #include "libpq/pqformat.h"
+#include "nodes/supportnodes.h"
 #include "port/pg_bitutils.h"
 #include "utils/array.h"
 #include "utils/builtins.h"
@@ -167,6 +168,36 @@ array_append(PG_FUNCTION_ARGS)
 	PG_RETURN_DATUM(result);
 }
 
+/*
+ * array_append_support()
+ *
+ * Planner support function for array_append()
+ */
+Datum
+array_append_support(PG_FUNCTION_ARGS)
+{
+	Node	   *rawreq = (Node *) PG_GETARG_POINTER(0);
+	Node	   *ret = NULL;
+
+	if (IsA(rawreq, SupportRequestModifyInPlace))
+	{
+		/*
+		 * We can optimize in-place appends if the function's array argument
+		 * is the array being assigned to.  We don't need to worry about array
+		 * references within the other argument.
+		 */
+		SupportRequestModifyInPlace *req = (SupportRequestModifyInPlace *) rawreq;
+		Param	   *arg = (Param *) linitial(req->args);
+
+		if (arg && IsA(arg, Param) &&
+			arg->paramkind == PARAM_EXTERN &&
+			arg->paramid == req->paramid)
+			ret = (Node *) arg;
+	}
+
+	PG_RETURN_POINTER(ret);
+}
+
 /*-----------------------------------------------------------------------------
  * array_prepend :
  *		push an element onto the front of a one-dimensional array
@@ -230,6 +261,36 @@ array_prepend(PG_FUNCTION_ARGS)
 	PG_RETURN_DATUM(result);
 }
 
+/*
+ * array_prepend_support()
+ *
+ * Planner support function for array_prepend()
+ */
+Datum
+array_prepend_support(PG_FUNCTION_ARGS)
+{
+	Node	   *rawreq = (Node *) PG_GETARG_POINTER(0);
+	Node	   *ret = NULL;
+
+	if (IsA(rawreq, SupportRequestModifyInPlace))
+	{
+		/*
+		 * We can optimize in-place prepends if the function's array argument
+		 * is the array being assigned to.  We don't need to worry about array
+		 * references within the other argument.
+		 */
+		SupportRequestModifyInPlace *req = (SupportRequestModifyInPlace *) rawreq;
+		Param	   *arg = (Param *) lsecond(req->args);
+
+		if (arg && IsA(arg, Param) &&
+			arg->paramkind == PARAM_EXTERN &&
+			arg->paramid == req->paramid)
+			ret = (Node *) arg;
+	}
+
+	PG_RETURN_POINTER(ret);
+}
+
 /*-----------------------------------------------------------------------------
  * array_cat :
  *		concatenate two nD arrays to form an nD array, or

src/backend/utils/adt/arraysubs.c

@@ -18,6 +18,7 @@
 #include "nodes/makefuncs.h"
 #include "nodes/nodeFuncs.h"
 #include "nodes/subscripting.h"
+#include "nodes/supportnodes.h"
 #include "parser/parse_coerce.h"
 #include "parser/parse_expr.h"
 #include "utils/array.h"
@@ -575,3 +576,36 @@ raw_array_subscript_handler(PG_FUNCTION_ARGS)
 
 	PG_RETURN_POINTER(&sbsroutines);
 }
+
+/*
+ * array_subscript_handler_support()
+ *
+ * Planner support function for array_subscript_handler()
+ */
+Datum
+array_subscript_handler_support(PG_FUNCTION_ARGS)
+{
+	Node	   *rawreq = (Node *) PG_GETARG_POINTER(0);
+	Node	   *ret = NULL;
+
+	if (IsA(rawreq, SupportRequestModifyInPlace))
+	{
+		/*
+		 * We can optimize in-place subscripted assignment if the refexpr is
+		 * the array being assigned to.  We don't need to worry about array
+		 * references within the refassgnexpr or the subscripts; however, if
+		 * there's no refassgnexpr then it's a fetch which there's no need to
+		 * optimize.
+		 */
+		SupportRequestModifyInPlace *req = (SupportRequestModifyInPlace *) rawreq;
+		Param	   *refexpr = (Param *) linitial(req->args);
+
+		if (refexpr && IsA(refexpr, Param) &&
+			refexpr->paramkind == PARAM_EXTERN &&
+			refexpr->paramid == req->paramid &&
+			lsecond(req->args) != NULL)
+			ret = (Node *) refexpr;
+	}
+
+	PG_RETURN_POINTER(ret);
+}

src/include/catalog/catversion.h

@@ -57,6 +57,6 @@
  */
 
 /*							yyyymmddN */
-#define CATALOG_VERSION_NO	202502071
+#define CATALOG_VERSION_NO	202502111
 
 #endif

src/include/catalog/pg_proc.dat

@@ -1598,14 +1598,20 @@
   proname => 'cardinality', prorettype => 'int4', proargtypes => 'anyarray',
   prosrc => 'array_cardinality' },
 { oid => '378', descr => 'append element onto end of array',
-  proname => 'array_append', proisstrict => 'f',
-  prorettype => 'anycompatiblearray',
+  proname => 'array_append', prosupport => 'array_append_support',
+  proisstrict => 'f', prorettype => 'anycompatiblearray',
   proargtypes => 'anycompatiblearray anycompatible', prosrc => 'array_append' },
+{ oid => '8680', descr => 'planner support for array_append',
+  proname => 'array_append_support', prorettype => 'internal',
+  proargtypes => 'internal', prosrc => 'array_append_support' },
 { oid => '379', descr => 'prepend element onto front of array',
-  proname => 'array_prepend', proisstrict => 'f',
-  prorettype => 'anycompatiblearray',
+  proname => 'array_prepend', prosupport => 'array_prepend_support',
+  proisstrict => 'f', prorettype => 'anycompatiblearray',
   proargtypes => 'anycompatible anycompatiblearray',
   prosrc => 'array_prepend' },
+{ oid => '8681', descr => 'planner support for array_prepend',
+  proname => 'array_prepend_support', prorettype => 'internal',
+  proargtypes => 'internal', prosrc => 'array_prepend_support' },
 { oid => '383',
   proname => 'array_cat', proisstrict => 'f',
   prorettype => 'anycompatiblearray',
@@ -12207,8 +12213,12 @@
 
 # subscripting support for built-in types
 { oid => '6179', descr => 'standard array subscripting support',
-  proname => 'array_subscript_handler', prorettype => 'internal',
+  proname => 'array_subscript_handler',
+  prosupport => 'array_subscript_handler_support', prorettype => 'internal',
   proargtypes => 'internal', prosrc => 'array_subscript_handler' },
+{ oid => '8682', descr => 'planner support for array_subscript_handler',
+  proname => 'array_subscript_handler_support', prorettype => 'internal',
+  proargtypes => 'internal', prosrc => 'array_subscript_handler_support' },
 { oid => '6180', descr => 'raw array subscripting support',
   proname => 'raw_array_subscript_handler', prorettype => 'internal',
   proargtypes => 'internal', prosrc => 'raw_array_subscript_handler' },

src/include/nodes/supportnodes.h

@@ -6,10 +6,10 @@
  * This file defines the API for "planner support functions", which
  * are SQL functions (normally written in C) that can be attached to
  * another "target" function to give the system additional knowledge
- * about the target function.  All the current capabilities have to do
- * with planning queries that use the target function, though it is
- * possible that future extensions will add functionality to be invoked
- * by the parser or executor.
+ * about the target function.  The name is now something of a misnomer,
+ * since some of the call sites are in the executor not the planner,
+ * but "function support function" would be a confusing name so we
+ * stick with "planner support function".
  *
  * A support function must have the SQL signature
  *		supportfn(internal) returns internal
@@ -343,4 +343,51 @@ typedef struct SupportRequestOptimizeWindowClause
 								 * optimizations are possible. */
 } SupportRequestOptimizeWindowClause;
 
+/*
+ * The ModifyInPlace request allows the support function to detect whether
+ * a call to its target function can be allowed to modify a read/write
+ * expanded object in-place.  The context is that we are considering a
+ * PL/pgSQL (or similar PL) assignment of the form "x := f(x, ...)" where
+ * the variable x is of a type that can be represented as an expanded object
+ * (see utils/expandeddatum.h).  If f() can usefully optimize by modifying
+ * the passed-in object in-place, then this request can be implemented to
+ * instruct PL/pgSQL to pass a read-write expanded pointer to the variable's
+ * value.  (Note that there is no guarantee that later calls to f() will
+ * actually do so.  If f() receives a read-only pointer, or a pointer to a
+ * non-expanded object, it must follow the usual convention of not modifying
+ * the pointed-to object.)  There are two requirements that must be met
+ * to make this safe:
+ * 1. f() must guarantee that it will not have modified the object if it
+ * fails.  Otherwise the variable's value might change unexpectedly.
+ * 2. If the other arguments to f() ("..." in the above example) contain
+ * references to x, f() must be able to cope with that; or if that's not
+ * safe, the support function must scan the other arguments to verify that
+ * there are no other references to x.  An example of the concern here is
+ * that in "arr := array_append(arr, arr[1])", if the array element type
+ * is pass-by-reference then array_append would receive a second argument
+ * that points into the array object it intends to modify.  array_append is
+ * coded to make that safe, but other functions might not be able to cope.
+ *
+ * "args" is a node tree list representing the function's arguments.
+ * One or more nodes within the node tree will be PARAM_EXTERN Params
+ * with ID "paramid", which represent the assignment target variable.
+ * (Note that such references are not necessarily at top level in the list,
+ * for example we might have "x := f(x, g(x))".  Generally it's only safe
+ * to optimize a reference that is at top level, else we're making promises
+ * about the behavior of g() as well as f().)
+ *
+ * If modify-in-place is safe, the support function should return the
+ * address of the Param node that is to return a read-write pointer.
+ * (At most one of the references is allowed to do so.)  Otherwise,
+ * return NULL.
+ */
+typedef struct SupportRequestModifyInPlace
+{
+	NodeTag		type;
+
+	Oid			funcid;			/* PG_PROC OID of the target function */
+	List	   *args;			/* Arguments to the function */
+	int			paramid;		/* ID of Param(s) representing variable */
+} SupportRequestModifyInPlace;
+
 #endif							/* SUPPORTNODES_H */

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$$;

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,65 +8481,52 @@ 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;
+		funcid = get_typsubscript(sbsref->refcontainertype, NULL);
 
-		/* We can optimize the refexpr if it's the target, otherwise not */
-		if (sbsref->refexpr && IsA(sbsref->refexpr, Param))
-		{
-			Param	   *param = (Param *) sbsref->refexpr;
-
-			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;
+
+		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);
 
-			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;
 	}
-		}
-	}
 }
 
 /*

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$$;

src/tools/pgindent/typedefs.list

@@ -2804,6 +2804,7 @@ SubscriptionRelState
 SummarizerReadLocalXLogPrivate
 SupportRequestCost
 SupportRequestIndexCondition
+SupportRequestModifyInPlace
 SupportRequestOptimizeWindowClause
 SupportRequestRows
 SupportRequestSelectivity