diff --git a/roaring/roaring.go b/roaring/roaring.go index cfcda97d1..16e6159d7 100644 --- a/roaring/roaring.go +++ b/roaring/roaring.go @@ -197,6 +197,7 @@ func (b *Bitmap) CountRange(start, end uint64) (n uint64) { if len(b.keys) == 0 { return } + skey := highbits(start) ekey := highbits(end) @@ -208,31 +209,28 @@ func (b *Bitmap) CountRange(start, end uint64) (n uint64) { return uint64(b.containers[i].countRange(int(lowbits(start)), int(lowbits(end)))) } - // Count first partial container. if i < 0 { - // start is before container, so we should start counting - // at first container that has value - if skey < b.keys[0] { - i = -1 - } else { - i = -i - } + // start's container did not exist + // set i to the index of the first container we have with values higher than start + i = -i - 1 } else { + // Count first partial container and advance i so we don't recount it n += uint64(b.containers[i].countRange(int(lowbits(start)), maxContainerVal+1)) + i += 1 } // Count last container. if j < 0 { - j = -j - if j > len(b.containers) { - j = len(b.containers) - } + // end's container did not exist + // set j to the index of the first container with values higher than end (or len(containers)) + j = -j - 1 } else { + // end's container exists, count it up to end n += uint64(b.containers[j].countRange(0, int(lowbits(end)))) } // Count containers in between. - for x := i + 1; x < j; x++ { + for x := i; x < j; x++ { n += uint64(b.containers[x].n) } @@ -2867,7 +2865,8 @@ func (*op) size() int { return 1 + 8 + 4 } func highbits(v uint64) uint64 { return uint64(v >> 16) } func lowbits(v uint64) uint16 { return uint16(v & 0xFFFF) } -// search32 returns the index of v in a. +// search32 returns the index of value in a. If value is not found, it works the +// same way as search64. func search32(a []uint16, value uint16) int { // Optimize for elements and the last element. n := len(a) @@ -2904,7 +2903,13 @@ func search32(a []uint16, value uint16) int { return -(lo + 1) } -// search64 returns the index of v in a. +// search64 returns the index of value in a. If value is not found, -1 * (1 + +// the index where v would be if it were inserted) is returned. This is done in +// order to both signal that value was not found (negative number), and also +// return information about where v would go if it were inserted. The +1 offset +// is necessary due to the case where v is not found, but would go at index 0. +// since negative 0 is no different from positive 0, we offset the returned +// negative indices by 1. See the test for this function for examples. func search64(a []uint64, value uint64) int { // Optimize for elements and the last element. n := len(a) diff --git a/roaring/roaring_internal_test.go b/roaring/roaring_internal_test.go index ea4e5f0ef..9cc2f2aa7 100644 --- a/roaring/roaring_internal_test.go +++ b/roaring/roaring_internal_test.go @@ -2310,3 +2310,81 @@ func Test_BufBitmapIterator_UnreadPanic(t *testing.T) { itr.unread() itr.unread() } + +func TestSearc64(t *testing.T) { + tests := []struct { + a []uint64 + value uint64 + exp int + }{ + { + a: []uint64{1, 5, 10, 12}, + value: 5, + exp: 1, + }, + { + a: []uint64{1, 5, 10, 12}, + value: 1, + exp: 0, + }, + { + a: []uint64{1, 5, 10, 12}, + value: 0, + exp: -1, + }, + { + a: []uint64{1, 5, 10, 12}, + value: 2, + exp: -2, + }, + { + a: []uint64{1, 5, 10, 12}, + value: 7, + exp: -3, + }, + { + a: []uint64{1, 5, 10, 12}, + value: 11, + exp: -4, + }, + { + a: []uint64{1, 5, 10, 12}, + value: 13, + exp: -5, + }, + { + a: []uint64{1, 5, 10, 12}, + value: 3843534, + exp: -5, + }, + { + a: []uint64{}, + value: 3843534, + exp: -1, + }, + { + a: []uint64{}, + value: 0, + exp: -1, + }, + { + a: []uint64{0}, + value: 0, + exp: 0, + }, + { + a: []uint64{0}, + value: 1, + exp: -2, + }, + } + + for _, test := range tests { + t.Run(fmt.Sprintf("%d in %v", test.value, test.a), func(t *testing.T) { + actual := search64(test.a, test.value) + if actual != test.exp { + t.Errorf("got: %d, exp: %d", actual, test.exp) + } + }) + } +} diff --git a/roaring/roaring_test.go b/roaring/roaring_test.go index c721f3804..7e524f51f 100644 --- a/roaring/roaring_test.go +++ b/roaring/roaring_test.go @@ -54,6 +54,104 @@ func TestContainerCount(t *testing.T) { } } +func TestCountRange(t *testing.T) { + tests := []struct { + name string + bitmap []uint64 + start uint64 + end uint64 + exp uint64 + }{ + { + name: "j < 0 : 1", + bitmap: []uint64{0, 1, 2, 3 * 65536}, + start: 0, + end: 65536, + exp: 3, + }, + { + name: "i < 0 : 1", + bitmap: []uint64{0, 1, 2, 2 * 65536, 3 * 65536}, + start: 65536, + end: 3 * 65536, + exp: 1, + }, + { + name: "single-container-run", + bitmap: []uint64{0, 2, 3, 4, 5, 2 * 65536, 3 * 65536}, + start: 2, + end: 5, + exp: 3, + }, + { + name: "single-container-beg", + bitmap: []uint64{1, 2, 3, 4, 5, 2 * 65536, 3 * 65536}, + start: 1, + end: 4, + exp: 3, + }, + { + name: "partial-start", + bitmap: []uint64{1, 2, 3, 4, 5, 2 * 65536, 3 * 65536}, + start: 5, + end: 3 * 65536, + exp: 2, + }, + { + name: "partial-end", + bitmap: []uint64{1, 2 * 65536, 3 * 65536, 3*65536 + 1, 3*65536 + 2}, + start: 0, + end: (3 * 65536) + 1, + exp: 3, + }, + { + name: "partial-both", + bitmap: []uint64{65536, 65537, 65538, 2 * 65536, 2*65536 + 1, 2*65536 + 2}, + start: 65537, + end: (2 * 65536) + 1, + exp: 3, + }, + { + name: "partial-both-bookends", + bitmap: []uint64{0, 65535, 65536, 65537, 65538, 2 * 65536, 2*65536 + 1, 2*65536 + 2, 3 * 65536}, + start: 65537, + end: (2 * 65536) + 1, + exp: 3, + }, + { + name: "empty-bookends", + bitmap: []uint64{1, 65535, 5 * 65536, 5*65536 + 1}, + start: 65536, + end: 5 * 65536, + exp: 0, + }, + { + name: "i not found, j found", + bitmap: []uint64{1, 65535, 5 * 65536}, + start: 2 * 65535, + end: 5*65536 + 1, + exp: 1, + }, + { + name: "i not found, j not found", + bitmap: []uint64{1, 65535, 5 * 65536, 7 * 65536}, + start: 2 * 65535, + end: 6 * 65536, + exp: 1, + }, + } + + for _, test := range tests { + t.Run(fmt.Sprintf("%s: %d to %d in '%v'", test.name, test.start, test.end, test.bitmap), func(t *testing.T) { + b := roaring.NewBitmap(test.bitmap...) + actual := b.CountRange(test.start, test.end) + if actual != test.exp { + t.Errorf("got: %d, exp: %d", actual, test.exp) + } + }) + } +} + func TestCheckBitmap(t *testing.T) { b := roaring.NewBitmap() x := 0