/*
 * This file and its contents are licensed under the Timescale License.
 * Please see the included NOTICE for copyright information and
 * LICENSE-TIMESCALE for a copy of the license.
 */

/*
 * Vectorized implementation of a single transition function that uses the
 * Int128AggState transition state. Note that the serialization format differs
 * based on whether the sum of squares is needed.
 */
#ifdef GENERATE_DISPATCH_TABLE
extern VectorAggFunctions FUNCTION_NAME(argdef);
AGG_CASES
return &FUNCTION_NAME(argdef);
#else

typedef struct
{
	int64 N;
	int128 sumX;
#ifdef NEED_SUMX2
	int128 sumX2;
#endif
} FUNCTION_NAME(state);

static void
FUNCTION_NAME(init)(void *restrict agg_states, int n)
{
	FUNCTION_NAME(state) *states = (FUNCTION_NAME(state) *) agg_states;
	for (int i = 0; i < n; i++)
	{
		states[i].N = 0;
		states[i].sumX = 0;
#ifdef NEED_SUMX2
		states[i].sumX2 = 0;
#endif
	}
}

static void
FUNCTION_NAME(emit)(void *agg_state, Datum *out_result, bool *out_isnull)
{
	FUNCTION_NAME(state) *state = (FUNCTION_NAME(state) *) agg_state;

	PgInt128AggState result = {
		.N = state->N,
		.sumX = state->sumX,
#ifdef NEED_SUMX2
		.sumX2 = state->sumX2,
#endif
	};

	/*
	 * The serialization functions insist on being called in aggregate context,
	 * but thankfully don't use it in any way so we can use this dummy.
	 */
	AggState agg_context = { .ss.ps.type = T_AggState };
	LOCAL_FCINFO(fcinfo, 1);
	InitFunctionCallInfoData(*fcinfo, NULL, 1, InvalidOid, (Node *) &agg_context, NULL);

	fcinfo->args[0].value = PointerGetDatum(&result);
	fcinfo->args[0].isnull = false;

#ifdef NEED_SUMX2
	*out_result = numeric_poly_serialize(fcinfo);
#else
	*out_result = int8_avg_serialize(fcinfo);
#endif

	*out_isnull = false;
}

static pg_attribute_always_inline void
FUNCTION_NAME(vector_impl)(void *agg_state, int n, const CTYPE *values, const uint64 *filter,
						   MemoryContext agg_extra_mctx)
{
	int64 N = 0;
	int128 sumX = 0;
#ifdef NEED_SUMX2
	int128 sumX2 = 0;
#endif
	for (int row = 0; row < n; row++)
	{
		const bool row_ok = arrow_row_is_valid(filter, row);
		const CTYPE value = values[row];
		N += row_ok;
		sumX += value * row_ok;
#ifdef NEED_SUMX2
		sumX2 += ((int128) value) * ((int128) value) * row_ok;
#endif
	}

	FUNCTION_NAME(state) *state = (FUNCTION_NAME(state) *) agg_state;
	state->N += N;
	state->sumX += sumX;
#ifdef NEED_SUMX2
	state->sumX2 += sumX2;
#endif
}

static pg_attribute_always_inline void
FUNCTION_NAME(one)(void *restrict agg_state, const CTYPE value)
{
	FUNCTION_NAME(state) *state = (FUNCTION_NAME(state) *) agg_state;
	state->N++;
	state->sumX += value;
#ifdef NEED_SUMX2
	state->sumX2 += ((int128) value) * ((int128) value);
#endif
}

#include "agg_many_vector_helper.c"
#include "agg_scalar_helper.c"
#include "agg_vector_validity_helper.c"

VectorAggFunctions FUNCTION_NAME(argdef) = {
	.state_bytes = sizeof(FUNCTION_NAME(state)),
	.agg_init = FUNCTION_NAME(init),
	.agg_emit = FUNCTION_NAME(emit),
	.agg_scalar = FUNCTION_NAME(scalar),
	.agg_vector = FUNCTION_NAME(vector),
	.agg_many_vector = FUNCTION_NAME(many_vector),
};

#endif

#undef PG_TYPE
#undef CTYPE
#undef DATUM_TO_CTYPE
#undef AGG_CASES
#undef AGG_NAME
#undef NEED_SUMX2
