Fix NextUnsetBit returning set bits in bitmap containers

bitmapcontainer.go

@@ -1198,9 +1198,7 @@ func (bc *bitmapContainer) NextUnsetBit(i uint) int {
 	if x >= length {
 		return int(i)
 	}
-	w := bc.bitmap[x]
-	w = w >> (i % 64)
-	w = ^w
+	w := (^bc.bitmap[x]) >> (i % 64)
 	if w != 0 {
 		return int(i) + bits.TrailingZeros64(w)
 	}

bitmapcontainer_test.go

@@ -441,6 +441,29 @@ func TestBitmapContainerAbsentValueAcrossWords(t *testing.T) {
 	})
 }
 
+func TestBitmapContainerNextUnsetBitAcrossWords(t *testing.T) {
+	t.Run("Next unset bit after a run of ones that reaches the end of a word", func(t *testing.T) {
+		bc := newBitmapContainer()
+		for i := 0; i < 5000; i++ {
+			bc.iadd(uint16(i))
+		}
+		assert.Equal(t, 5000, bc.NextUnsetBit(0))
+		assert.Equal(t, 5000, bc.NextUnsetBit(10))
+		assert.Equal(t, 5000, bc.NextUnsetBit(70))
+		assert.Equal(t, 5001, bc.NextUnsetBit(5001))
+	})
+
+	t.Run("Next unset bit of a full container", func(t *testing.T) {
+		bc := newBitmapContainer()
+		for i := 0; i < 1<<16; i++ {
+			bc.iadd(uint16(i))
+		}
+		assert.Equal(t, 1<<16, bc.NextUnsetBit(0))
+		assert.Equal(t, 1<<16, bc.NextUnsetBit(64))
+		assert.Equal(t, 1<<16, bc.NextUnsetBit(65535))
+	})
+}
+
 func TestBitMapContainerValidate(t *testing.T) {
 	bc := newBitmapContainer()
 

iter_test.go

@@ -1138,3 +1138,51 @@ func TestRangesEarlyTermination(t *testing.T) {
 		breakAfter(t, b, 2)
 	})
 }
+
+// Issue #592: the unset iterator of a bitmap container yielded values that
+// were present in the bitmap when the rest of a word was all ones.
+func TestUnsetBitmapContainer(t *testing.T) {
+	d := NewBitmap()
+	d.AddMany([]uint32{0})
+	d.AddRange(2, 4098)
+	_, isBitmap := d.highlowcontainer.containers[0].(*bitmapContainer)
+	assert.True(t, isBitmap)
+
+	var got []uint32
+	for v := range Unset(d, 2, 4099) {
+		got = append(got, v)
+	}
+	assert.Equal(t, []uint32{4098, 4099}, got)
+
+	b := NewBitmap()
+	values := make([]uint32, 0, 65534-100)
+	for v := uint32(100); v < 65534; v++ {
+		values = append(values, v)
+	}
+	b.AddMany(values)
+	b.Remove(6011)
+	b.Remove(23000)
+	_, isBitmap = b.highlowcontainer.containers[0].(*bitmapContainer)
+	assert.True(t, isBitmap)
+
+	for _, r := range [][2]uint32{{0, 65535}, {50, 200}, {6000, 30000}, {23814, 65535}} {
+		var expected []uint32
+		for v := r[0]; v <= r[1]; v++ {
+			if !b.Contains(v) {
+				expected = append(expected, v)
+			}
+		}
+		got = got[:0]
+		for v := range Unset(b, r[0], r[1]) {
+			got = append(got, v)
+		}
+		assert.Equal(t, expected, got, "range [%d, %d]", r[0], r[1])
+	}
+
+	// AdvanceIfNeeded goes through NextUnsetBit as well.
+	it := b.UnsetIterator(0, 65536)
+	it.AdvanceIfNeeded(6001)
+	assert.Equal(t, uint32(6011), it.Next())
+	assert.Equal(t, uint32(23000), it.Next())
+	assert.Equal(t, uint32(65534), it.Next())
+}