bit map

utils/bitmap.go

@@ -0,0 +1,105 @@
+// Copyright 2023 LiveKit, Inc.
+//
+// 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 utils
+
+import (
+	"unsafe"
+)
+
+type bitmapNumber interface {
+	uint16 | uint32 | uint64
+}
+
+type Bitmap[T bitmapNumber] struct {
+	halfRange T
+	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),
+	}
+}
+
+func (b *Bitmap[T]) Set(val T) {
+	b.SetRange(val, val)
+}
+
+func (b *Bitmap[T]) SetRange(min, max T) {
+	if (max - min) > b.halfRange {
+		// out-of-order
+		return
+	}
+
+	sm, ls, rs, lo, ro := b.getSlotsAndOffsets(min, max)
+	if ls == rs {
+		b.bits[ls&sm] |= (((1 << (ro - lo + 1)) - 1) << lo)
+	} else {
+		b.bits[ls&sm] |= ^((1 << lo) - 1)
+		for i := ls + 1; i < rs; i++ {
+			b.bits[i&sm] = ^uint64(0)
+		}
+		b.bits[rs&sm] |= (1 << (ro + 1)) - 1
+	}
+}
+
+func (b *Bitmap[T]) Clear(val T) {
+	b.ClearRange(val, val)
+}
+
+func (b *Bitmap[T]) ClearRange(min, max T) {
+	if (max - min) > b.halfRange {
+		// out-of-order
+		return
+	}
+
+	sm, ls, rs, lo, ro := b.getSlotsAndOffsets(min, max)
+	if ls == rs {
+		b.bits[ls&sm] &= ^(((1 << (ro - lo + 1)) - 1) << lo)
+	} else {
+		b.bits[ls&sm] &= ^uint64(0) >> (64 - lo)
+		for i := ls + 1; i < rs; i++ {
+			b.bits[i&sm] = 0
+		}
+		b.bits[rs&sm] &= ^uint64(0) << (ro + 1)
+	}
+}
+
+func (b *Bitmap[T]) IsSet(val T) bool {
+	sm := len(b.bits) - 1 // slot mask
+	s := int(val >> 6)    // slot
+	o := int(val & 0x3f)  // offset
+	return b.bits[s&sm]&(1<<o) != 0
+}
+
+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
+	lo = int(min & 0x3f) // left offset
+	ro = int(max & 0x3f) // right offset
+	return
+}

utils/bitmap_test.go

@@ -0,0 +1,78 @@
+// Copyright 2023 LiveKit, Inc.
+//
+// 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 utils
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestBitmap(t *testing.T) {
+	b := NewBitmap[uint32](153)
+	require.Equal(t, 4, cap(b.bits)) // rounded up to next power of 2, i. e. 256
+	require.Equal(t, 4, len(b.bits))
+
+	e := make([]uint64, 4)
+	require.Equal(t, e, b.bits)
+
+	b.Set(666666)
+	require.True(t, b.IsSet(666666))
+	require.True(t, b.IsSet(666666+256)) // aliased
+	require.True(t, b.IsSet(666666-256)) // aliased
+	require.False(t, b.IsSet(666666-1))
+	require.False(t, b.IsSet(666666+1))
+
+	copy(e, b.bits)
+	b.Set(42) // same as above - aliased
+	require.Equal(t, e, b.bits)
+
+	// same slot range
+	b.SetRange(24, 63)
+	// different slot range
+	b.SetRange(64, 240)
+	require.False(t, b.IsSet(241))
+	require.False(t, b.IsSet(23))
+	require.True(t, b.IsSet(24))
+	require.True(t, b.IsSet(240))
+	e[0] = 0xFFFF_FFFF_FF00_0000
+	e[1] = ^uint64(0)
+	e[2] = ^uint64(0)
+	e[3] = 0x0001_FFFF_FFFF_FFFF
+	require.Equal(t, e, b.bits)
+
+	b.Clear(6700) // aliases to 44
+	e[0] = 0xFFFF_EFFF_FF00_0000
+	require.Equal(t, e, b.bits)
+	require.False(t, b.IsSet(44))
+
+	// same slot range
+	b.ClearRange(24, 34)
+	e[0] = 0xFFFF_EFF8_0000_0000
+	require.Equal(t, e, b.bits)
+	require.False(t, b.IsSet(24))
+	require.True(t, b.IsSet(35))
+
+	// different slot range
+	b.ClearRange(95, 234)
+	e[1] = 0x0000_0000_7FFF_FFFF
+	e[2] = 0x0
+	e[3] = 0x0001_F800_0000_0000
+	require.Equal(t, e, b.bits)
+	require.True(t, b.IsSet(94))
+	require.False(t, b.IsSet(95))
+	require.False(t, b.IsSet(234))
+	require.True(t, b.IsSet(235))
+}

utils/timed_aggregator.go

@@ -28,15 +28,15 @@ var (
 
 // ------------------------------------------------
 
-type number interface {
+type timedAggregatorNumber interface {
 	int64 | float64
 }
 
 type TimedAggregatorParams struct {
 	CapNegativeValues bool
 }
 
-type TimedAggregator[T number] struct {
+type TimedAggregator[T timedAggregatorNumber] struct {
 	params TimedAggregatorParams
 
 	lock              sync.RWMutex
@@ -46,7 +46,7 @@ type TimedAggregator[T number] struct {
 	aggregateDuration time.Duration
 }
 
-func NewTimedAggregator[T number](params TimedAggregatorParams) *TimedAggregator[T] {
+func NewTimedAggregator[T timedAggregatorNumber](params TimedAggregatorParams) *TimedAggregator[T] {
 	return &TimedAggregator[T]{
 		params: params,
 	}