Add user-specified SRF state placement

Author: Nikita Glukhov

src/backend/utils/fmgr/funcapi.c

@@ -50,7 +50,7 @@ static TypeFuncClass get_type_func_class(Oid typid, Oid *base_typeid);
  * and error checking
  */
 FuncCallContext *
-init_MultiFuncCall(PG_FUNCTION_ARGS)
+init_MultiFuncCall(PG_FUNCTION_ARGS, FuncCallContext **pfuncctx)
 {
 	FuncCallContext *retval;
 
@@ -62,7 +62,7 @@ init_MultiFuncCall(PG_FUNCTION_ARGS)
 				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 				 errmsg("set-valued function called in context that cannot accept a set")));
 
-	if (fcinfo->flinfo->fn_extra == NULL)
+	if (*pfuncctx == NULL)
 	{
 		/*
 		 * First call
@@ -97,15 +97,15 @@ init_MultiFuncCall(PG_FUNCTION_ARGS)
 		/*
 		 * save the pointer for cross-call use
 		 */
-		fcinfo->flinfo->fn_extra = retval;
+		*pfuncctx = retval;
 
 		/*
 		 * Ensure we will get shut down cleanly if the exprcontext is not run
 		 * to completion.
 		 */
 		RegisterExprContextCallback(rsi->econtext,
 									shutdown_MultiFuncCall,
-									PointerGetDatum(fcinfo->flinfo));
+									PointerGetDatum(pfuncctx));
 	}
 	else
 	{
@@ -125,9 +125,9 @@ init_MultiFuncCall(PG_FUNCTION_ARGS)
  * Do Multi-function per-call setup
  */
 FuncCallContext *
-per_MultiFuncCall(PG_FUNCTION_ARGS)
+per_MultiFuncCall(PG_FUNCTION_ARGS, FuncCallContext **pfuncctx)
 {
-	FuncCallContext *retval = (FuncCallContext *) fcinfo->flinfo->fn_extra;
+	FuncCallContext *retval = *pfuncctx;
 
 	return retval;
 }
@@ -137,17 +137,18 @@ per_MultiFuncCall(PG_FUNCTION_ARGS)
  * Clean up after init_MultiFuncCall
  */
 void
-end_MultiFuncCall(PG_FUNCTION_ARGS, FuncCallContext *funcctx)
+end_MultiFuncCall(PG_FUNCTION_ARGS, FuncCallContext *funcctx,
+				  FuncCallContext **pfuncctx)
 {
 	ReturnSetInfo *rsi = (ReturnSetInfo *) fcinfo->resultinfo;
 
 	/* Deregister the shutdown callback */
 	UnregisterExprContextCallback(rsi->econtext,
 								  shutdown_MultiFuncCall,
-								  PointerGetDatum(fcinfo->flinfo));
+								  PointerGetDatum(pfuncctx));
 
 	/* But use it to do the real work */
-	shutdown_MultiFuncCall(PointerGetDatum(fcinfo->flinfo));
+	shutdown_MultiFuncCall(PointerGetDatum(pfuncctx));
 }
 
 /*
@@ -157,11 +158,11 @@ end_MultiFuncCall(PG_FUNCTION_ARGS, FuncCallContext *funcctx)
 static void
 shutdown_MultiFuncCall(Datum arg)
 {
-	FmgrInfo   *flinfo = (FmgrInfo *) DatumGetPointer(arg);
-	FuncCallContext *funcctx = (FuncCallContext *) flinfo->fn_extra;
+	FuncCallContext **pfuncctx = (FuncCallContext **) DatumGetPointer(arg);
+	FuncCallContext *funcctx = *pfuncctx;
 
 	/* unbind from flinfo */
-	flinfo->fn_extra = NULL;
+	*pfuncctx = NULL;
 
 	/*
 	 * Delete context that holds all multi-call data, including the

src/include/funcapi.h

@@ -276,15 +276,20 @@ extern Datum HeapTupleHeaderGetDatum(HeapTupleHeader tuple);
  */
 
 /* from funcapi.c */
-extern FuncCallContext *init_MultiFuncCall(PG_FUNCTION_ARGS);
-extern FuncCallContext *per_MultiFuncCall(PG_FUNCTION_ARGS);
-extern void end_MultiFuncCall(PG_FUNCTION_ARGS, FuncCallContext *funcctx);
+extern FuncCallContext *init_MultiFuncCall(PG_FUNCTION_ARGS, FuncCallContext **pfuncctx);
+extern FuncCallContext *per_MultiFuncCall(PG_FUNCTION_ARGS, FuncCallContext **pfuncctx);
+extern void end_MultiFuncCall(PG_FUNCTION_ARGS, FuncCallContext *funcctx, FuncCallContext **pfuncctx);
 
-#define SRF_IS_FIRSTCALL() (fcinfo->flinfo->fn_extra == NULL)
+#define SRF_DEFAULT_FCTX ((FuncCallContext **) &fcinfo->flinfo->fn_extra)
 
-#define SRF_FIRSTCALL_INIT() init_MultiFuncCall(fcinfo)
+#define SRF_IS_FIRSTCALL_EXT(_pfuncctx) (*(_pfuncctx) == NULL)
+#define SRF_IS_FIRSTCALL()		SRF_IS_FIRSTCALL_EXT(SRF_DEFAULT_FCTX)
 
-#define SRF_PERCALL_SETUP() per_MultiFuncCall(fcinfo)
+#define SRF_FIRSTCALL_INIT_EXT(_pfuncctx) init_MultiFuncCall(fcinfo, _pfuncctx)
+#define SRF_FIRSTCALL_INIT()	SRF_FIRSTCALL_INIT_EXT(SRF_DEFAULT_FCTX)
+
+#define SRF_PERCALL_SETUP_EXT(_pfuncctx) per_MultiFuncCall(fcinfo, _pfuncctx)
+#define SRF_PERCALL_SETUP()		SRF_PERCALL_SETUP_EXT(SRF_DEFAULT_FCTX)
 
 #define SRF_RETURN_NEXT(_funcctx, _result) \
 	do { \
@@ -304,15 +309,18 @@ extern void end_MultiFuncCall(PG_FUNCTION_ARGS, FuncCallContext *funcctx);
 		PG_RETURN_NULL(); \
 	} while (0)
 
-#define  SRF_RETURN_DONE(_funcctx) \
+#define  SRF_RETURN_DONE_EXT(_funcctx, _pfuncctx) \
 	do { \
 		ReturnSetInfo	   *rsi; \
-		end_MultiFuncCall(fcinfo, _funcctx); \
+		end_MultiFuncCall(fcinfo, _funcctx, _pfuncctx); \
 		rsi = (ReturnSetInfo *) fcinfo->resultinfo; \
 		rsi->isDone = ExprEndResult; \
 		PG_RETURN_NULL(); \
 	} while (0)
 
+#define  SRF_RETURN_DONE(_funcctx) \
+	SRF_RETURN_DONE_EXT(_funcctx, SRF_DEFAULT_FCTX)
+
 /*----------
  *	Support to ease writing of functions dealing with VARIADIC inputs
  *----------