From 5d0115b068ecb01d791eab5f8a78a6d25b9cf45c Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Wed, 14 Sep 2022 12:38:01 +0000
Subject: [PATCH v6 1/3] Support pg_lsearch8_eq and pg_lsearch8_ge

---
 src/include/port/pg_lfind.h                   |  71 ++++++++
 src/include/port/simd.h                       | 155 +++++++++++++++++-
 .../test_lfind/expected/test_lfind.out        |  12 ++
 .../modules/test_lfind/sql/test_lfind.sql     |   2 +
 .../modules/test_lfind/test_lfind--1.0.sql    |   8 +
 src/test/modules/test_lfind/test_lfind.c      | 139 ++++++++++++++++
 6 files changed, 378 insertions(+), 9 deletions(-)

diff --git a/src/include/port/pg_lfind.h b/src/include/port/pg_lfind.h
index 0625cac6b5..583f204763 100644
--- a/src/include/port/pg_lfind.h
+++ b/src/include/port/pg_lfind.h
@@ -80,6 +80,77 @@ pg_lfind8_le(uint8 key, uint8 *base, uint32 nelem)
 	return false;
 }
 
+/*
+ * pg_lsearch8
+ *
+ * Return the index of the element in 'base' that equals to 'key', otherwise return
+ * -1.
+ */
+static inline int
+pg_lsearch8(uint8 key, uint8 *base, uint32 nelem)
+{
+	uint32		i;
+
+	/* round down to multiple of vector length */
+	uint32		tail_idx = nelem & ~(sizeof(Vector8) - 1);
+	Vector8		chunk;
+
+	for (i = 0; i < tail_idx; i += sizeof(Vector8))
+	{
+		int idx;
+
+		vector8_load(&chunk, &base[i]);
+		if ((idx = vector8_search_eq(chunk, key)) != -1)
+			return i + idx;
+	}
+
+	/* Process the remaining elements one at a time. */
+	for (; i < nelem; i++)
+	{
+		if (key == base[i])
+			return i;
+	}
+
+	return -1;
+}
+
+
+/*
+ * pg_lsearch8_ge
+ *
+ * Return the index of the first element in 'base' that is greater than or equal to
+ * 'key'. Return nelem if there is no such element.
+ *
+ * Note that this function assumes the elements in 'base' are sorted.
+ */
+static inline int
+pg_lsearch8_ge(uint8 key, uint8 *base, uint32 nelem)
+{
+	uint32		i;
+
+	/* round down to multiple of vector length */
+	uint32		tail_idx = nelem & ~(sizeof(Vector8) - 1);
+	Vector8		chunk;
+
+	for (i = 0; i < tail_idx; i += sizeof(Vector8))
+	{
+		int idx;
+
+		vector8_load(&chunk, &base[i]);
+		if ((idx = vector8_search_ge(chunk, key)) != sizeof(Vector8))
+			return i + idx;
+	}
+
+	/* Process the remaining elements one at a time. */
+	for (; i < nelem; i++)
+	{
+		if (base[i] >= key)
+			break;
+	}
+
+	return i;
+}
+
 /*
  * pg_lfind32
  *
diff --git a/src/include/port/simd.h b/src/include/port/simd.h
index 61ae4ecf60..e2a99578a5 100644
--- a/src/include/port/simd.h
+++ b/src/include/port/simd.h
@@ -18,6 +18,8 @@
 #ifndef SIMD_H
 #define SIMD_H
 
+#include "port/pg_bitutils.h"
+
 #if (defined(__x86_64__) || defined(_M_AMD64))
 /*
  * SSE2 instructions are part of the spec for the 64-bit x86 ISA. We assume
@@ -88,14 +90,9 @@ static inline Vector32 vector32_or(const Vector32 v1, const Vector32 v2);
 static inline Vector8 vector8_ssub(const Vector8 v1, const Vector8 v2);
 #endif
 
-/*
- * comparisons between vectors
- *
- * Note: These return a vector rather than boolean, which is why we don't
- * have non-SIMD implementations.
- */
-#ifndef USE_NO_SIMD
+/* comparisons between vectors */
 static inline Vector8 vector8_eq(const Vector8 v1, const Vector8 v2);
+#ifndef USE_NO_SIMD
 static inline Vector32 vector32_eq(const Vector32 v1, const Vector32 v2);
 #endif
 
@@ -277,6 +274,140 @@ vector8_is_highbit_set(const Vector8 v)
 #endif
 }
 
+/*
+ * Return the bitmak of the high-bit of each element.
+ */
+static inline uint32
+vector8_highbit_mask(const Vector8 v)
+{
+#ifdef USE_SSE2
+	return (uint32) _mm_movemask_epi8(v);
+#elif defined(USE_NEON)
+	static const uint8 mask[16] = {
+        1 << 0, 1 << 1, 1 << 2, 1 << 3,
+        1 << 4, 1 << 5, 1 << 6, 1 << 7,
+        1 << 0, 1 << 1, 1 << 2, 1 << 3,
+        1 << 4, 1 << 5, 1 << 6, 1 << 7,
+      };
+
+    uint8x16_t masked = vandq_u8(vld1q_u8(mask), (uint8x16_t) vshrq_n_s8(v, 7));
+    uint8x16_t maskedhi = vextq_u8(masked, masked, 8);
+
+    return (uint32) vaddvq_u16((uint16x8_t) vzip1q_u8(masked, maskedhi));
+#else
+	uint32 mask = 0;
+
+	for (Size i = 0; i < sizeof(Vector8); i++)
+		mask |= (((const uint8 *) &v)[i] >> 7) << i;
+
+	return mask;
+#endif
+}
+
+/*
+ * Compare the given vectors and return the vector of minimum elements.
+ */
+static inline Vector8
+vector8_min(const Vector8 v1, const Vector8 v2)
+{
+#ifdef USE_SSE2
+	return _mm_min_epu8(v1, v2);
+#elif defined(USE_NEON)
+	return vminq_u8(v1, v2);
+#else /* USE_NO_SIMD */
+	Vector8 r = 0;
+	uint8 *rp = (uint8 *) &r;
+
+	for (Size i = 0; i < sizeof(Vector8); i++)
+		rp[i] = Min(((const uint8 *) &v1)[i], ((const uint8 *) &v2)[i]);
+
+	return r;
+#endif
+}
+
+/*
+ * Return the index of the element in the vector that equal to the given
+ * scalar. Otherwise, return -1.
+ */
+static inline int
+vector8_search_eq(const Vector8 v, const uint8 c)
+{
+	Vector8 keys = vector8_broadcast(c);
+	Vector8	cmp;
+	uint32	mask;
+	int		result;
+
+#ifdef USE_ASSERT_CHECKING
+	int		assert_result = -1;
+
+	for (Size i = 0; i < sizeof(Vector8); i++)
+	{
+		if (((const uint8 *) &v)[i] == c)
+		{
+			assert_result = i;
+			break;
+		}
+	}
+#endif							/* USE_ASSERT_CHECKING */
+
+	cmp = vector8_eq(keys, v);
+	mask = vector8_highbit_mask(cmp);
+
+	if (mask)
+		result = pg_rightmost_one_pos32(mask);
+	else
+		result = -1;
+
+	Assert(assert_result == result);
+	return result;
+}
+
+/*
+ * Return the index of the first element in the vector that is greater than
+ * or eual to the given scalar. Return sizeof(Vector8) if there is no such
+ * element.
+ *
+ * Note that this function assumes the elements in the vector are sorted.
+ */
+static inline int
+vector8_search_ge(const Vector8 v, const uint8 c)
+{
+	Vector8 keys = vector8_broadcast(c);
+	Vector8 min;
+	Vector8	cmp;
+	uint32	mask;
+	int		result;
+
+#ifdef USE_ASSERT_CHECKING
+	int		assert_result = -1;
+	Size	i;
+
+	for (i = 0; i < sizeof(Vector8); i++)
+	{
+		if (((const uint8 *) &v)[i] >= c)
+			break;
+	}
+	assert_result = i;
+#endif							/* USE_ASSERT_CHECKING */
+
+	/*
+	 * There is a bit more complicated than vector8_search_eq(), because
+	 * until recently no unsigned uint8 compasion instruction existed.
+	 * Therefore, we need to use vector8_min() to effectively get <= elements.
+	 */
+	min = vector8_min(v, keys);
+	cmp = vector8_eq(keys, min);
+	mask = vector8_highbit_mask(cmp);
+
+	if (mask)
+		result = pg_rightmost_one_pos32(mask);
+	else
+		result = sizeof(Vector8);
+
+	Assert(assert_result == result);
+	return result;
+}
+
 /*
  * Exactly like vector8_is_highbit_set except for the input type, so it
  * looks at each byte separately.
@@ -348,7 +479,6 @@ vector8_ssub(const Vector8 v1, const Vector8 v2)
  * Return a vector with all bits set in each lane where the the corresponding
  * lanes in the inputs are equal.
  */
-#ifndef USE_NO_SIMD
 static inline Vector8
 vector8_eq(const Vector8 v1, const Vector8 v2)
 {
@@ -356,9 +486,16 @@ vector8_eq(const Vector8 v1, const Vector8 v2)
 	return _mm_cmpeq_epi8(v1, v2);
 #elif defined(USE_NEON)
 	return vceqq_u8(v1, v2);
+#else /* USE_NO_SIMD */
+	Vector8 r = 0;
+	uint8 *rp = (uint8 *) &r;
+
+	for (Size i = 0; i < sizeof(Vector8); i++)
+		rp[i] = (((const uint8 *) &v1)[i] == ((const uint8 *) &v2)[i]) ? 0xFF : 0;
+
+	return r;
 #endif
 }
-#endif							/* ! USE_NO_SIMD */
 
 #ifndef USE_NO_SIMD
 static inline Vector32
diff --git a/src/test/modules/test_lfind/expected/test_lfind.out b/src/test/modules/test_lfind/expected/test_lfind.out
index 1d4b14e703..9416161955 100644
--- a/src/test/modules/test_lfind/expected/test_lfind.out
+++ b/src/test/modules/test_lfind/expected/test_lfind.out
@@ -22,3 +22,15 @@ SELECT test_lfind32();
  
 (1 row)
 
+SELECT test_lsearch8();
+ test_lsearch8 
+---------------
+ 
+(1 row)
+
+SELECT test_lsearch8_ge();
+ test_lsearch8_ge 
+------------------
+ 
+(1 row)
+
diff --git a/src/test/modules/test_lfind/sql/test_lfind.sql b/src/test/modules/test_lfind/sql/test_lfind.sql
index 766c640831..d0dbb142ec 100644
--- a/src/test/modules/test_lfind/sql/test_lfind.sql
+++ b/src/test/modules/test_lfind/sql/test_lfind.sql
@@ -8,3 +8,5 @@ CREATE EXTENSION test_lfind;
 SELECT test_lfind8();
 SELECT test_lfind8_le();
 SELECT test_lfind32();
+SELECT test_lsearch8();
+SELECT test_lsearch8_ge();
diff --git a/src/test/modules/test_lfind/test_lfind--1.0.sql b/src/test/modules/test_lfind/test_lfind--1.0.sql
index 81801926ae..13857cec3b 100644
--- a/src/test/modules/test_lfind/test_lfind--1.0.sql
+++ b/src/test/modules/test_lfind/test_lfind--1.0.sql
@@ -14,3 +14,11 @@ CREATE FUNCTION test_lfind8()
 CREATE FUNCTION test_lfind8_le()
 	RETURNS pg_catalog.void
 	AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION test_lsearch8()
+	RETURNS pg_catalog.void
+	AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION test_lsearch8_ge()
+	RETURNS pg_catalog.void
+	AS 'MODULE_PATHNAME' LANGUAGE C;
diff --git a/src/test/modules/test_lfind/test_lfind.c b/src/test/modules/test_lfind/test_lfind.c
index 82673d54c6..c494c27436 100644
--- a/src/test/modules/test_lfind/test_lfind.c
+++ b/src/test/modules/test_lfind/test_lfind.c
@@ -14,6 +14,7 @@
 #include "postgres.h"
 
 #include "fmgr.h"
+#include "lib/stringinfo.h"
 #include "port/pg_lfind.h"
 
 /*
@@ -115,6 +116,144 @@ test_lfind8_le(PG_FUNCTION_ARGS)
 	PG_RETURN_VOID();
 }
 
+static void
+test_lsearch8_internal(uint8 key)
+{
+	uint8		charbuf[LEN_WITH_TAIL(Vector8)];
+	const int	len_no_tail = LEN_NO_TAIL(Vector8);
+	const int	len_with_tail = LEN_WITH_TAIL(Vector8);
+	int			keypos;
+
+	memset(charbuf, 0xFF, len_with_tail);
+	/* search tail to test one-byte-at-a-time path */
+	keypos = len_with_tail - 1;
+	charbuf[keypos] = key;
+	if (key > 0x00 && (pg_lsearch8(key - 1, charbuf, len_with_tail) != -1))
+		elog(ERROR, "pg_lsearch8() found nonexistent element '0x%x'", key - 1);
+	if (key < 0xFF && (pg_lsearch8(key, charbuf, len_with_tail) != keypos))
+		elog(ERROR, "pg_lsearch8() did not find existing element '0x%x'", key);
+	if (key < 0xFE && (pg_lsearch8(key + 1, charbuf, len_with_tail) != -1))
+		elog(ERROR, "pg_lsearch8() found nonexistent element '0x%x'", key + 1);
+
+	memset(charbuf, 0xFF, len_with_tail);
+	/* search with vector operations */
+	keypos = len_no_tail - 1;
+	charbuf[keypos] = key;
+	if (key > 0x00 && (pg_lsearch8(key - 1, charbuf, len_no_tail) != -1))
+		elog(ERROR, "pg_lsearch8() found nonexistent element '0x%x'", key - 1);
+	if (key < 0xFF && (pg_lsearch8(key, charbuf, len_no_tail) != keypos))
+		elog(ERROR, "pg_lsearch8() did not find existing element '0x%x'", key);
+	if (key < 0xFE && (pg_lsearch8(key + 1, charbuf, len_no_tail) != -1))
+		elog(ERROR, "pg_lsearch8() found nonexistent element '0x%x'", key + 1);
+}
+
+PG_FUNCTION_INFO_V1(test_lsearch8);
+Datum
+test_lsearch8(PG_FUNCTION_ARGS)
+{
+	test_lsearch8_internal(0);
+	test_lsearch8_internal(1);
+	test_lsearch8_internal(0x7F);
+	test_lsearch8_internal(0x80);
+	test_lsearch8_internal(0x81);
+	test_lsearch8_internal(0xFD);
+	test_lsearch8_internal(0xFE);
+	test_lsearch8_internal(0xFF);
+
+	PG_RETURN_VOID();
+}
+
+static void
+report_lsearch8_error(uint8 *buf, int size, uint8 key, int result, int expected)
+{
+	StringInfoData bufstr;
+	char *sep = "";
+
+	initStringInfo(&bufstr);
+
+	for (int i = 0; i < size; i++)
+	{
+		appendStringInfo(&bufstr, "%s0x%02x", sep, buf[i]);
+		sep = ",";
+	}
+
+	elog(ERROR,
+		 "pg_lsearch8_ge returned %d, expected %d, key 0x%02x buffer %s",
+		 result, expected, key, bufstr.data);
+}
+
+/* workhorse for test_lsearch8_ge */
+static void
+test_lsearch8_ge_internal(uint8 *buf, uint8 key)
+{
+	const int	len_no_tail = LEN_NO_TAIL(Vector8);
+	const int	len_with_tail = LEN_WITH_TAIL(Vector8);
+	int			expected;
+	int			result;
+	int			i;
+
+	/* search tail to test one-byte-at-a-time path */
+	for (i = 0; i < len_with_tail; i++)
+	{
+		if (buf[i] >= key)
+			break;
+	}
+	expected = i;
+	result = pg_lsearch8_ge(key, buf, len_with_tail);
+
+	if (result != expected)
+		report_lsearch8_error(buf, len_with_tail, key, result, expected);
+
+	/* search with vector operations */
+	for (i = 0; i < len_no_tail; i++)
+	{
+		if (buf[i] >= key)
+			break;
+	}
+	expected = i;
+	result = pg_lsearch8_ge(key, buf, len_no_tail);
+
+	if (result != expected)
+		report_lsearch8_error(buf, len_no_tail, key, result, expected);
+}
+
+static int
+cmp(const void *p1, const void *p2)
+{
+	uint8	v1 = *((const uint8 *) p1);
+	uint8	v2 = *((const uint8 *) p2);
+
+	if (v1 < v2)
+		return -1;
+	if (v1 > v2)
+		return 1;
+	return 0;
+}
+
+PG_FUNCTION_INFO_V1(test_lsearch8_ge);
+Datum
+test_lsearch8_ge(PG_FUNCTION_ARGS)
+{
+	uint8		charbuf[LEN_WITH_TAIL(Vector8)];
+	const int	len_with_tail = LEN_WITH_TAIL(Vector8);
+
+	for (int i = 0; i < len_with_tail; i++)
+		charbuf[i] = (uint8) rand();
+
+	qsort(charbuf, len_with_tail, sizeof(uint8), cmp);
+
+	test_lsearch8_ge_internal(charbuf, 0);
+	test_lsearch8_ge_internal(charbuf, 1);
+	test_lsearch8_ge_internal(charbuf, 0x7F);
+	test_lsearch8_ge_internal(charbuf, 0x80);
+	test_lsearch8_ge_internal(charbuf, 0x81);
+	test_lsearch8_ge_internal(charbuf, 0xFD);
+	test_lsearch8_ge_internal(charbuf, 0xFE);
+	test_lsearch8_ge_internal(charbuf, 0xFF);
+
+	PG_RETURN_VOID();
+}
+
 PG_FUNCTION_INFO_V1(test_lfind32);
 Datum
 test_lfind32(PG_FUNCTION_ARGS)
-- 
2.31.1

