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) + } +}