diff --git a/roaring/roaring.go b/roaring/roaring.go index e0810a6c2..86d810450 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) + n += popcount(c.bitmap[i] >> off) + i++ } // Count words in between. @@ -916,9 +924,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..41d0d167e --- /dev/null +++ b/roaring/roaring_internal_test.go @@ -0,0 +1,43 @@ +// Copyright 2017 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +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) + } + } +}