Files
protocol/utils/bitmap.go
Raja Subramanian cf9f95141e bit map (#473)
copying Paul's go playground mostly.
2023-09-16 01:53:28 +05:30

106 lines
2.3 KiB
Go

// 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
}