diff --git a/roaring/roaring.go b/roaring/roaring.go index e74739311..0a91c4493 100644 --- a/roaring/roaring.go +++ b/roaring/roaring.go @@ -890,9 +890,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) + n += popcount(c.bitmap[i] >> off) + i++ } // Count words in between. @@ -902,9 +910,8 @@ func (c *container) bitmapCountRange(start, end uint32) int { // Count partial ending word. if int(j) < len(c.bitmap) { - if off := end % 64; off != 0 { - n += popcount(c.bitmap[j] >> off) - } + off := 64 - (end % 64) + n += popcount(c.bitmap[j] << off) } return int(n) diff --git a/roaring/roaring_internal_test.go b/roaring/roaring_internal_test.go new file mode 100644 index 000000000..5cff4a30f --- /dev/null +++ b/roaring/roaring_internal_test.go @@ -0,0 +1,29 @@ +package roaring + +import ( + "testing" +) + +func TestBitmapCountRange(t *testing.T) { + c := container{} + tests := []struct { + start uint32 + end uint32 + bitmap []uint64 + exp int + }{ + {start: 0, end: 1, bitmap: []uint64{1}, exp: 1}, + {start: 2, end: 7, bitmap: []uint64{0xFFFFFFFFFFFFFF18}, exp: 2}, + {start: 67, end: 68, bitmap: []uint64{0, 0x8}, exp: 1}, + {start: 1, end: 68, bitmap: []uint64{0x3, 0x8, 0xF}, exp: 2}, + {start: 1, end: 258, bitmap: []uint64{0xF, 0x8, 0xA, 0x4, 0xFFFFFFFFFFFFFFFF}, exp: 9}, + {start: 66, end: 71, bitmap: []uint64{0xF, 0xFFFFFFFFFFFFFF18}, exp: 2}, + {start: 63, end: 64, bitmap: []uint64{0x8000000000000000}, exp: 1}, + } + for i, test := range tests { + c.bitmap = test.bitmap + if ret := c.bitmapCountRange(test.start, test.end); ret != test.exp { + t.Fatalf("test #%v count of %v from %v to %v should be %v but got %v", i, test.bitmap, test.start, test.end, test.exp, ret) + } + } +}