github.com/livekit/protocol@v1.39.3/utils/bitmap.go (about)

     1  // Copyright 2023 LiveKit, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package utils
    16  
    17  import "math/bits"
    18  
    19  type bitmapNumber interface {
    20  	uint8 | uint16 | uint32 | uint64
    21  }
    22  
    23  type Bitmap[T bitmapNumber] struct {
    24  	bits []uint64
    25  }
    26  
    27  func NewBitmap[T bitmapNumber](size int) *Bitmap[T] {
    28  	return &Bitmap[T]{
    29  		bits: make([]uint64, 1<<bits.Len64(uint64(size+63)/64)),
    30  	}
    31  }
    32  
    33  func (b *Bitmap[T]) Len() int {
    34  	return len(b.bits) << 6
    35  }
    36  
    37  func (b *Bitmap[T]) Set(val T) {
    38  	sm, s, o := b.getSlotAndOffset(val)
    39  	b.bits[s&sm] |= 1 << o
    40  }
    41  
    42  func (b *Bitmap[T]) GetAndSet(val T) bool {
    43  	sm, s, o := b.getSlotAndOffset(val)
    44  	prev := b.bits[s&sm]&(1<<o) != 0
    45  	b.bits[s&sm] |= 1 << o
    46  	return prev
    47  }
    48  
    49  func (b *Bitmap[T]) SetRange(min, max T) {
    50  	if max < min {
    51  		return
    52  	}
    53  
    54  	sm, ls, rs, lo, ro := b.getSlotsAndOffsets(min, max)
    55  	if ls == rs {
    56  		b.bits[ls&sm] |= (((1 << (ro - lo + 1)) - 1) << lo)
    57  	} else {
    58  		b.bits[ls&sm] |= ^((1 << lo) - 1)
    59  		for i := ls + 1; i < rs; i++ {
    60  			b.bits[i&sm] = ^uint64(0)
    61  		}
    62  		b.bits[rs&sm] |= (1 << (ro + 1)) - 1
    63  	}
    64  }
    65  
    66  func (b *Bitmap[T]) Clear(val T) {
    67  	sm, s, o := b.getSlotAndOffset(val)
    68  	b.bits[s&sm] &= ^(1 << o)
    69  }
    70  
    71  func (b *Bitmap[T]) ClearRange(min, max T) {
    72  	if max < min {
    73  		return
    74  	}
    75  
    76  	sm, ls, rs, lo, ro := b.getSlotsAndOffsets(min, max)
    77  	if ls == rs {
    78  		b.bits[ls&sm] &= ^(((1 << (ro - lo + 1)) - 1) << lo)
    79  	} else {
    80  		b.bits[ls&sm] &= ^uint64(0) >> (64 - lo)
    81  		for i := ls + 1; i < rs; i++ {
    82  			b.bits[i&sm] = 0
    83  		}
    84  		b.bits[rs&sm] &= ^uint64(0) << (ro + 1)
    85  	}
    86  }
    87  
    88  func (b *Bitmap[T]) IsSet(val T) bool {
    89  	sm, s, o := b.getSlotAndOffset(val)
    90  	return b.bits[s&sm]&(1<<o) != 0
    91  }
    92  
    93  func (b *Bitmap[T]) getSlotAndOffset(val T) (sm, s, o int) {
    94  	sm = len(b.bits) - 1 // slot mask
    95  	s = int(val >> 6)    // slot
    96  	o = int(val & 0x3f)  // offset
    97  	return
    98  }
    99  
   100  func (b *Bitmap[T]) getSlotsAndOffsets(min, max T) (sm, ls, rs, lo, ro int) {
   101  	sm = len(b.bits) - 1 // slot mask
   102  
   103  	ls = int(min >> 6) // left slot
   104  	rs = int(max >> 6) // right slot
   105  
   106  	if rs-ls > len(b.bits) {
   107  		rs = ls + len(b.bits)
   108  		return
   109  	}
   110  
   111  	lo = int(min & 0x3f) // left offset
   112  	ro = int(max & 0x3f) // right offset
   113  	return
   114  }