From bff879a76a16d28bba4a3859ad11650a642917d2 Mon Sep 17 00:00:00 2001
From: Michel Pelletier <michel@onesparse.com>
Date: Sat, 23 Nov 2024 21:01:28 -0800
Subject: [PATCH] Add example test module for expanded objects.

This is a very simple template for creating expanded objects
that keeps track of the number of times it has been expanded.

Future feature support for expanded objects should showecase
those features here.

Discussion: https://www.postgresql.org/message-id/CACxu%3DvJNMj1MqqUiwATuazoewireaN%3D7nskD5V-CBEcrs_K6Vg%40mail.gmail.com
---
 src/test/modules/pgexpanded/.gitignore        |   3 +
 src/test/modules/pgexpanded/Makefile          |  21 +++
 src/test/modules/pgexpanded/README.md         |  21 +++
 .../pgexpanded/expected/pgexpanded.out        |  92 +++++++++
 .../modules/pgexpanded/pgexpanded--1.0.sql    |  48 +++++
 src/test/modules/pgexpanded/pgexpanded.c      | 177 ++++++++++++++++++
 .../modules/pgexpanded/pgexpanded.control     |   5 +
 src/test/modules/pgexpanded/pgexpanded.h      |  87 +++++++++
 .../modules/pgexpanded/sql/pgexpanded.sql     |   5 +
 9 files changed, 459 insertions(+)
 create mode 100644 src/test/modules/pgexpanded/.gitignore
 create mode 100644 src/test/modules/pgexpanded/Makefile
 create mode 100644 src/test/modules/pgexpanded/README.md
 create mode 100644 src/test/modules/pgexpanded/expected/pgexpanded.out
 create mode 100644 src/test/modules/pgexpanded/pgexpanded--1.0.sql
 create mode 100644 src/test/modules/pgexpanded/pgexpanded.c
 create mode 100644 src/test/modules/pgexpanded/pgexpanded.control
 create mode 100644 src/test/modules/pgexpanded/pgexpanded.h
 create mode 100644 src/test/modules/pgexpanded/sql/pgexpanded.sql

diff --git a/src/test/modules/pgexpanded/.gitignore b/src/test/modules/pgexpanded/.gitignore
new file mode 100644
index 0000000000..44d119cfcc
--- /dev/null
+++ b/src/test/modules/pgexpanded/.gitignore
@@ -0,0 +1,3 @@
+# Generated subdirectories
+/log/
+/results/
diff --git a/src/test/modules/pgexpanded/Makefile b/src/test/modules/pgexpanded/Makefile
new file mode 100644
index 0000000000..27aa55ddaa
--- /dev/null
+++ b/src/test/modules/pgexpanded/Makefile
@@ -0,0 +1,21 @@
+# src/test/modules/pgexpanded/Makefile
+
+MODULES = pgexpanded
+
+EXTENSION = pgexpanded
+DATA = pgexpanded--1.0.sql
+PGFILEDESC = "pgexpanded - template for expanded datum"
+
+REGRESS = pgexpanded
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/pgexpanded
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
+
diff --git a/src/test/modules/pgexpanded/README.md b/src/test/modules/pgexpanded/README.md
new file mode 100644
index 0000000000..20b11bc690
--- /dev/null
+++ b/src/test/modules/pgexpanded/README.md
@@ -0,0 +1,21 @@
+# pgexpanded
+
+This is an example postgres extension that shows how to implement an
+"expanded" data type in C as described [in this
+documentation](https://www.postgresql.org/docs/current/xtypes.html):
+
+*"Another feature that's enabled by TOAST support is the possibility of
+having an expanded in-memory data representation that is more
+convenient to work with than the format that is stored on disk. The
+regular or “flat” varlena storage format is ultimately just a blob of
+bytes; it cannot for example contain pointers, since it may get copied
+to other locations in memory. For complex data types, the flat format
+may be quite expensive to work with, so PostgreSQL provides a way to
+“expand” the flat format into a representation that is more suited to
+computation, and then pass that format in-memory between functions of
+the data type."*
+
+This repository provides a simple, compilable and runnable example
+expanded data type that can be used as a basis for other extensions.
+By way of trivial example, it shows how to expand a data type that
+keeps track of the number of expansions it's gone through.
diff --git a/src/test/modules/pgexpanded/expected/pgexpanded.out b/src/test/modules/pgexpanded/expected/pgexpanded.out
new file mode 100644
index 0000000000..748dbc7f36
--- /dev/null
+++ b/src/test/modules/pgexpanded/expected/pgexpanded.out
@@ -0,0 +1,92 @@
+SET client_min_messages = 'debug1';
+CREATE EXTENSION pgexpanded;
+DEBUG:  executing extension script for "pgexpanded" version '1.0'
+SELECT '0'::exobj;
+DEBUG:  exobj_in
+LINE 1: SELECT '0'::exobj;
+               ^
+DEBUG:  new_expanded_exobj
+LINE 1: SELECT '0'::exobj;
+               ^
+DEBUG:  exobj_get_flat_size
+LINE 1: SELECT '0'::exobj;
+               ^
+DEBUG:  exobj_flatten_into
+LINE 1: SELECT '0'::exobj;
+               ^
+DEBUG:  DatumGetExobj
+DEBUG:  new_expanded_exobj
+DEBUG:  exobj_out
+DEBUG:  context_callback_exobj_free
+DEBUG:  context_callback_exobj_free
+ exobj 
+-------
+ 2
+(1 row)
+
+SELECT test_expand('0'::exobj);
+DEBUG:  exobj_in
+LINE 1: SELECT test_expand('0'::exobj);
+                           ^
+DEBUG:  new_expanded_exobj
+LINE 1: SELECT test_expand('0'::exobj);
+                           ^
+DEBUG:  exobj_get_flat_size
+LINE 1: SELECT test_expand('0'::exobj);
+                           ^
+DEBUG:  exobj_flatten_into
+LINE 1: SELECT test_expand('0'::exobj);
+                           ^
+DEBUG:  exobj_info
+DEBUG:  DatumGetExobj
+DEBUG:  new_expanded_exobj
+DEBUG:  context_callback_exobj_free
+NOTICE:  expand count 2
+DEBUG:  DatumGetExobj
+DEBUG:  new_expanded_exobj
+DEBUG:  exobj_out
+DEBUG:  context_callback_exobj_free
+DEBUG:  context_callback_exobj_free
+ test_expand 
+-------------
+ 2
+(1 row)
+
+SELECT test_expand_expand('0'::exobj);
+DEBUG:  exobj_in
+LINE 1: SELECT test_expand_expand('0'::exobj);
+                                  ^
+DEBUG:  new_expanded_exobj
+LINE 1: SELECT test_expand_expand('0'::exobj);
+                                  ^
+DEBUG:  exobj_get_flat_size
+LINE 1: SELECT test_expand_expand('0'::exobj);
+                                  ^
+DEBUG:  exobj_flatten_into
+LINE 1: SELECT test_expand_expand('0'::exobj);
+                                  ^
+DEBUG:  exobj_info
+DEBUG:  DatumGetExobj
+DEBUG:  new_expanded_exobj
+DEBUG:  context_callback_exobj_free
+NOTICE:  expand expand count 2
+DEBUG:  exobj_info
+DEBUG:  DatumGetExobj
+DEBUG:  new_expanded_exobj
+DEBUG:  context_callback_exobj_free
+NOTICE:  expand count 2
+DEBUG:  exobj_info
+DEBUG:  DatumGetExobj
+DEBUG:  new_expanded_exobj
+DEBUG:  context_callback_exobj_free
+NOTICE:  expand count 2
+DEBUG:  DatumGetExobj
+DEBUG:  new_expanded_exobj
+DEBUG:  exobj_out
+DEBUG:  context_callback_exobj_free
+DEBUG:  context_callback_exobj_free
+ test_expand_expand 
+--------------------
+ 2
+(1 row)
+
diff --git a/src/test/modules/pgexpanded/pgexpanded--1.0.sql b/src/test/modules/pgexpanded/pgexpanded--1.0.sql
new file mode 100644
index 0000000000..7f5d7e75ad
--- /dev/null
+++ b/src/test/modules/pgexpanded/pgexpanded--1.0.sql
@@ -0,0 +1,48 @@
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION pgexpanded" to load this file. \quit
+
+CREATE TYPE exobj;
+
+CREATE FUNCTION exobj_in(cstring)
+RETURNS exobj
+AS '$libdir/pgexpanded', 'exobj_in'
+LANGUAGE C IMMUTABLE STRICT;
+
+CREATE FUNCTION exobj_out(exobj)
+RETURNS cstring
+AS '$libdir/pgexpanded', 'exobj_out'
+LANGUAGE C IMMUTABLE STRICT;
+
+CREATE TYPE exobj (
+    input = exobj_in,
+    output = exobj_out,
+    alignment = int4,
+    storage = 'extended',
+    internallength = -1
+);
+
+CREATE FUNCTION info(exobj)
+RETURNS bigint
+AS '$libdir/pgexpanded', 'exobj_info'
+LANGUAGE C STABLE;
+
+create or replace function test_expand(obj exobj) returns exobj language plpgsql as
+    $$
+    declare
+        i bigint = info(obj);
+    begin
+        raise notice 'expand count %', i;
+        return obj;
+    end;
+    $$;
+
+create or replace function test_expand_expand(obj exobj) returns exobj language plpgsql as
+    $$
+    declare
+        i bigint = info(obj);
+    begin
+        raise notice 'expand expand count %', i;
+        obj = test_expand(obj);
+        return test_expand(obj);
+    end;
+    $$;
diff --git a/src/test/modules/pgexpanded/pgexpanded.c b/src/test/modules/pgexpanded/pgexpanded.c
new file mode 100644
index 0000000000..562a1f2860
--- /dev/null
+++ b/src/test/modules/pgexpanded/pgexpanded.c
@@ -0,0 +1,177 @@
+#include "pgexpanded.h"
+PG_MODULE_MAGIC;
+
+/* Compute flattened size of storage needed for a exobj */
+static Size
+exobj_get_flat_size(ExpandedObjectHeader *eohptr) {
+	pgexpanded_Exobj *A = (pgexpanded_Exobj*) eohptr;
+	Size nbytes;
+
+	LOGF();
+
+	/* This is a sanity check that the object is initialized */
+	Assert(A->em_magic == exobj_MAGIC);
+
+	/* Use cached value if already computed */
+	if (A->flat_size) {
+		return A->flat_size;
+	}
+
+	// Add the overhead of the flat header to the size of the data
+	// payload
+	nbytes = PGEXPANDED_EXOBJ_OVERHEAD();
+	nbytes += sizeof(uint64_t);
+
+	/* Cache this value in the expanded object */
+	A->flat_size = nbytes;
+	return nbytes;
+}
+
+/* Flatten exobj into a pre-allocated result buffer that is
+   allocated_size in bytes.  */
+static void
+exobj_flatten_into(ExpandedObjectHeader *eohptr,
+				   void *result, Size allocated_size)  {
+	void *data;
+
+	/* Cast EOH pointer to expanded object, and result pointer to flat
+	   object */
+	pgexpanded_Exobj *A = (pgexpanded_Exobj *) eohptr;
+	pgexpanded_FlatExobj *flat = (pgexpanded_FlatExobj *) result;
+
+	LOGF();
+
+	/* Sanity check the object is valid */
+	Assert(A->em_magic == exobj_MAGIC);
+	Assert(allocated_size == A->flat_size);
+
+	/* Zero out the whole allocated buffer */
+	memset(flat, 0, allocated_size);
+
+	/* Get the pointer to the start of the flattened data and copy the
+	   expanded value into it */
+	data = PGEXPANDED_EXOBJ_DATA(flat);
+	memcpy(data, A->value, sizeof(int64_t));
+
+	/* Set the size of the varlena object */
+	SET_VARSIZE(flat, allocated_size);
+}
+
+/* Expand a flat exobj in to an Expanded one, return as Postgres Datum. */
+pgexpanded_Exobj *
+new_expanded_exobj(int64_t value, MemoryContext parentcontext) {
+	pgexpanded_Exobj *A;
+
+	MemoryContext objcxt, oldcxt;
+	MemoryContextCallback *ctxcb;
+
+	LOGF();
+
+	/* Create a new context that will hold the expanded object. */
+	objcxt = AllocSetContextCreate(parentcontext,
+								   "expanded exobj",
+								   ALLOCSET_DEFAULT_SIZES);
+
+	/* Allocate a new expanded exobj */
+	A = (pgexpanded_Exobj*)MemoryContextAlloc(objcxt,
+											  sizeof(pgexpanded_Exobj));
+
+	/* Initialize the ExpandedObjectHeader member with flattening
+	 * methods and the new object context */
+	EOH_init_header(&A->hdr, &exobj_methods, objcxt);
+
+	/* Used for debugging checks */
+	A->em_magic = exobj_MAGIC;
+
+	/* Switch to new object context */
+	oldcxt = MemoryContextSwitchTo(objcxt);
+
+	/* Get value from flat object and increment it */
+	A->value = palloc(sizeof(int64_t));
+	*(A->value) = value + 1;
+
+	/* Setting flat size to zero tells us the object has been written. */
+	A->flat_size = 0;
+
+	/* Create a context callback to free exobj when context is cleared */
+	ctxcb = MemoryContextAlloc(objcxt, sizeof(MemoryContextCallback));
+
+	ctxcb->func = context_callback_exobj_free;
+	ctxcb->arg = A;
+	MemoryContextRegisterResetCallback(objcxt, ctxcb);
+
+	/* Switch back to old context */
+	MemoryContextSwitchTo(oldcxt);
+	return A;
+}
+
+/* MemoryContextCallback function to free exobj data when their
+   context goes out of scope. */
+static void
+context_callback_exobj_free(void* ptr) {
+	pgexpanded_Exobj *A = (pgexpanded_Exobj *) ptr;
+	LOGF();
+	pfree(A->value);
+}
+
+/* Helper function to always expanded datum
+
+   This is used by PG_GETARG_EXOBJ */
+pgexpanded_Exobj *
+DatumGetExobj(Datum d) {
+	pgexpanded_Exobj *A;
+	pgexpanded_FlatExobj *flat;
+	int64_t *value;
+
+	LOGF();
+	if (VARATT_IS_EXTERNAL_EXPANDED(DatumGetPointer(d))) {
+		A = ExobjGetEOHP(d);
+		Assert(A->em_magic == exobj_MAGIC);
+		return A;
+	}
+	flat = (pgexpanded_FlatExobj*)PG_DETOAST_DATUM(d);
+	value = PGEXPANDED_EXOBJ_DATA(flat);
+	A = new_expanded_exobj(*value, CurrentMemoryContext);
+	return A;
+}
+
+Datum
+exobj_in(PG_FUNCTION_ARGS) {
+	char *input;
+	pgexpanded_Exobj *result;
+	int64_t value;
+	LOGF();
+	input = PG_GETARG_CSTRING(0);
+	value = strtoll(input, NULL, 10);
+	result = new_expanded_exobj(value, CurrentMemoryContext);
+	PGEXPANDED_RETURN_EXOBJ(result);
+}
+
+Datum
+exobj_out(PG_FUNCTION_ARGS)
+{
+	char *result;
+	pgexpanded_Exobj *A = PGEXPANDED_GETARG_EXOBJ(0);
+	LOGF();
+	result = palloc(32);
+	snprintf(result, sizeof(result), "%lld", (long long int) *A->value);
+	PG_RETURN_CSTRING(result);
+}
+
+Datum
+exobj_info(PG_FUNCTION_ARGS) {
+	pgexpanded_Exobj *A;
+	LOGF();
+	A = PGEXPANDED_GETARG_EXOBJ(0);
+	return Int64GetDatum(*A->value);
+}
+
+void
+_PG_init(void)
+{
+	LOGF();
+}
+/* Local Variables: */
+/* mode: c */
+/* c-file-style: "postgresql" */
+/* End: */
diff --git a/src/test/modules/pgexpanded/pgexpanded.control b/src/test/modules/pgexpanded/pgexpanded.control
new file mode 100644
index 0000000000..2e9d133484
--- /dev/null
+++ b/src/test/modules/pgexpanded/pgexpanded.control
@@ -0,0 +1,5 @@
+# pgexpanded extension
+comment = 'Example Postgres extension for expanded data types.'
+default_version = '1.0'
+relocatable = true
+requires = ''
diff --git a/src/test/modules/pgexpanded/pgexpanded.h b/src/test/modules/pgexpanded/pgexpanded.h
new file mode 100644
index 0000000000..4adebace5f
--- /dev/null
+++ b/src/test/modules/pgexpanded/pgexpanded.h
@@ -0,0 +1,87 @@
+#ifndef PGEXPANDED_H
+#define PGEXPANDED_H
+
+#include "postgres.h"
+#include "funcapi.h"
+#include "utils/expandeddatum.h"
+
+/* ID for debugging crosschecks */
+#define exobj_MAGIC 689276813
+
+#define LOGF() elog(DEBUG1, __func__)
+
+/* Flattened representation of exobj, used to store to disk.
+
+   The first 32 bits must the length of the data.  Actual flattened data
+   is appended after this struct and cannot exceed 1GB.
+*/
+typedef struct pgexpanded_FlatExobj {
+	int32 vl_len_;
+} pgexpanded_FlatExobj;
+
+/* Expanded representation of exobj.
+
+   When loaded from storage, the flattened representation is used to
+   build the exobj.  In this case, it's just a pointer to an integer.
+*/
+typedef struct pgexpanded_Exobj  {
+	ExpandedObjectHeader hdr;
+	int em_magic;
+	Size flat_size;
+	int64_t *value;
+} pgexpanded_Exobj;
+
+/* Callback function for freeing exobj arrays. */
+static void
+context_callback_exobj_free(void*);
+
+/* Expanded Object Header "methods" for flattening for storage */
+static Size
+exobj_get_flat_size(ExpandedObjectHeader *eohptr);
+
+static void
+exobj_flatten_into(ExpandedObjectHeader *eohptr,
+				   void *result, Size allocated_size);
+
+static const ExpandedObjectMethods exobj_methods = {
+	exobj_get_flat_size,
+	exobj_flatten_into
+};
+
+/* Create a new exobj datum. */
+pgexpanded_Exobj *
+new_expanded_exobj(int64_t value,  MemoryContext parentcontext);
+
+/* Helper function that either detoasts or expands. */
+pgexpanded_Exobj *DatumGetExobj(Datum d);
+
+/* Helper macro to detoast and expand exobjs arguments */
+#define PGEXPANDED_GETARG_EXOBJ(n)  DatumGetExobj(PG_GETARG_DATUM(n))
+
+/* Helper macro to return Expanded Object Header Pointer from exobj. */
+#define PGEXPANDED_RETURN_EXOBJ(A) return EOHPGetRWDatum(&(A)->hdr)
+
+/* Helper macro to compute flat exobj header size */
+#define PGEXPANDED_EXOBJ_OVERHEAD() MAXALIGN(sizeof(pgexpanded_FlatExobj))
+
+/* Helper macro to get pointer to beginning of exobj data. */
+#define PGEXPANDED_EXOBJ_DATA(a) ((int64_t *)(((char *) (a)) + PGEXPANDED_EXOBJ_OVERHEAD()))
+
+/* Help macro to cast generic Datum header pointer to expanded Exobj */
+#define ExobjGetEOHP(d) (pgexpanded_Exobj *) DatumGetEOHP(d);
+
+/* Public API functions */
+
+PG_FUNCTION_INFO_V1(exobj);
+PG_FUNCTION_INFO_V1(exobj_in);
+PG_FUNCTION_INFO_V1(exobj_out);
+PG_FUNCTION_INFO_V1(exobj_info);
+
+void
+_PG_init(void);
+
+#endif /* PGEXPANDED_H */
+/* Local Variables: */
+/* mode: c */
+/* c-file-style: "postgresql" */
+/* End: */
diff --git a/src/test/modules/pgexpanded/sql/pgexpanded.sql b/src/test/modules/pgexpanded/sql/pgexpanded.sql
new file mode 100644
index 0000000000..8858851b32
--- /dev/null
+++ b/src/test/modules/pgexpanded/sql/pgexpanded.sql
@@ -0,0 +1,5 @@
+SET client_min_messages = 'debug1';
+CREATE EXTENSION pgexpanded;
+SELECT '0'::exobj;
+SELECT test_expand('0'::exobj);
+SELECT test_expand_expand('0'::exobj);
-- 
2.34.1

