Create an official API function for C functions to use to check if they are

being called as aggregates, and to get the aggregate transition state memory
context if needed.  Use it instead of poking directly into AggState and
WindowAggState in places that shouldn't know so much.

We should have done this in 8.4, probably, but better late than never.

Revised version of a patch by Hitoshi Harada.

5 files changed, 45 insertions, 74 deletions

diff --git a/src/backend/utils/adt/array_userfuncs.c b/src/backend/utils/adt/array_userfuncs.c
index 0bdb3d8a8a8..7d4ea11bf8b 100644
--- a/src/backend/utils/adt/array_userfuncs.c
+++ b/src/backend/utils/adt/array_userfuncs.c
@@ -6,13 +6,12 @@
  * Copyright (c) 2003-2010, PostgreSQL Global Development Group
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/utils/adt/array_userfuncs.c,v 1.33 2010/01/02 16:57:53 momjian Exp $
+ *	  $PostgreSQL: pgsql/src/backend/utils/adt/array_userfuncs.c,v 1.34 2010/02/08 20:39:51 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
 #include "postgres.h"
 
-#include "nodes/execnodes.h"
 #include "utils/array.h"
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
@@ -484,15 +483,10 @@ array_agg_transfn(PG_FUNCTION_ARGS)
 				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 				 errmsg("could not determine input data type")));
 
-	if (fcinfo->context && IsA(fcinfo->context, AggState))
-		aggcontext = ((AggState *) fcinfo->context)->aggcontext;
-	else if (fcinfo->context && IsA(fcinfo->context, WindowAggState))
-		aggcontext = ((WindowAggState *) fcinfo->context)->wincontext;
-	else
+	if (!AggCheckCallContext(fcinfo, &aggcontext))
 	{
 		/* cannot be called directly because of internal-type argument */
 		elog(ERROR, "array_agg_transfn called in non-aggregate context");
-		aggcontext = NULL;		/* keep compiler quiet */
 	}
 
 	state = PG_ARGISNULL(0) ? NULL : (ArrayBuildState *) PG_GETARG_POINTER(0);
@@ -528,9 +522,7 @@ array_agg_finalfn(PG_FUNCTION_ARGS)
 		PG_RETURN_NULL();		/* returns null iff no input values */
 
 	/* cannot be called directly because of internal-type argument */
-	Assert(fcinfo->context &&
-		   (IsA(fcinfo->context, AggState) ||
-			IsA(fcinfo->context, WindowAggState)));
+	Assert(AggCheckCallContext(fcinfo, NULL));
 
 	state = (ArrayBuildState *) PG_GETARG_POINTER(0);
 
diff --git a/src/backend/utils/adt/float.c b/src/backend/utils/adt/float.c
index 9f194f2785f..82e95704fb1 100644
--- a/src/backend/utils/adt/float.c
+++ b/src/backend/utils/adt/float.c
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/utils/adt/float.c,v 1.164 2010/01/02 16:57:53 momjian Exp $
+ *	  $PostgreSQL: pgsql/src/backend/utils/adt/float.c,v 1.165 2010/02/08 20:39:51 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -1765,13 +1765,11 @@ float8_accum(PG_FUNCTION_ARGS)
 	CHECKFLOATVAL(sumX2, isinf(transvalues[2]) || isinf(newval), true);
 
 	/*
-	 * If we're invoked by nodeAgg, we can cheat and modify our first
+	 * If we're invoked as an aggregate, we can cheat and modify our first
 	 * parameter in-place to reduce palloc overhead. Otherwise we construct a
 	 * new array with the updated transition data and return it.
 	 */
-	if (fcinfo->context &&
-		(IsA(fcinfo->context, AggState) ||
-		 IsA(fcinfo->context, WindowAggState)))
+	if (AggCheckCallContext(fcinfo, NULL))
 	{
 		transvalues[0] = N;
 		transvalues[1] = sumX;
@@ -1820,13 +1818,11 @@ float4_accum(PG_FUNCTION_ARGS)
 	CHECKFLOATVAL(sumX2, isinf(transvalues[2]) || isinf(newval), true);
 
 	/*
-	 * If we're invoked by nodeAgg, we can cheat and modify our first
+	 * If we're invoked as an aggregate, we can cheat and modify our first
 	 * parameter in-place to reduce palloc overhead. Otherwise we construct a
 	 * new array with the updated transition data and return it.
 	 */
-	if (fcinfo->context &&
-		(IsA(fcinfo->context, AggState) ||
-		 IsA(fcinfo->context, WindowAggState)))
+	if (AggCheckCallContext(fcinfo, NULL))
 	{
 		transvalues[0] = N;
 		transvalues[1] = sumX;
@@ -2039,13 +2035,11 @@ float8_regr_accum(PG_FUNCTION_ARGS)
 				  isinf(newvalY), true);
 
 	/*
-	 * If we're invoked by nodeAgg, we can cheat and modify our first
+	 * If we're invoked as an aggregate, we can cheat and modify our first
 	 * parameter in-place to reduce palloc overhead. Otherwise we construct a
 	 * new array with the updated transition data and return it.
 	 */
-	if (fcinfo->context &&
-		(IsA(fcinfo->context, AggState) ||
-		 IsA(fcinfo->context, WindowAggState)))
+	if (AggCheckCallContext(fcinfo, NULL))
 	{
 		transvalues[0] = N;
 		transvalues[1] = sumX;
diff --git a/src/backend/utils/adt/int8.c b/src/backend/utils/adt/int8.c
index 6707b79e548..1482017561d 100644
--- a/src/backend/utils/adt/int8.c
+++ b/src/backend/utils/adt/int8.c
@@ -7,7 +7,7 @@
  * Portions Copyright (c) 1994, Regents of the University of California
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/utils/adt/int8.c,v 1.77 2010/01/07 04:53:34 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/backend/utils/adt/int8.c,v 1.78 2010/02/08 20:39:51 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -19,7 +19,6 @@
 
 #include "funcapi.h"
 #include "libpq/pqformat.h"
-#include "nodes/nodes.h"
 #include "utils/int8.h"
 
 
@@ -654,15 +653,13 @@ int8inc(PG_FUNCTION_ARGS)
 {
 	/*
 	 * When int8 is pass-by-reference, we provide this special case to avoid
-	 * palloc overhead for COUNT(): when called from nodeAgg, we know that the
-	 * argument is modifiable local storage, so just update it in-place. (If
-	 * int8 is pass-by-value, then of course this is useless as well as
-	 * incorrect, so just ifdef it out.)
+	 * palloc overhead for COUNT(): when called as an aggregate, we know that
+	 * the argument is modifiable local storage, so just update it
+	 * in-place. (If int8 is pass-by-value, then of course this is useless as
+	 * well as incorrect, so just ifdef it out.)
 	 */
 #ifndef USE_FLOAT8_BYVAL		/* controls int8 too */
-	if (fcinfo->context &&
-		(IsA(fcinfo->context, AggState) ||
-		 IsA(fcinfo->context, WindowAggState)))
+	if (AggCheckCallContext(fcinfo, NULL))
 	{
 		int64	   *arg = (int64 *) PG_GETARG_POINTER(0);
 		int64		result;
@@ -680,7 +677,7 @@ int8inc(PG_FUNCTION_ARGS)
 	else
 #endif
 	{
-		/* Not called by nodeAgg, so just do it the dumb way */
+		/* Not called as an aggregate, so just do it the dumb way */
 		int64		arg = PG_GETARG_INT64(0);
 		int64		result;
 
diff --git a/src/backend/utils/adt/numeric.c b/src/backend/utils/adt/numeric.c
index 2892b5d2fbf..1b9f7944959 100644
--- a/src/backend/utils/adt/numeric.c
+++ b/src/backend/utils/adt/numeric.c
@@ -14,7 +14,7 @@
  * Copyright (c) 1998-2010, PostgreSQL Global Development Group
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/utils/adt/numeric.c,v 1.121 2010/01/07 04:53:34 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/backend/utils/adt/numeric.c,v 1.122 2010/02/08 20:39:51 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -2679,16 +2679,14 @@ int2_sum(PG_FUNCTION_ARGS)
 	}
 
 	/*
-	 * If we're invoked by nodeAgg, we can cheat and modify our first
+	 * If we're invoked as an aggregate, we can cheat and modify our first
 	 * parameter in-place to avoid palloc overhead. If not, we need to return
 	 * the new value of the transition variable. (If int8 is pass-by-value,
 	 * then of course this is useless as well as incorrect, so just ifdef it
 	 * out.)
 	 */
 #ifndef USE_FLOAT8_BYVAL		/* controls int8 too */
-	if (fcinfo->context &&
-		(IsA(fcinfo->context, AggState) ||
-		 IsA(fcinfo->context, WindowAggState)))
+	if (AggCheckCallContext(fcinfo, NULL))
 	{
 		int64	   *oldsum = (int64 *) PG_GETARG_POINTER(0);
 
@@ -2730,16 +2728,14 @@ int4_sum(PG_FUNCTION_ARGS)
 	}
 
 	/*
-	 * If we're invoked by nodeAgg, we can cheat and modify our first
+	 * If we're invoked as an aggregate, we can cheat and modify our first
 	 * parameter in-place to avoid palloc overhead. If not, we need to return
 	 * the new value of the transition variable. (If int8 is pass-by-value,
 	 * then of course this is useless as well as incorrect, so just ifdef it
 	 * out.)
 	 */
 #ifndef USE_FLOAT8_BYVAL		/* controls int8 too */
-	if (fcinfo->context &&
-		(IsA(fcinfo->context, AggState) ||
-		 IsA(fcinfo->context, WindowAggState)))
+	if (AggCheckCallContext(fcinfo, NULL))
 	{
 		int64	   *oldsum = (int64 *) PG_GETARG_POINTER(0);
 
@@ -2782,7 +2778,7 @@ int8_sum(PG_FUNCTION_ARGS)
 	}
 
 	/*
-	 * Note that we cannot special-case the nodeAgg case here, as we do for
+	 * Note that we cannot special-case the aggregate case here, as we do for
 	 * int2_sum and int4_sum: numeric is of variable size, so we cannot modify
 	 * our first parameter in-place.
 	 */
@@ -2820,13 +2816,11 @@ int2_avg_accum(PG_FUNCTION_ARGS)
 	Int8TransTypeData *transdata;
 
 	/*
-	 * If we're invoked by nodeAgg, we can cheat and modify our first
+	 * If we're invoked as an aggregate, we can cheat and modify our first
 	 * parameter in-place to reduce palloc overhead. Otherwise we need to make
 	 * a copy of it before scribbling on it.
 	 */
-	if (fcinfo->context &&
-		(IsA(fcinfo->context, AggState) ||
-		 IsA(fcinfo->context, WindowAggState)))
+	if (AggCheckCallContext(fcinfo, NULL))
 		transarray = PG_GETARG_ARRAYTYPE_P(0);
 	else
 		transarray = PG_GETARG_ARRAYTYPE_P_COPY(0);
@@ -2850,13 +2844,11 @@ int4_avg_accum(PG_FUNCTION_ARGS)
 	Int8TransTypeData *transdata;
 
 	/*
-	 * If we're invoked by nodeAgg, we can cheat and modify our first
+	 * If we're invoked as an aggregate, we can cheat and modify our first
 	 * parameter in-place to reduce palloc overhead. Otherwise we need to make
 	 * a copy of it before scribbling on it.
 	 */
-	if (fcinfo->context &&
-		(IsA(fcinfo->context, AggState) ||
-		 IsA(fcinfo->context, WindowAggState)))
+	if (AggCheckCallContext(fcinfo, NULL))
 		transarray = PG_GETARG_ARRAYTYPE_P(0);
 	else
 		transarray = PG_GETARG_ARRAYTYPE_P_COPY(0);
diff --git a/src/backend/utils/adt/varlena.c b/src/backend/utils/adt/varlena.c
index 096ca75bf93..652e1e6add4 100644
--- a/src/backend/utils/adt/varlena.c
+++ b/src/backend/utils/adt/varlena.c
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/utils/adt/varlena.c,v 1.175 2010/02/01 03:14:43 itagaki Exp $
+ *	  $PostgreSQL: pgsql/src/backend/utils/adt/varlena.c,v 1.176 2010/02/08 20:39:51 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -21,7 +21,6 @@
 #include "libpq/md5.h"
 #include "libpq/pqformat.h"
 #include "miscadmin.h"
-#include "nodes/execnodes.h"
 #include "parser/scansup.h"
 #include "regex/regex.h"
 #include "utils/builtins.h"
@@ -74,7 +73,7 @@ static bytea *bytea_substring(Datum str,
 				int L,
 				bool length_not_specified);
 static bytea *bytea_overlay(bytea *t1, bytea *t2, int sp, int sl);
-static StringInfo makeStringAggState(fmNodePtr context);
+static StringInfo makeStringAggState(FunctionCallInfo fcinfo);
 
 
 /*****************************************************************************
@@ -3327,25 +3326,25 @@ pg_column_size(PG_FUNCTION_ARGS)
  * actually used at all, and on subsequent calls the delimiter precedes
  * the associated value.
  */
+
+/* subroutine to initialize state */
 static StringInfo
-makeStringAggState(fmNodePtr context)
+makeStringAggState(FunctionCallInfo fcinfo)
 {
 	StringInfo		state;
 	MemoryContext	aggcontext;
 	MemoryContext	oldcontext;
 
-	if (context && IsA(context, AggState))
-		aggcontext = ((AggState *) context)->aggcontext;
-	else if (context && IsA(context, WindowAggState))
-		aggcontext = ((WindowAggState *) context)->wincontext;
-	else
+	if (!AggCheckCallContext(fcinfo, &aggcontext))
 	{
 		/* cannot be called directly because of internal-type argument */
 		elog(ERROR, "string_agg_transfn called in non-aggregate context");
-		aggcontext = NULL;		/* keep compiler quiet */
 	}
 
-	/* Create state in aggregate context */
+	/*
+	 * Create state in aggregate context.  It'll stay there across subsequent
+	 * calls.
+	 */
 	oldcontext = MemoryContextSwitchTo(aggcontext);
 	state = makeStringInfo();
 	MemoryContextSwitchTo(oldcontext);
@@ -3360,11 +3359,11 @@ string_agg_transfn(PG_FUNCTION_ARGS)
 
 	state = PG_ARGISNULL(0) ? NULL : (StringInfo) PG_GETARG_POINTER(0);
 
-	/* Append the element unless not null. */
+	/* Append the element unless null. */
 	if (!PG_ARGISNULL(1))
 	{
 		if (state == NULL)
-			state = makeStringAggState(fcinfo->context);
+			state = makeStringAggState(fcinfo);
 		appendStringInfoText(state, PG_GETARG_TEXT_PP(1));	/* value */
 	}
 
@@ -3382,11 +3381,12 @@ string_agg_delim_transfn(PG_FUNCTION_ARGS)
 
 	state = PG_ARGISNULL(0) ? NULL : (StringInfo) PG_GETARG_POINTER(0);
 
-	/* Append the value unless not null. */
+	/* Append the value unless null. */
 	if (!PG_ARGISNULL(1))
 	{
+		/* On the first time through, we ignore the delimiter. */
 		if (state == NULL)
-			state = makeStringAggState(fcinfo->context);
+			state = makeStringAggState(fcinfo);
 		else if (!PG_ARGISNULL(2))
 			appendStringInfoText(state, PG_GETARG_TEXT_PP(2));	/* delimiter */
 
@@ -3405,15 +3405,11 @@ string_agg_finalfn(PG_FUNCTION_ARGS)
 {
 	StringInfo		state;
 
-	if (PG_ARGISNULL(0))
-		PG_RETURN_NULL();
-
 	/* cannot be called directly because of internal-type argument */
-	Assert(fcinfo->context &&
-		   (IsA(fcinfo->context, AggState) ||
-			IsA(fcinfo->context, WindowAggState)));
+	Assert(AggCheckCallContext(fcinfo, NULL));
+
+	state = PG_ARGISNULL(0) ? NULL : (StringInfo) PG_GETARG_POINTER(0);
 
-	state = (StringInfo) PG_GETARG_POINTER(0);
 	if (state != NULL)
 		PG_RETURN_TEXT_P(cstring_to_text(state->data));
 	else