feat: add SetRange method to BitSet

Author: T-J-L

bitset.go

@@ -289,6 +289,48 @@ func (b *BitSet) SetTo(i uint, value bool) *BitSet {
 	return b.Clear(i)
 }
 
+// SetRange sets bits in [start, end) to 1, the capacity of the bitset is
+// automatically increased accordingly.
+// Warning: using a very large value for 'end'
+// may lead to a memory shortage and a panic: the caller is responsible
+// for providing sensible parameters in line with their memory capacity.
+func (b *BitSet) SetRange(start, end uint) *BitSet {
+	if start >= end {
+		return b
+	}
+
+	if end-1 >= b.length {
+		b.extendSet(end - 1)
+	}
+
+	startWord := start >> log2WordSize
+	endWord := (end - 1) >> log2WordSize // inclusive, the word holding bit end-1
+
+	// e.g. start = 71 -> wordsIndex(start) = 7
+	//  firstMask = ^uint64(0) << 7 = 0b111111....11110000000
+	// keeps the bits below start untouched
+	firstMask := ^uint64(0) << wordsIndex(start)
+
+	// e.g. end = 135 -> wordsIndex(-end) = 57, see FlipRange for the
+	// modular arithmetic of the unary minus
+	//  lastMask = ^uint64(0) >> 57 = 0b00000....0001111111
+	// keeps the bits from end on untouched
+	lastMask := ^uint64(0) >> wordsIndex(-end)
+
+	if startWord == endWord { // the whole range lives in a single word
+		b.set[startWord] |= firstMask & lastMask
+		return b
+	}
+
+	b.set[startWord] |= firstMask
+	for i := startWord + 1; i < endWord; i++ {
+		b.set[i] = ^uint64(0)
+	}
+	b.set[endWord] |= lastMask
+
+	return b
+}
+
 // Flip bit at i.
 // Warning: using a very large value for 'i'
 // may lead to a memory shortage and a panic: the caller is responsible

bitset_benchmark_test.go

@@ -670,3 +670,28 @@ func BenchmarkRank(b *testing.B) {
 		})
 	}
 }
+
+func BenchmarkSetVsSetRange(b *testing.B) {
+	const (
+		sz         = 100_000
+		rangeStart = 5_000
+		rangeEnd   = 10_000
+	)
+
+	s := New(uint(sz))
+
+	b.Run("Set", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			s.ClearAll()
+			for j := uint(rangeStart); j < rangeEnd; j++ {
+				s.Set(j)
+			}
+		}
+	})
+	b.Run("SetRange", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			s.ClearAll()
+			s.SetRange(rangeStart, rangeEnd)
+		}
+	})
+}

bitset_test.go

@@ -1969,6 +1969,49 @@ func TestFlipRange(t *testing.T) {
 	}
 }
 
+func TestSetRange(t *testing.T) {
+	b := new(BitSet)
+	c := b.SetRange(4, 25)
+	if c.length != 25 {
+		t.Error("Unexpected value: ", c.length)
+		return
+	}
+	if c.Count() != 21 {
+		t.Error("Unexpected value: ", c.Count())
+		return
+	}
+	// an empty range is a no-op, even beyond the current length
+	if d := c.SetRange(100, 100); d.length != 25 || d.Count() != 21 {
+		t.Error("Unexpected value: ", d.length, d.Count())
+		return
+	}
+	//
+	for i := uint(0); i < 256; i++ {
+		for j := uint(0); j <= i; j++ {
+			bits := New(i)
+			bits.SetRange(0, j)
+			c := bits.Count()
+			if c != j {
+				t.Error("Unexpected value: ", c, " expected: ", j)
+				return
+			}
+		}
+	}
+	// every sub-range of a 256-bit set, checked bit by bit
+	for start := uint(0); start < 256; start++ {
+		for end := start; end < 256; end++ {
+			bits := New(256)
+			bits.SetRange(start, end)
+			for i := uint(0); i < 256; i++ {
+				if bits.Test(i) != (i >= start && i < end) {
+					t.Error("Unexpected value at ", i, " for range [", start, ",", end, ")")
+					return
+				}
+			}
+		}
+	}
+}
+
 func TestCopy(t *testing.T) {
 	a := New(10)
 	if a.Copy(nil) != 0 {