Add support for halting the iteration

arraycontainer.go

@@ -24,10 +24,14 @@ func (ac *arrayContainer) fillLeastSignificant16bits(x []uint32, i int, mask uin
 	}
 }
 
-func (ac *arrayContainer) iterate(cb func(x uint16)) {
+func (ac *arrayContainer) iterate(cb func(x uint16) bool) bool {
 	for i := 0; i < len(ac.content); i++ {
-		cb(ac.content[i])
+		if !cb(ac.content[i]) {
+			return false
+		}
 	}
+
+	return true
 }
 
 func (ac *arrayContainer) getShortIterator() shortPeekable {

benchmark_test.go

@@ -419,8 +419,9 @@ func BenchmarkIterateRoaring(b *testing.B) {
 
 		for j := 0; j < b.N; j++ {
 			c9 = uint(0)
-			s.Iterate(func(x uint32) {
+			s.Iterate(func(x uint32) bool {
 				c9++
+				return true
 			})
 		}
 	})
@@ -434,8 +435,9 @@ func BenchmarkIterateRoaring(b *testing.B) {
 
 		for j := 0; j < b.N; j++ {
 			c9 = uint(0)
-			s.Iterate(func(x uint32) {
+			s.Iterate(func(x uint32) bool {
 				c9++
+				return true
 			})
 		}
 	})

bitmapcontainer.go

@@ -96,12 +96,15 @@ func (bc *bitmapContainer) maximum() uint16 {
 	return uint16(0)
 }
 
-func (bc *bitmapContainer) iterate(cb func(x uint16)) {
+func (bc *bitmapContainer) iterate(cb func(x uint16) bool) bool {
 	for i := bc.NextSetBit(0); i >= 0; {
 		j := i
 		i = bc.NextSetBit(i + 1)
-		cb(uint16(j))
+		if !cb(uint16(j)) {
+			return false
+		}
 	}
+	return true
 }
 
 type bitmapContainerShortIterator struct {

roaring.go

@@ -416,31 +416,36 @@ func (rb *Bitmap) String() string {
 	return buffer.String()
 }
 
-// Iterate iterates over the bitmap, calling the given callback with each value in the bitmap.
+// Iterate iterates over the bitmap, calling the given callback with each value in the bitmap.  If the callback returns
+// false, the iteration is halted.
 // The iteration results are undefined if the bitmap is modified (e.g., with Add or Remove).
 // There is no guarantee as to what order the values will be iterated
-func (rb *Bitmap) Iterate(cb func(x uint32)) {
+func (rb *Bitmap) Iterate(cb func(x uint32) bool) {
 	var hs uint32
 	for i := 0; i < rb.highlowcontainer.size(); i++ {
 		hs = uint32(rb.highlowcontainer.getKeyAtIndex(i)) << 16
 
 		c := rb.highlowcontainer.getContainerAtIndex(i)
 
+		var shouldContinue bool
 		// This is hacky but it avoids allocations from invoking an interface method with a closure
 		switch t := c.(type) {
 		case *arrayContainer:
-			t.iterate(func(x uint16) {
-				cb(uint32(x) | hs)
+			shouldContinue = t.iterate(func(x uint16) bool {
+				return cb(uint32(x) | hs)
 			})
 		case *runContainer16:
-			t.iterate(func(x uint16) {
-				cb(uint32(x) | hs)
+			shouldContinue = t.iterate(func(x uint16) bool {
+				return cb(uint32(x) | hs)
 			})
 		case *bitmapContainer:
-			t.iterate(func(x uint16) {
-				cb(uint32(x) | hs)
+			shouldContinue = t.iterate(func(x uint16) bool {
+				return cb(uint32(x) | hs)
 			})
 		}
+		if !shouldContinue {
+			return
+		}
 	}
 }
 

roaring_test.go

@@ -2317,8 +2317,9 @@ func TestIterate(t *testing.T) {
 	}
 
 	var values []uint32
-	rb.Iterate(func(x uint32) {
+	rb.Iterate(func(x uint32) bool {
 		values = append(values, x)
+		return true
 	})
 
 	assert.Equal(t, rb.ToArray(), values)
@@ -2334,8 +2335,9 @@ func TestIterateCompressed(t *testing.T) {
 	rb.RunOptimize()
 
 	var values []uint32
-	rb.Iterate(func(x uint32) {
+	rb.Iterate(func(x uint32) bool {
 		values = append(values, x)
+		return true
 	})
 
 	assert.Equal(t, rb.ToArray(), values)
@@ -2350,9 +2352,35 @@ func TestIterateLargeValues(t *testing.T) {
 	}
 
 	var values []uint32
-	rb.Iterate(func(x uint32) {
+	rb.Iterate(func(x uint32) bool {
 		values = append(values, x)
+		return true
 	})
 
 	assert.Equal(t, rb.ToArray(), values)
 }
+
+func TestIterateHalt(t *testing.T) {
+	rb := NewBitmap()
+
+	// This range of values ensures that all different types of containers will be used
+	for i := 150000; i < 450000; i++ {
+		rb.Add(uint32(i))
+	}
+
+	var values []uint32
+	count := uint64(0)
+	stopAt := rb.GetCardinality() - 1
+	rb.Iterate(func(x uint32) bool {
+		values = append(values, x)
+		count++
+		if count == stopAt {
+			return false
+		}
+		return true
+	})
+
+	expected := rb.ToArray()
+	expected = expected[0 : len(expected)-1]
+	assert.Equal(t, expected, values)
+}

roaringarray.go

@@ -39,7 +39,7 @@ type container interface {
 	inot(firstOfRange, endx int) container // i stands for inplace, range is [firstOfRange,endx)
 	xor(r container) container
 	getShortIterator() shortPeekable
-	iterate(cb func(x uint16))
+	iterate(cb func(x uint16) bool) bool
 	getReverseIterator() shortIterable
 	getManyIterator() manyIterable
 	contains(i uint16) bool

runcontainer.go

@@ -1162,7 +1162,7 @@ func (rc *runContainer16) newRunIterator16() *runIterator16 {
 	return &runIterator16{rc: rc, curIndex: 0, curPosInIndex: 0}
 }
 
-func (rc *runContainer16) iterate(cb func(x uint16)) {
+func (rc *runContainer16) iterate(cb func(x uint16) bool) bool {
 	curIndex := int64(0)
 	curPosInIndex := uint16(0)
 
@@ -1181,9 +1181,12 @@ func (rc *runContainer16) iterate(cb func(x uint16)) {
 			curPosInIndex++
 		}
 
-		cb(next)
+		if !cb(next) {
+			return false
+		}
 	}
 
+	return true
 }
 
 // hasNext returns false if calling next will panic. It