Merge pull request #522 from e2b-dev/efficient-range-iteration

Add Ranges iterator

Author: lemire

iter.go

@@ -42,3 +42,108 @@ func Unset(b *Bitmap, min, max uint32) iter.Seq[uint32] {
 		}
 	}
 }
+
+// Ranges iterates contiguous ranges of values present in the bitmap as
+// half-open [start, endExclusive) pairs. endExclusive is uint64 to represent
+// ranges that include MaxUint32. Ranges spanning container boundaries are merged.
+func (b *Bitmap) Ranges() iter.Seq2[uint32, uint64] {
+	return func(yield func(uint32, uint64) bool) {
+		ra := &b.highlowcontainer
+		keys := ra.keys
+		containers := ra.containers
+		n := len(keys)
+
+		var pendingStart, pendingEnd uint64
+		hasPending := false
+
+		emit := func(rStart, rEnd uint64) bool {
+			if hasPending && rStart <= pendingEnd {
+				if rEnd > pendingEnd {
+					pendingEnd = rEnd
+				}
+				return true
+			}
+			if hasPending {
+				if !yield(uint32(pendingStart), pendingEnd) {
+					return false
+				}
+			}
+			pendingStart = rStart
+			pendingEnd = rEnd
+			hasPending = true
+			return true
+		}
+
+		for idx := 0; idx < n; idx++ {
+			hs := uint64(keys[idx]) << 16
+			c := containers[idx]
+
+			switch t := c.(type) {
+			case *runContainer16:
+				for _, iv := range t.iv {
+					if !emit(hs+uint64(iv.start), hs+uint64(iv.start)+uint64(iv.length)+1) {
+						return
+					}
+				}
+
+			case *bitmapContainer:
+				bm := t.bitmap
+				length := uint(len(bm))
+				pos := uint(0)
+
+				for pos < length {
+					if bm[pos] == 0 {
+						pos++
+						continue
+					}
+
+					w := bm[pos]
+					lo := uint(countTrailingZeros(w))
+					bitStart := pos*64 + lo
+
+					ones := uint(countTrailingOnes(w >> lo))
+					if lo+ones < 64 {
+						if !emit(hs|uint64(bitStart), hs|uint64(bitStart+ones)) {
+							return
+						}
+						pos = (bitStart + ones) / 64
+					} else {
+						pos++
+						for pos < length && bm[pos] == 0xFFFFFFFFFFFFFFFF {
+							pos++
+						}
+						var bitEnd uint
+						if pos < length {
+							bitEnd = pos*64 + uint(countTrailingOnes(bm[pos]))
+						} else {
+							bitEnd = length * 64
+						}
+						if !emit(hs|uint64(bitStart), hs|uint64(bitEnd)) {
+							return
+						}
+					}
+				}
+
+			case *arrayContainer:
+				content := t.content
+				i := 0
+				for i < len(content) {
+					start := uint64(content[i])
+					end := start + 1
+					i++
+					for i < len(content) && uint64(content[i]) == end {
+						end++
+						i++
+					}
+					if !emit(hs|start, hs|end) {
+						return
+					}
+				}
+			}
+		}
+
+		if hasPending {
+			yield(uint32(pendingStart), pendingEnd)
+		}
+	}
+}

iter_test.go

@@ -480,3 +480,42 @@ func TestUnsetIteratorPeekable(t *testing.T) {
 		assert.False(t, it.HasNext())
 	})
 }
+
+func TestRanges(t *testing.T) {
+	b := New()
+	b.AddRange(5, 10)
+	b.AddRange(20, 25)
+	b.AddRange(100, 105)
+	var ranges [][2]uint64
+	for start, end := range b.Ranges() {
+		ranges = append(ranges, [2]uint64{uint64(start), end})
+	}
+	assert.Equal(t, [][2]uint64{{5, 10}, {20, 25}, {100, 105}}, ranges)
+
+	// Merges across container boundaries
+	b2 := New()
+	b2.AddRange(0xFFF0, 0x10010)
+	ranges = nil
+	for start, end := range b2.Ranges() {
+		ranges = append(ranges, [2]uint64{uint64(start), end})
+	}
+	assert.Equal(t, [][2]uint64{{0xFFF0, 0x10010}}, ranges)
+
+	// Scattered values become individual ranges
+	b3 := New()
+	b3.Add(1)
+	b3.Add(3)
+	b3.Add(5)
+	ranges = nil
+	for start, end := range b3.Ranges() {
+		ranges = append(ranges, [2]uint64{uint64(start), end})
+	}
+	assert.Equal(t, [][2]uint64{{1, 2}, {3, 4}, {5, 6}}, ranges)
+
+	// Empty bitmap
+	count := 0
+	for range New().Ranges() {
+		count++
+	}
+	assert.Equal(t, 0, count)
+}