From fb4651cdb8092e1dfad3b0384b57d306ce351282 Mon Sep 17 00:00:00 2001
From: Matt Jaffee <jaffee@umbel.com>
Date: Thu, 18 May 2017 12:04:37 -0500
Subject: [PATCH] fix 3 separate bugs in bitmapCountRange

in order of the diff:

1. When the start and end of the range fall in the same word, special handling
is needed to "mask" off the beginning and end of the word simultaneously to
avoid counting bits at the beginning or end of the word that aren't in the
range.

2. `i++` is needed at the end of the first partial word to avoid counting this
word in the next block.

3. the shift amount for right shifts is 64 - (end % 64) rather than just end
% 64. If end is (e.g.) 68, then 68 - 64 is 4 and we are only interested in the
first 4 bits of the word, so we must right shift by 60 bits, not 4 bits.
---
 roaring/roaring.go               | 10 +++++++++-
 roaring/roaring_internal_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 roaring/roaring_internal_test.go

diff --git a/roaring/roaring.go b/roaring/roaring.go
index e0810a6c2..1a4d8d486 100644
--- a/roaring/roaring.go
+++ b/roaring/roaring.go
@@ -904,9 +904,17 @@ func (c *container) bitmapCountRange(start, end uint32) int {
 	var n uint64
 	i, j := start/64, end/64
 
+	// Special case when start and end fall in the same word.
+	if i == j {
+		offi, offj := start%64, 64-end%64
+		n += popcount((c.bitmap[i] << offi) >> (offj + offi))
+		return int(n)
+	}
+
 	// Count partial starting word.
 	if off := start % 64; off != 0 {
 		n += popcount(c.bitmap[i] << off)
+		i++
 	}
 
 	// Count words in between.
@@ -916,7 +924,7 @@ func (c *container) bitmapCountRange(start, end uint32) int {
 
 	// Count partial ending word.
 	if int(j) < len(c.bitmap) {
-		if off := end % 64; off != 0 {
+		if off := 64 - (end % 64); off != 0 {
 			n += popcount(c.bitmap[j] >> off)
 		}
 	}
diff --git a/roaring/roaring_internal_test.go b/roaring/roaring_internal_test.go
new file mode 100644
index 000000000..17a35c5e1
--- /dev/null
+++ b/roaring/roaring_internal_test.go
@@ -0,0 +1,31 @@
+package roaring
+
+import (
+	"testing"
+)
+
+func TestBitmapCountRange(t *testing.T) {
+	c := container{bitmap: []uint64{1}}
+	cnt := c.bitmapCountRange(63, 65)
+	if cnt != 1 {
+		t.Fatalf("count of %v from 63 to 65 should be 1, but got %v", c.bitmap, cnt)
+	}
+
+	c = container{bitmap: []uint64{0, 0x8000000000000000}}
+	cnt = c.bitmapCountRange(65, 66)
+	if cnt != 0 {
+		t.Fatalf("count of %v from 65 to 66 should be 0, but got %v", c.bitmap, cnt)
+	}
+
+	c = container{bitmap: []uint64{0, 0xF000000000000000}}
+	cnt = c.bitmapCountRange(65, 66)
+	if cnt != 1 {
+		t.Fatalf("count of %v from 65 to 66 should be 1, but got %v", c.bitmap, cnt)
+	}
+
+	c = container{bitmap: []uint64{0x1, 0xFF00000000000000}}
+	cnt = c.bitmapCountRange(62, 66)
+	if cnt != 3 {
+		t.Fatalf("count of %v from 62 to 66 should be 3, but got %v", c.bitmap, cnt)
+	}
+}