From 233fafb93efe2daf8e94178e0b51c47ee776f73c Mon Sep 17 00:00:00 2001
From: Maxim Orlov <orlovmg@gmail.com>
Date: Thu, 3 Nov 2022 18:08:35 +0300
Subject: [PATCH v1 2/2] Add tests for palloc_aligned

---
 src/test/regress/expected/misc.out |  14 ++++
 src/test/regress/regress.c         | 101 +++++++++++++++++++++++++++++
 src/test/regress/sql/misc.sql      |  11 ++++
 3 files changed, 126 insertions(+)

diff --git a/src/test/regress/expected/misc.out b/src/test/regress/expected/misc.out
index 6e816c57f1..3ee36d9d68 100644
--- a/src/test/regress/expected/misc.out
+++ b/src/test/regress/expected/misc.out
@@ -396,3 +396,17 @@ SELECT *, (equipment(CAST((h.*) AS hobbies_r))).name FROM hobbies_r h;
 --
 -- rewrite rules
 --
+--
+-- Test aligned memory allocation
+--
+CREATE FUNCTION test_palloc_aligned()
+   RETURNS void
+   AS :'regresslib'
+   LANGUAGE C STRICT;
+SELECT test_palloc_aligned();
+ test_palloc_aligned 
+---------------------
+ 
+(1 row)
+
+DROP FUNCTION test_palloc_aligned();
diff --git a/src/test/regress/regress.c b/src/test/regress/regress.c
index 548afb4438..499f6e3ac2 100644
--- a/src/test/regress/regress.c
+++ b/src/test/regress/regress.c
@@ -44,6 +44,7 @@
 #include "utils/geo_decls.h"
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
+#include "utils/palloc.h"
 #include "utils/rel.h"
 #include "utils/typcache.h"
 
@@ -1257,3 +1258,103 @@ get_columns_length(PG_FUNCTION_ARGS)
 
 	PG_RETURN_INT32(column_offset);
 }
+
+static void
+test_palloc_aligned_validate_and_free(void *p, Size size, Size alignto,
+									  bool check_zeroed)
+{
+	char   *itr;
+	Size	i;
+
+	if (p == NULL)
+		elog(ERROR, "could not allocate %zu byte(s)", size);
+
+	if ((uintptr_t) p % alignto != 0)
+		elog(ERROR, "memory block are not aligned to %zu", alignto);
+
+	if (check_zeroed)
+	{
+		for (i = 0, itr = p; i < size; ++i, ++itr)
+			if (*itr != 0)
+				elog(ERROR, "memory block are not zeroed");
+	}
+
+	pfree(p);
+}
+
+static void
+test_palloc_aligned_alloc_and_test(Size size, Size alignto, int flags)
+{
+	void   *p;
+
+	p = palloc_aligned(size, alignto, flags);
+	test_palloc_aligned_validate_and_free(p, size, alignto,
+										  (flags & MCXT_ALLOC_ZERO) != 0);
+
+	p = MemoryContextAllocAligned(CurrentMemoryContext, size, alignto, flags);
+	test_palloc_aligned_validate_and_free(p, size, alignto,
+										  (flags & MCXT_ALLOC_ZERO) != 0);
+}
+
+/*
+ * Test aligned memory allocation.
+ */
+PG_FUNCTION_INFO_V1(test_palloc_aligned);
+Datum
+test_palloc_aligned(PG_FUNCTION_ARGS)
+{
+	Size	alignments[] = {2, 4, 8, 16, 32},
+			sizes[] = {1, 100, 1000},
+			i,
+			j,
+			alignment,
+			n;
+	void   *p;
+	int		flags;
+
+	for (i = 0; i < lengthof(alignments); ++i)
+	{
+		alignment = alignments[i];
+		/*
+		 * This behavior is implementation defined, but we should not fail here.
+		 */
+		p = palloc_aligned(0, alignment, MCXT_ALLOC_NO_OOM);
+		if (p != NULL)
+			pfree(p);
+
+		n = 0;
+		p = palloc_aligned(0, alignment, MCXT_ALLOC_NO_OOM | MCXT_ALLOC_ZERO);
+		if (p != NULL)
+			pfree(p);
+
+		p = palloc_aligned(0, alignment, MCXT_ALLOC_NO_OOM | MCXT_ALLOC_ZERO |
+										 MCXT_ALLOC_HUGE);
+		if (p != NULL)
+			pfree(p);
+
+		/*
+		 * "Normal" cases.
+		 */
+		for (j = 0; j < lengthof(sizes); ++j)
+		{
+			n = sizes[j];
+			/* No flags */
+			flags = 0;
+			test_palloc_aligned_alloc_and_test(n, alignment, flags);
+
+			/* MCXT_ALLOC_NO_OOM */
+			flags |= MCXT_ALLOC_NO_OOM;
+			test_palloc_aligned_alloc_and_test(n, alignment, flags);
+
+			/* MCXT_ALLOC_NO_OOM | MCXT_ALLOC_ZERO */
+			flags |= MCXT_ALLOC_ZERO;
+			test_palloc_aligned_alloc_and_test(n, alignment, flags);
+
+			/* MCXT_ALLOC_NO_OOM | MCXT_ALLOC_ZERO | MCXT_ALLOC_HUGE */
+			flags |= MCXT_ALLOC_HUGE;
+			test_palloc_aligned_alloc_and_test(n, alignment, flags);
+		}
+	}
+
+	PG_RETURN_VOID();
+}
diff --git a/src/test/regress/sql/misc.sql b/src/test/regress/sql/misc.sql
index 165a2e175f..43f3bb3476 100644
--- a/src/test/regress/sql/misc.sql
+++ b/src/test/regress/sql/misc.sql
@@ -273,3 +273,14 @@ SELECT *, (equipment(CAST((h.*) AS hobbies_r))).name FROM hobbies_r h;
 --
 -- rewrite rules
 --
+
+--
+-- Test aligned memory allocation
+--
+CREATE FUNCTION test_palloc_aligned()
+   RETURNS void
+   AS :'regresslib'
+   LANGUAGE C STRICT;
+
+SELECT test_palloc_aligned();
+DROP FUNCTION test_palloc_aligned();
-- 
2.37.0 (Apple Git-136)