diff --git a/utils/bitmap.go b/utils/bitmap.go index 4ad206a7..3ab5e37d 100644 --- a/utils/bitmap.go +++ b/utils/bitmap.go @@ -14,32 +14,19 @@ package utils -import ( - "unsafe" -) +import "math/bits" type bitmapNumber interface { uint16 | uint32 | uint64 } type Bitmap[T bitmapNumber] struct { - halfRange T - bits []uint64 + bits []uint64 } func NewBitmap[T bitmapNumber](size int) *Bitmap[T] { - pof2 := 1 - for { - if pof2 >= 64 && pof2 > size { - break - } - pof2 <<= 1 - } - - var t T return &Bitmap[T]{ - halfRange: 1 << (unsafe.Sizeof(t)*8 - 1), - bits: make([]uint64, pof2/64), + bits: make([]uint64, 1< b.halfRange { - // out-of-order + if max < min { return } @@ -70,8 +56,7 @@ func (b *Bitmap[T]) Clear(val T) { } func (b *Bitmap[T]) ClearRange(min, max T) { - if (max - min) > b.halfRange { - // out-of-order + if max < min { return } @@ -97,8 +82,14 @@ func (b *Bitmap[T]) IsSet(val T) bool { func (b *Bitmap[T]) getSlotsAndOffsets(min, max T) (sm int, ls int, rs int, lo int, ro int) { sm = len(b.bits) - 1 // slot mask - ls = int(min >> 6) // left slot - rs = int(max >> 6) // right slot + ls = int(min >> 6) // left slot + rs = int(max >> 6) // right slot + + if rs-ls > len(b.bits) { + rs = ls + len(b.bits) + return + } + lo = int(min & 0x3f) // left offset ro = int(max & 0x3f) // right offset return diff --git a/utils/bitmap_test.go b/utils/bitmap_test.go index 9abadc98..da351bb2 100644 --- a/utils/bitmap_test.go +++ b/utils/bitmap_test.go @@ -15,6 +15,7 @@ package utils import ( + "math" "testing" "github.com/stretchr/testify/require" @@ -75,4 +76,24 @@ func TestBitmap(t *testing.T) { require.False(t, b.IsSet(95)) require.False(t, b.IsSet(234)) require.True(t, b.IsSet(235)) + + // set large range + b.SetRange(0, 1000) + e[0] = 0xFFFF_FFFF_FFFF_FFFF + e[1] = 0xFFFF_FFFF_FFFF_FFFF + e[2] = 0xFFFF_FFFF_FFFF_FFFF + e[3] = 0xFFFF_FFFF_FFFF_FFFF + require.Equal(t, e, b.bits) + + // clear large range + b.ClearRange(0, 1000) + e[0] = 0x0000_0000_0000_0000 + e[1] = 0x0000_0000_0000_0000 + e[2] = 0x0000_0000_0000_0000 + e[3] = 0x0000_0000_0000_0000 + require.Equal(t, e, b.bits) + + // large range changes touch each word once + sm, ls, rs, lo, ro := b.getSlotsAndOffsets(0, math.MaxUint32) + require.Equal(t, []int{3, 0, 4, 0, 0}, []int{sm, ls, rs, lo, ro}) }