diff --git a/bitmap.go b/bitmap.go index 3e11abf..d2276fb 100644 --- a/bitmap.go +++ b/bitmap.go @@ -123,18 +123,21 @@ func (dst Bitmap) CountTo(until uint32) int { return 0 } - // Figure out the index of the last block - blkUntil := int(until >> 6) - bitUntil := int(until % 64) - if blkUntil >= len(dst) { - blkUntil = len(dst) - 1 + if maxUntil := uint32(len(dst) << 6); until > maxUntil { + until = maxUntil } + // Figure out the index of the last block + blkUntil := until >> 6 + bitUntil := until % 64 + // Count the bits right before the last block sum := dst[:blkUntil].Count() // Count the bits at the end - sum += bits.OnesCount64(dst[blkUntil] << (64 - uint64(bitUntil))) + if bitUntil > 0 { + sum += bits.OnesCount64(dst[blkUntil] << (64 - uint64(bitUntil))) + } return sum } diff --git a/bitmap_test.go b/bitmap_test.go index 96a1524..14ab9db 100644 --- a/bitmap_test.go +++ b/bitmap_test.go @@ -392,7 +392,12 @@ func TestCount(t *testing.T) { assert.Equal(t, 4, b.CountTo(100)) assert.Equal(t, 4, b.CountTo(101)) assert.Equal(t, 5, b.CountTo(102)) + assert.Equal(t, 8, b.CountTo(128)) assert.Equal(t, 8, b.CountTo(math.MaxUint32)) + + b.Set(127) + + assert.Equal(t, 9, b.CountTo(128)) } func TestGrow(t *testing.T) {