From 39ed7de5c44f20d9f4634eb220ca2e4b6194ca10 Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Wed, 7 Sep 2016 10:13:31 -0600
Subject: [PATCH] Optimize IntersectionCount for Array+Bitmap

Reworks the `roaring.intersectionCountArrayBitmap()` call to avoid
using an iterator. Performance of the included benchmark went from
2.5ms to 1.1ms.

Some of the issue with intersectionCount is the increased size of
bitmaps and slices and I need to do additional testing with various
sizes.
---
 roaring/roaring.go      | 58 +++++++++++++++++++++++++++--------------
 roaring/roaring_test.go | 29 +++++++++++++++++++++
 2 files changed, 68 insertions(+), 19 deletions(-)

diff --git a/roaring/roaring.go b/roaring/roaring.go
index cdf6dfe3a..146908157 100644
--- a/roaring/roaring.go
+++ b/roaring/roaring.go
@@ -475,7 +475,7 @@ func (b *Bitmap) UnmarshalBinary(data []byte) error {
 
 		// Map byte slice directly to the container data.
 		c := b.containers[i]
-		if c.n <= arrayMaxSize {
+		if c.n <= ArrayMaxSize {
 			c.array = (*[0xFFFFFFF]uint32)(unsafe.Pointer(&data[offset]))[:c.n]
 			opsOffset = int(offset) + len(c.array)*4
 		} else {
@@ -704,7 +704,7 @@ func (itr *BufIterator) Unread() {
 }
 
 // The maximum size of array containers.
-const arrayMaxSize = (1 << 20)
+const ArrayMaxSize = (1 << 20)
 
 // container represents a container for uint32 integers.
 //
@@ -803,7 +803,7 @@ func (c *container) add(v uint32) bool {
 
 func (c *container) arrayAdd(v uint32) bool {
 	// Optimize appending to the end of an array container.
-	if c.n > 0 && c.n < arrayMaxSize && c.isArray() && c.array[c.n-1] < v {
+	if c.n > 0 && c.n < ArrayMaxSize && c.isArray() && c.array[c.n-1] < v {
 		c.unmap()
 		c.array = append(c.array, v)
 		c.n++
@@ -817,7 +817,7 @@ func (c *container) arrayAdd(v uint32) bool {
 	}
 
 	// Convert to a bitmap container if too many values are in an array container.
-	if c.n >= arrayMaxSize {
+	if c.n >= ArrayMaxSize {
 		c.convertToBitmap()
 		return c.bitmapAdd(v)
 	}
@@ -889,7 +889,7 @@ func (c *container) bitmapRemove(v uint32) bool {
 	c.bitmap[v/64] &^= (uint64(1) << (v % 64))
 
 	// Convert to array if we go below the threshold.
-	if c.n == arrayMaxSize {
+	if c.n == ArrayMaxSize {
 		c.convertToArray()
 	}
 	return true
@@ -1061,24 +1061,44 @@ func intersectionCountArrayArray(a, b *container) (n uint64) {
 }
 
 func intersectionCountArrayBitmap(a, b *container) (n uint64) {
-	itr := newBufIterator(newBitmapIterator(b.bitmap))
-	for i := 0; i < len(a.array); {
-		va := a.array[i]
-		vb, eof := itr.next()
-		if eof {
-			break
+	// Copy array header so we can shrink it.
+	array := a.array
+	if len(array) == 0 {
+		return 0
+	}
+
+	// Iterate over bitmap and find matching bits.
+	for i, bn := uint32(0), uint32(len(b.bitmap)); i < bn; i++ {
+		v := b.bitmap[i]
+
+		// Ignore if bytes are empty or array is done.
+		if v == 0 {
+			continue
 		}
 
-		if va < vb {
-			i++
-			itr.unread()
-		} else if va > vb {
-			// nop
-		} else {
-			n++
-			i++
+		// Check each bit.
+		for j := uint32(0); j < 64; j++ {
+			if v&(1<<j) == 0 {
+				continue
+			}
+
+			// Search array until match.
+			bv := (i * 64) + j
+			for {
+				if len(array) == 0 {
+					return n
+				} else if array[0] < bv {
+					array = array[1:]
+				} else if array[0] == bv {
+					n++
+					break
+				} else {
+					break
+				}
+			}
 		}
 	}
+
 	return n
 }
 
diff --git a/roaring/roaring_test.go b/roaring/roaring_test.go
index a4cbd08ff..cb03eb7d3 100644
--- a/roaring/roaring_test.go
+++ b/roaring/roaring_test.go
@@ -274,6 +274,35 @@ func TestIterator(t *testing.T) {
 	}
 }
 
+var benchmarkBitmapIntersectionCountData struct {
+	a, b *roaring.Bitmap
+}
+
+func BenchmarkBitmap_IntersectionCount_ArrayBitmap(b *testing.B) {
+	data := &benchmarkBitmapIntersectionCountData
+	if data.a == nil {
+		const max = (1 << 24) / 64
+
+		// Build bitmap with array container.
+		data.a = roaring.NewBitmap()
+		for i, n := 0, rand.Intn(roaring.ArrayMaxSize); i < n; i++ {
+			data.a.Add(uint64(rand.Intn(max)))
+		}
+
+		// Build bitmap with bitmap container.
+		data.b = roaring.NewBitmap()
+		for i, n := 0, roaring.ArrayMaxSize*2; i < n; i++ {
+			data.b.Add(uint64(i * 3))
+		}
+	}
+
+	// Reset timer & benchmark.
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		data.a.IntersectionCount(data.b)
+	}
+}
+
 // GenerateUint64Slice generates between [0, n) random uint64 numbers between min and max.
 func GenerateUint64Slice(n int, min, max uint64, sorted bool, rand *rand.Rand) []uint64 {
 	a := make([]uint64, rand.Intn(n))