From 154e570a6cc0d1950b0b327a445f7f3c12e00892 Mon Sep 17 00:00:00 2001
From: Melanie Plageman <melanieplageman@gmail.com>
Date: Sat, 6 Apr 2024 12:41:05 -0400
Subject: [PATCH v19 10/21] Add UnifiedTBMIterator: common interface for
 TBMIterators

Add a new struct, UnifiedTBMIterator, which encapsulates the serial and
parallel TID Bitmap iterator interfaces. This simplifies call sites for
callers supporting both parallel and serial TID bitmap access.

Author: Melanie Plageman
Reviewed-by: Tomas Vondra
Discussion: https://postgr.es/m/CAAKRu_ZwCwWFeL_H3ia26bP2e7HiKLWt0ZmGXPVwPO6uXq0vaA%40mail.gmail.com
---
 src/backend/nodes/tidbitmap.c    | 55 ++++++++++++++++++++++++++++++++
 src/include/nodes/tidbitmap.h    | 17 ++++++++++
 src/tools/pgindent/typedefs.list |  1 +
 3 files changed, 73 insertions(+)

diff --git a/src/backend/nodes/tidbitmap.c b/src/backend/nodes/tidbitmap.c
index e8ab5d78fcc..c0628fa8a56 100644
--- a/src/backend/nodes/tidbitmap.c
+++ b/src/backend/nodes/tidbitmap.c
@@ -1556,3 +1556,58 @@ tbm_calculate_entries(double maxbytes)
 
 	return nbuckets;
 }
+
+/*
+ * Start iteration on a shared or non-shared bitmap iterator. Note that tbm
+ * will only be provided by serial BitmapHeapScan callers. dsa and dsp will
+ * only be provided by parallel BitmapHeapScan callers.
+ */
+void
+unified_tbm_begin_iterate(UnifiedTBMIterator *iterator, TIDBitmap *tbm,
+						  dsa_area *dsa, dsa_pointer dsp)
+{
+	Assert(iterator);
+
+	iterator->serial = NULL;
+	iterator->parallel = NULL;
+	iterator->exhausted = false;
+
+	/* Allocate a private iterator and attach the shared state to it */
+	if (DsaPointerIsValid(dsp))
+		iterator->parallel = tbm_attach_shared_iterate(dsa, dsp);
+	else
+		iterator->serial = tbm_begin_iterate(tbm);
+}
+
+/*
+ * Clean up shared or non-shared bitmap iterator.
+ */
+void
+unified_tbm_end_iterate(UnifiedTBMIterator *iterator)
+{
+	Assert(iterator);
+
+	if (iterator->serial)
+		tbm_end_iterate(iterator->serial);
+	else if (iterator->parallel)
+		tbm_end_shared_iterate(iterator->parallel);
+
+	iterator->serial = NULL;
+	iterator->parallel = NULL;
+	iterator->exhausted = true;
+}
+
+/*
+ * Get the next TBMIterateResult from the shared or non-shared bitmap iterator.
+ */
+TBMIterateResult *
+unified_tbm_iterate(UnifiedTBMIterator *iterator)
+{
+	Assert(iterator);
+	Assert(!iterator->exhausted);
+
+	if (iterator->serial)
+		return tbm_iterate(iterator->serial);
+	else
+		return tbm_shared_iterate(iterator->parallel);
+}
diff --git a/src/include/nodes/tidbitmap.h b/src/include/nodes/tidbitmap.h
index 1945f0639bf..40dbcf7d3d8 100644
--- a/src/include/nodes/tidbitmap.h
+++ b/src/include/nodes/tidbitmap.h
@@ -46,6 +46,17 @@ typedef struct TBMIterateResult
 	OffsetNumber offsets[FLEXIBLE_ARRAY_MEMBER];
 } TBMIterateResult;
 
+/*
+ * Callers with both serial and parallel implementations can use this unified
+ * API.
+ */
+typedef struct UnifiedTBMIterator
+{
+	TBMIterator *serial;
+	TBMSharedIterator *parallel;
+	bool		exhausted;
+} UnifiedTBMIterator;
+
 /* function prototypes in nodes/tidbitmap.c */
 
 extern TIDBitmap *tbm_create(long maxbytes, dsa_area *dsa);
@@ -72,4 +83,10 @@ extern TBMSharedIterator *tbm_attach_shared_iterate(dsa_area *dsa,
 													dsa_pointer dp);
 extern long tbm_calculate_entries(double maxbytes);
 
+extern void unified_tbm_begin_iterate(UnifiedTBMIterator *iterator, TIDBitmap *tbm,
+									  dsa_area *dsa, dsa_pointer dsp);
+extern void unified_tbm_end_iterate(UnifiedTBMIterator *iterator);
+extern TBMIterateResult *unified_tbm_iterate(UnifiedTBMIterator *iterator);
+
+
 #endif							/* TIDBITMAP_H */
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index f87e8b80ec0..71686e79abb 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2996,6 +2996,7 @@ UV
 UVersionInfo
 UnicodeNormalizationForm
 UnicodeNormalizationQC
+UnifiedTBMIterator
 Unique
 UniquePath
 UniquePathMethod
-- 
2.40.1