github.com/database64128/shadowsocks-go@v1.10.2-0.20240315062903-143a773533f1/bitset/bitset.go (about)

     1  package bitset
     2  
     3  import (
     4  	"fmt"
     5  	"math/bits"
     6  )
     7  
     8  const blockBits = bits.UintSize
     9  
    10  // BitSet is a set of bits.
    11  type BitSet struct {
    12  	blocks   []uint
    13  	capacity uint
    14  }
    15  
    16  // NewBitSet returns a new BitSet with the given capacity.
    17  func NewBitSet(capacity uint) BitSet {
    18  	return BitSet{
    19  		blocks:   make([]uint, (capacity+blockBits-1)/blockBits),
    20  		capacity: capacity,
    21  	}
    22  }
    23  
    24  // Capacity returns the capacity of the BitSet.
    25  func (s BitSet) Capacity() uint {
    26  	return s.capacity
    27  }
    28  
    29  // Count returns the number of bits set to 1.
    30  func (s BitSet) Count() (count uint) {
    31  	for i := range s.blocks {
    32  		count += uint(bits.OnesCount(s.blocks[i]))
    33  	}
    34  	return
    35  }
    36  
    37  func (s BitSet) checkIndex(index uint) {
    38  	if index >= s.capacity {
    39  		panic(fmt.Sprintf("bitset: index out of range [%d] with capacity %d", index, s.capacity))
    40  	}
    41  }
    42  
    43  func (s BitSet) blockIndex(index uint) uint {
    44  	return index / blockBits
    45  }
    46  
    47  func (s BitSet) bitIndex(index uint) uint {
    48  	return index % blockBits
    49  }
    50  
    51  // IsSet returns whether the bit at the given index is set to 1.
    52  func (s BitSet) IsSet(index uint) bool {
    53  	s.checkIndex(index)
    54  	return s.blocks[s.blockIndex(index)]&(1<<s.bitIndex(index)) != 0
    55  }
    56  
    57  // Set sets the bit at the given index to 1.
    58  func (s *BitSet) Set(index uint) {
    59  	s.checkIndex(index)
    60  	s.blocks[s.blockIndex(index)] |= 1 << s.bitIndex(index)
    61  }
    62  
    63  // Unset sets the bit at the given index to 0.
    64  func (s *BitSet) Unset(index uint) {
    65  	s.checkIndex(index)
    66  	s.blocks[s.blockIndex(index)] &^= 1 << s.bitIndex(index)
    67  }
    68  
    69  // Flip flips the bit at the given index.
    70  func (s *BitSet) Flip(index uint) {
    71  	s.checkIndex(index)
    72  	s.blocks[s.blockIndex(index)] ^= 1 << s.bitIndex(index)
    73  }
    74  
    75  // SetAll sets all bits to 1.
    76  func (s *BitSet) SetAll() {
    77  	fullBlocks := s.blockIndex(s.capacity)
    78  	for i := range s.blocks[:fullBlocks] {
    79  		s.blocks[i] = ^uint(0)
    80  	}
    81  	if fullBlocks < uint(len(s.blocks)) {
    82  		s.blocks[fullBlocks] = ^(^uint(0) << s.bitIndex(s.capacity))
    83  	}
    84  }
    85  
    86  // UnsetAll sets all bits to 0.
    87  func (s *BitSet) UnsetAll() {
    88  	clear(s.blocks)
    89  }
    90  
    91  // FlipAll flips all bits.
    92  func (s *BitSet) FlipAll() {
    93  	fullBlocks := s.blockIndex(s.capacity)
    94  	for i := range s.blocks[:fullBlocks] {
    95  		s.blocks[i] = ^s.blocks[i]
    96  	}
    97  	if fullBlocks < uint(len(s.blocks)) {
    98  		s.blocks[fullBlocks] ^= ^(^uint(0) << s.bitIndex(s.capacity))
    99  	}
   100  }