more efficient UnsetIterator implementation

Also add some property-based tests with a small corpus
of values.

Signed-off-by: Roger Peppe <[email protected]>

Author: rogpeppe

arraycontainer.go

@@ -62,6 +62,58 @@ func (ac *arrayContainer) getManyIterator() manyIterable {
 	return &shortIterator{ac.content, 0}
 }
 
+type arrayContainerUnsetIterator struct {
+	content []uint16
+	pos     int
+	nextVal int
+}
+
+func (acui *arrayContainerUnsetIterator) next() uint16 {
+	val := acui.nextVal
+	acui.nextVal++
+	for acui.pos < len(acui.content) && uint16(acui.nextVal) == acui.content[acui.pos] {
+		acui.nextVal++
+		acui.pos++
+	}
+	return uint16(val)
+}
+
+func (acui *arrayContainerUnsetIterator) hasNext() bool {
+	return acui.nextVal < 65536
+}
+
+func (acui *arrayContainerUnsetIterator) peekNext() uint16 {
+	return uint16(acui.nextVal)
+}
+
+func (acui *arrayContainerUnsetIterator) advanceIfNeeded(minval uint16) {
+	if !acui.hasNext() || acui.peekNext() >= minval {
+		return
+	}
+	acui.nextVal = int(minval)
+	acui.pos = binarySearch(acui.content, minval)
+	if acui.pos < 0 {
+		acui.pos = -acui.pos - 1
+	}
+	for acui.pos < len(acui.content) && uint16(acui.nextVal) == acui.content[acui.pos] {
+		acui.nextVal++
+		acui.pos++
+	}
+}
+
+func newArrayContainerUnsetIterator(a *arrayContainer) *arrayContainerUnsetIterator {
+	acui := &arrayContainerUnsetIterator{content: a.content, pos: 0, nextVal: 0}
+	for acui.pos < len(acui.content) && uint16(acui.nextVal) == acui.content[acui.pos] {
+		acui.nextVal++
+		acui.pos++
+	}
+	return acui
+}
+
+func (ac *arrayContainer) getUnsetIterator() shortPeekable {
+	return newArrayContainerUnsetIterator(ac)
+}
+
 func (ac *arrayContainer) minimum() uint16 {
 	return ac.content[0] // assume not empty
 }

benchmark_test.go

@@ -600,6 +600,22 @@ func BenchmarkIterateRoaring(b *testing.B) {
 			})
 		}
 	})
+	b.Run("unsetIterator", func(b *testing.B) {
+		b.ReportAllocs()
+
+		s := Flip(newBitmap(), 0, 0x100000000)
+
+		b.ResetTimer()
+
+		for j := 0; j < b.N; j++ {
+			c9 = uint(0)
+			i := s.UnsetIterator(0, 0xffffffff)
+			for i.HasNext() {
+				i.Next()
+				c9++
+			}
+		}
+	})
 }
 
 // go test -bench BenchmarkSparseIterate -run -

bitmapcontainer.go

@@ -262,6 +262,39 @@ func (bc *bitmapContainer) getManyIterator() manyIterable {
 	return newBitmapContainerManyIterator(bc)
 }
 
+type bitmapContainerUnsetIterator struct {
+	ptr *bitmapContainer
+	i   int
+}
+
+func (bcui *bitmapContainerUnsetIterator) next() uint16 {
+	j := bcui.i
+	bcui.i = bcui.ptr.NextUnsetBit(uint(bcui.i) + 1)
+	return uint16(j)
+}
+
+func (bcui *bitmapContainerUnsetIterator) hasNext() bool {
+	return bcui.i >= 0 && bcui.i < 65536
+}
+
+func (bcui *bitmapContainerUnsetIterator) peekNext() uint16 {
+	return uint16(bcui.i)
+}
+
+func (bcui *bitmapContainerUnsetIterator) advanceIfNeeded(minval uint16) {
+	if bcui.hasNext() && bcui.peekNext() < minval {
+		bcui.i = bcui.ptr.NextUnsetBit(uint(minval))
+	}
+}
+
+func newBitmapContainerUnsetIterator(a *bitmapContainer) *bitmapContainerUnsetIterator {
+	return &bitmapContainerUnsetIterator{a, a.NextUnsetBit(0)}
+}
+
+func (bc *bitmapContainer) getUnsetIterator() shortPeekable {
+	return newBitmapContainerUnsetIterator(bc)
+}
+
 func (bc *bitmapContainer) getSizeInBytes() int {
 	return len(bc.bitmap) * 8
 }
@@ -1113,6 +1146,29 @@ func (bc *bitmapContainer) NextSetBit(i uint) int {
 	return -1
 }
 
+func (bc *bitmapContainer) NextUnsetBit(i uint) int {
+	var (
+		x      = i / 64
+		length = uint(len(bc.bitmap))
+	)
+	if x >= length {
+		return int(i)
+	}
+	w := bc.bitmap[x]
+	w = w >> uint(i%64)
+	w = ^w
+	if w != 0 {
+		return int(i) + countTrailingZeros(w)
+	}
+	x++
+	for ; x < length; x++ {
+		if bc.bitmap[x] != 0xFFFFFFFFFFFFFFFF {
+			return int(x*64) + countTrailingZeros(^bc.bitmap[x])
+		}
+	}
+	return int(length * 64)
+}
+
 // PrevSetBit returns the previous set bit e.g the previous int packed into the bitmaparray
 func (bc *bitmapContainer) PrevSetBit(i int) int {
 	if i < 0 {