From 0facf126f040a29cee58f3d3f151b0cc1f29a6a1 Mon Sep 17 00:00:00 2001
From: Seebs <seebs@molecula.com>
Date: Tue, 14 Sep 2021 15:48:19 -0500
Subject: [PATCH] callback logic fixes for intersectionCallback and
 containerCallback

The inner loop of intersectionCallbackArrayArray's "fast"
case has
	for len(ca) > 0 && ca[0] < va {
	}
so we do not leave that loop unless len(ca) is 0, or
ca[0] >= va.

We then return from the whole function if len(ca) is 0,
so the only way we finish one iteration of the outer for
loop is if ca[0] >= va. Thus, this can be an `if` rather
than a `for`.

We also fix the logic for ArrayRun to make it require fewer
tests and be clearer about why the tests work and clearer about
always making progress.

And, finally, the bitmap/range callback logic, and the underlying
"callback per bit in word" logic, were both badly broken. In
particular, if a range started and ended in the same word, it would
hit the values in that word twice, once with them incorrectly
shifted, but then it would further garble any offsets past the first
in a word. Eww.
---
 roaring/roaring.go               | 51 +++++++++++++++++++++++---------
 roaring/roaring_internal_test.go | 12 ++++++++
 2 files changed, 49 insertions(+), 14 deletions(-)

diff --git a/roaring/roaring.go b/roaring/roaring.go
index 32d573911..07751d18e 100644
--- a/roaring/roaring.go
+++ b/roaring/roaring.go
@@ -2952,6 +2952,7 @@ func ArrayCountRange(array []uint16, start, end int32) (n int32) {
 	return n
 }
 
+// BitmapCountRange counts bits set in [start,end).
 func BitmapCountRange(bitmap []uint64, start, end int32) int32 {
 	if roaringParanoia {
 		if start > end {
@@ -2988,15 +2989,16 @@ func BitmapCountRange(bitmap []uint64, start, end int32) int32 {
 }
 
 func callbackBits(w uint64, base uint16, fn func(uint16)) {
-	bit := uint16(0)
 	for w != 0 {
-		trail := bits.TrailingZeros64(w)
-		bit += uint16(trail)
-		w >>= (trail + 1)
-		fn(base + bit)
+		trail := uint16(bits.TrailingZeros64(w))
+		fn(base + trail)
+		base += trail + 1
+		w >>= trail + 1
 	}
 }
 
+// bitmapCallbackRange calls the provided function for every bit set in
+// bitmap in the range [start,end).
 func bitmapCallbackRange(bitmap []uint64, start, end int32, fn func(uint16)) {
 	if roaringParanoia {
 		if start > end {
@@ -3006,11 +3008,30 @@ func bitmapCallbackRange(bitmap []uint64, start, end int32, fn func(uint16)) {
 	i, j := start/64, end/64
 	// Special case when start and end fall in the same word.
 	if i == j {
-		offi, offj := uint(start%64), uint(64-end%64)
-		w := (bitmap[i] >> offi) << (offj + offi)
+		// So, we want to know the offsets. For instance, if start and end
+		// are 65 and 69, we might want i=1, offi=1, j=1, offj=5. Then we
+		// compute masks from offi (masking out 0x1, or (1<<offi)-1), and
+		// offj (masking out everything *but* (1<<offj)-1).
+		//
+		// But we can also use "x >> offi << offi" to trim the lowest offi
+		// bits, and "x << (64-offj) >> (64-offj)" to trim all but the
+		// lowest offj bits.
+		//
+		// We can then simplify slightly further: we use the inverted value
+		// as offj, and compute (w << offi) >> (offi + offj) << offi.
+		//
+		// But wait, you ask. What if offi+offj is too large! Well, then
+		// start and end were in the wrong order. We have 0 <= i <= j < 64.
+		// If x+i > 64, then x > (64-i). Thus, if (64-j)+i > 64, it
+		// follows that (64-j) > (64-i). So they'd have been in the wrong order.
+		// In which case, we correctly yield a value of (0 << offi), or 0,
+		// because nothing is between them.
+		offi, offj := uint(start%64), uint(64-(end%64))
+		w := (bitmap[i] << offj) >> (offi + offj) << offi
 		if w != 0 {
 			callbackBits(w, uint16(i)*64, fn)
 		}
+		return
 	}
 
 	// Count partial starting word.
@@ -4447,7 +4468,7 @@ func intersectionCallbackArrayArray(a, b *Container, fn func(uint16)) {
 	}
 	if (na << 2) < nb {
 		for _, va := range ca {
-			for cb[0] < va {
+			if cb[0] < va {
 				// try to skip ahead a bit faster
 				for len(cb) > 7 && cb[7] < va {
 					cb = cb[8:]
@@ -4485,13 +4506,15 @@ func intersectionCallbackArrayRun(a, b *Container, fn func(uint16)) {
 	na, nb := len(array), len(runs)
 	for i, j := 0, 0; i < na && j < nb; {
 		va, vb := array[i], runs[j]
-		if va < vb.Start {
-			i++
-		} else if va >= vb.Start && va <= vb.Last {
-			i++
-			fn(va)
-		} else if va > vb.Last {
+		if va > vb.Last {
 			j++
+			continue
+		}
+		// If we got here, va is either before or in the current run,
+		// so we're definitely done with this member of the array.
+		i++
+		if va >= vb.Start {
+			fn(va)
 		}
 	}
 }
diff --git a/roaring/roaring_internal_test.go b/roaring/roaring_internal_test.go
index 9902c4373..ecc2b6d97 100644
--- a/roaring/roaring_internal_test.go
+++ b/roaring/roaring_internal_test.go
@@ -4636,3 +4636,15 @@ func TestContainer_unionInPlace_ArrayUnionRun(t *testing.T) {
 		}
 	}
 }
+
+func TestIntersectionCallback(t *testing.T) {
+	var hits []uint16
+	cb := func(u uint16) {
+		hits = append(hits, u)
+	}
+	bm := []uint64{0, 5, 0}
+	bitmapCallbackRange(bm, 64, 69, cb)
+	if len(hits) != 2 || hits[0] != 64 || hits[1] != 66 {
+		t.Fatalf("expected 64, 66, got %d", hits)
+	}
+}