From f32760f730048b35c1045f005a1b2c9898a759e6 Mon Sep 17 00:00:00 2001
From: David Rowley <dgrowley@gmail.com>
Date: Mon, 6 Nov 2017 14:39:33 +1300
Subject: [PATCH] Add bms_add_range to add members within the specified range

The same behavior could be obtained by looping and using bms_add_member,
however, using bms_add_range operates at the bitmapword level and should
be far faster when the range is large.

Author: David Rowley
---
 src/backend/nodes/bitmapset.c | 74 +++++++++++++++++++++++++++++++++++++++++++
 src/include/nodes/bitmapset.h |  1 +
 2 files changed, 75 insertions(+)

diff --git a/src/backend/nodes/bitmapset.c b/src/backend/nodes/bitmapset.c
index bf8545d..34d242b 100644
--- a/src/backend/nodes/bitmapset.c
+++ b/src/backend/nodes/bitmapset.c
@@ -785,6 +785,80 @@ bms_add_members(Bitmapset *a, const Bitmapset *b)
 }
 
 /*
+ * bms_add_range
+ *		Add members in the range of 'lower' to 'upper' to the set.
+ *
+ * Note this could also be done by calling bms_add_member in a loop, however,
+ * using this function will be faster when the range is large as we work with
+ * at the bitmapword level rather than at bit level.
+ */
+Bitmapset *
+bms_add_range(Bitmapset *a, int lower, int upper)
+{
+	int			lwordnum,
+				uwordnum,
+				wordnum;
+
+	if (lower < 0 || upper < 0)
+		elog(ERROR, "negative bitmapset member not allowed");
+	if (lower > upper)
+		elog(ERROR, "lower range must not be above upper range");
+	uwordnum = WORDNUM(upper);
+
+	if (a == NULL)
+	{
+		a = (Bitmapset *) palloc0(BITMAPSET_SIZE(uwordnum + 1));
+		a->nwords = uwordnum + 1;
+	}
+
+	/* ensure we have enough words to store the upper bit */
+	else if (uwordnum >= a->nwords)
+	{
+		int			oldnwords = a->nwords;
+		int			i;
+
+		a = (Bitmapset *) repalloc(a, BITMAPSET_SIZE(uwordnum + 1));
+		a->nwords = uwordnum + 1;
+		/* zero out the enlarged portion */
+		for (i = oldnwords; i < a->nwords; i++)
+			a->words[i] = 0;
+	}
+
+	wordnum = lwordnum = WORDNUM(lower);
+
+	/*
+	 * Starting at lower's wordnum, loop over each word up to upper's wordnum.
+	 * Along the way set all bits inclusively between lower and upper to 1. We
+	 * only need to handle the lwordnum and uwordnum specially so we don't set
+	 * any bits outside of the range.
+	 */
+	while (wordnum <= uwordnum)
+	{
+		bitmapword mask = (bitmapword) ~0;
+
+		/* If working on the lower word, zero out bits below 'lower'. */
+		if (wordnum == lwordnum)
+		{
+			int lbitnum = BITNUM(lower);
+			mask >>= lbitnum;
+			mask <<= lbitnum;
+		}
+
+		/* Likewise, if working on the upper word, zero bits above 'upper' */
+		if (wordnum == uwordnum)
+		{
+			int ushiftbits = BITS_PER_BITMAPWORD - (BITNUM(upper) + 1);
+			mask <<= ushiftbits;
+			mask >>= ushiftbits;
+		}
+
+		a->words[wordnum++] |= mask;
+	}
+
+	return a;
+}
+
+/*
  * bms_int_members - like bms_intersect, but left input is recycled
  */
 Bitmapset *
diff --git a/src/include/nodes/bitmapset.h b/src/include/nodes/bitmapset.h
index aa3fb25..3b62a97 100644
--- a/src/include/nodes/bitmapset.h
+++ b/src/include/nodes/bitmapset.h
@@ -90,6 +90,7 @@ extern bool bms_is_empty(const Bitmapset *a);
 extern Bitmapset *bms_add_member(Bitmapset *a, int x);
 extern Bitmapset *bms_del_member(Bitmapset *a, int x);
 extern Bitmapset *bms_add_members(Bitmapset *a, const Bitmapset *b);
+extern Bitmapset *bms_add_range(Bitmapset *a, int lower, int upper);
 extern Bitmapset *bms_int_members(Bitmapset *a, const Bitmapset *b);
 extern Bitmapset *bms_del_members(Bitmapset *a, const Bitmapset *b);
 extern Bitmapset *bms_join(Bitmapset *a, Bitmapset *b);
-- 
1.9.5.msysgit.1