/*

create function drive_bloom(num_oids int, dup_freq int, count int) returns void
strict volatile language c as '/path/to/drive_bloom.so';

\timing on

select drive_bloom(100, 0, 100000);

 */

#include "postgres.h"

#include "fmgr.h"
#include "lib/bloomfilter.h"
#include "miscadmin.h"
#include "tcop/tcopprot.h"
#include "utils/builtins.h"
#include "utils/memutils.h"

PG_MODULE_MAGIC;

#define ROLES_LIST_BLOOM_THRESHOLD 1024
#define ROLES_LIST_BLOOM_SIZE (1024 * 1024)


static inline List *
roles_list_append(List *roles_list, Oid role, bloom_filter **roles_list_bf)
{
	unsigned char *roleptr = (unsigned char *) &role;

	/*
	 * If there is a previously-created Bloom filter, use it to determine
	 * whether the role is missing from the list.  Otherwise, do an ordinary
	 * linear search through the existing role list.
	 */
	if ((*roles_list_bf &&
		 bloom_lacks_element(*roles_list_bf, roleptr, sizeof(Oid))) ||
		!list_member_oid(roles_list, role))
	{
		/*
		 * If the list is large, we take on the overhead of creating and
		 * populating a Bloom filter to speed up future calls to this
		 * function.
		 */
		if (!*roles_list_bf &&
			list_length(roles_list) > ROLES_LIST_BLOOM_THRESHOLD)
		{
			*roles_list_bf = bloom_create(ROLES_LIST_BLOOM_SIZE,
										 work_mem, 0);
			foreach_oid(roleid, roles_list)
				bloom_add_element(*roles_list_bf,
								  (unsigned char *) &roleid,
								  sizeof(Oid));
		}

		/*
		 * Finally, add the role to the list and the Bloom filter, if it
		 * exists.
		 */
		roles_list = lappend_oid(roles_list, role);
		if (*roles_list_bf)
			bloom_add_element(*roles_list_bf, roleptr, sizeof(Oid));
	}

	return roles_list;
}

/*
 * drive_bloom(num_oids int, dup_freq int, count int) returns void
 *
 * num_oids: number of OIDs to de-duplicate
 * dup_freq: if > 0, every dup_freq'th OID is duplicated
 * count: overall repetition count; choose large enough to get reliable timing
 */
PG_FUNCTION_INFO_V1(drive_bloom);
Datum
drive_bloom(PG_FUNCTION_ARGS)
{
	int32		num_oids = PG_GETARG_INT32(0);
	int32		dup_freq = PG_GETARG_INT32(1);
	int32		count = PG_GETARG_INT32(2);
	MemoryContext mycontext;

	mycontext = AllocSetContextCreate(CurrentMemoryContext,
									  "drive_bloom work cxt",
									  ALLOCSET_DEFAULT_SIZES);

	while (count-- > 0)
	{
		List *roles_list = NIL;
		Oid nextrole = 1;
		bloom_filter *roles_list_bf = NULL;

		MemoryContext oldcontext;

		oldcontext = MemoryContextSwitchTo(mycontext);

		for (int i = 0; i < num_oids; i++)
		{
			roles_list = roles_list_append(roles_list, nextrole,
										   &roles_list_bf);
			if (dup_freq > 0 && i % dup_freq == 0)
				roles_list = roles_list_append(roles_list, nextrole,
											   &roles_list_bf);
			nextrole++;
		}

		if (roles_list_bf)
			bloom_free(roles_list_bf);

		MemoryContextSwitchTo(oldcontext);

		MemoryContextReset(mycontext);

		CHECK_FOR_INTERRUPTS();
	}

	PG_RETURN_VOID();
}