github.com/livekit/protocol@v1.16.1-0.20240517185851-47e4c6bba773/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  	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]) SetRange(min, max T) {
    43  	if max < min {
    44  		return
    45  	}
    46  
    47  	sm, ls, rs, lo, ro := b.getSlotsAndOffsets(min, max)
    48  	if ls == rs {
    49  		b.bits[ls&sm] |= (((1 << (ro - lo + 1)) - 1) << lo)
    50  	} else {
    51  		b.bits[ls&sm] |= ^((1 << lo) - 1)
    52  		for i := ls + 1; i < rs; i++ {
    53  			b.bits[i&sm] = ^uint64(0)
    54  		}
    55  		b.bits[rs&sm] |= (1 << (ro + 1)) - 1
    56  	}
    57  }
    58  
    59  func (b *Bitmap[T]) Clear(val T) {
    60  	sm, s, o := b.getSlotAndOffset(val)
    61  	b.bits[s&sm] &= ^(1 << o)
    62  }
    63  
    64  func (b *Bitmap[T]) ClearRange(min, max T) {
    65  	if max < min {
    66  		return
    67  	}
    68  
    69  	sm, ls, rs, lo, ro := b.getSlotsAndOffsets(min, max)
    70  	if ls == rs {
    71  		b.bits[ls&sm] &= ^(((1 << (ro - lo + 1)) - 1) << lo)
    72  	} else {
    73  		b.bits[ls&sm] &= ^uint64(0) >> (64 - lo)
    74  		for i := ls + 1; i < rs; i++ {
    75  			b.bits[i&sm] = 0
    76  		}
    77  		b.bits[rs&sm] &= ^uint64(0) << (ro + 1)
    78  	}
    79  }
    80  
    81  func (b *Bitmap[T]) IsSet(val T) bool {
    82  	sm, s, o := b.getSlotAndOffset(val)
    83  	return b.bits[s&sm]&(1<<o) != 0
    84  }
    85  
    86  func (b *Bitmap[T]) getSlotAndOffset(val T) (sm, s, o int) {
    87  	sm = len(b.bits) - 1 // slot mask
    88  	s = int(val >> 6)    // slot
    89  	o = int(val & 0x3f)  // offset
    90  	return
    91  }
    92  
    93  func (b *Bitmap[T]) getSlotsAndOffsets(min, max T) (sm, ls, rs, lo, ro int) {
    94  	sm = len(b.bits) - 1 // slot mask
    95  
    96  	ls = int(min >> 6) // left slot
    97  	rs = int(max >> 6) // right slot
    98  
    99  	if rs-ls > len(b.bits) {
   100  		rs = ls + len(b.bits)
   101  		return
   102  	}
   103  
   104  	lo = int(min & 0x3f) // left offset
   105  	ro = int(max & 0x3f) // right offset
   106  	return
   107  }