diff --git a/utils/bitmap.go b/utils/bitmap.go new file mode 100644 index 00000000..4ad206a7 --- /dev/null +++ b/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<> 6) // left slot + rs = int(max >> 6) // right slot + 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 new file mode 100644 index 00000000..9abadc98 --- /dev/null +++ b/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)) +} diff --git a/utils/timed_aggregator.go b/utils/timed_aggregator.go index d5ba4cd0..f524c6bc 100644 --- a/utils/timed_aggregator.go +++ b/utils/timed_aggregator.go @@ -28,7 +28,7 @@ var ( // ------------------------------------------------ -type number interface { +type timedAggregatorNumber interface { int64 | float64 } @@ -36,7 +36,7 @@ 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, }