simplify bitmap range ops (#486)

* simplify bitmap range ops * tidy
livekit · Sep 30, 2023 · 43dd3e4 · 43dd3e4
1 parent 281e00a
commit 43dd3e4
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 22 deletions.
diff --git 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<<bits.Len64(uint64(size+63)/64)),
 	}
 }
 
@@ -48,8 +35,7 @@ func (b *Bitmap[T]) Set(val T) {
 }
 
 func (b *Bitmap[T]) SetRange(min, max T) {
-	if (max - min) > 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
@@ -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})
 }