bitset.gno

  1package bitset
  2
  3import "strings"
  4
  5const bitsPerWord = 64
  6
  7// BitSet implements an arbitrary-size bit array.
  8type BitSet struct {
  9	words []uint64
 10}
 11
 12// New creates a new BitSet pre-allocated to hold at least size bits.
 13func New(size uint64) BitSet {
 14	if size == 0 {
 15		return BitSet{}
 16	}
 17	return BitSet{words: make([]uint64, wordIndex(size-1)+1)}
 18}
 19
 20// Set turns on the bit at a given position.
 21func (b *BitSet) Set(i uint64) {
 22	idx := wordIndex(i)
 23	if idx >= len(b.words) {
 24		b.grow(idx + 1)
 25	}
 26	b.words[idx] |= bitMask(i) // Turn bit on
 27}
 28
 29// Clear turns off the bit at a given position.
 30func (b *BitSet) Clear(i uint64) {
 31	idx := wordIndex(i)
 32	if idx < len(b.words) {
 33		b.words[idx] &^= bitMask(i) // Turn bit off
 34	}
 35}
 36
 37// ClearAll turns off all bits.
 38// The size of the bitset remains unchanged.
 39func (b *BitSet) ClearAll() {
 40	for i := range b.words {
 41		b.words[i] = 0
 42	}
 43}
 44
 45// Compact reclaims memory by removing trailing zero words.
 46func (b *BitSet) Compact() {
 47	end := len(b.words)
 48	for ; end > 0 && b.words[end-1] == 0; end-- {
 49	}
 50	words := make([]uint64, end)
 51	copy(words, b.words)
 52	b.words = words
 53}
 54
 55// IsSet checks whether the bit at a given position is set.
 56func (b BitSet) IsSet(i uint64) bool {
 57	idx := wordIndex(i)
 58	if idx >= len(b.words) {
 59		return false // Bit has not been set yet
 60	}
 61	return b.words[idx]&bitMask(i) != 0
 62}
 63
 64// Size returns the total number of bits (including unset) currently stored.
 65// It's the total space allocated within the set in bits.
 66func (b BitSet) Size() int {
 67	return len(b.words) * bitsPerWord
 68}
 69
 70// Len returns the number of set bits.
 71func (b BitSet) Len() int {
 72	var count int
 73	for _, w := range b.words {
 74		count += countSetBits(w)
 75	}
 76	return count
 77}
 78
 79// And performs an in-place AND with other bitset.
 80// Bits beyond the other bitset are cleared.
 81func (b *BitSet) And(other BitSet) {
 82	for i := range b.words {
 83		if i < len(other.words) {
 84			b.words[i] &= other.words[i]
 85		} else {
 86			b.words[i] = 0
 87		}
 88	}
 89}
 90
 91// Or performs an in-place OR (union) with other bitset.
 92// Current bitset grows if the other bitset is bigger.
 93func (b *BitSet) Or(other BitSet) {
 94	otherLen := len(other.words)
 95	if otherLen > len(b.words) {
 96		b.grow(otherLen)
 97	}
 98
 99	for i := range other.words {
100		b.words[i] |= other.words[i]
101	}
102}
103
104// Xor performs an in-place XOR with other.
105// Current bitset grows if the other bitset is bigger.
106func (b *BitSet) Xor(other BitSet) {
107	otherLen := len(other.words)
108	if otherLen > len(b.words) {
109		b.grow(otherLen)
110	}
111
112	for i := range other.words {
113		b.words[i] ^= other.words[i]
114	}
115}
116
117// Equal checks whether other bitset have exactly the same bits set.
118func (b BitSet) Equal(other BitSet) bool {
119	short, long := b.words, other.words
120	if len(short) > len(long) {
121		short, long = long, short
122	}
123
124	for i := range short {
125		if short[i] != long[i] {
126			return false
127		}
128	}
129
130	// Any trailing words in the longer set must be zero
131	for i := len(short); i < len(long); i++ {
132		if long[i] != 0 {
133			return false
134		}
135	}
136	return true
137}
138
139// PaddedString returns a binary representation of the bitset with zero padding.
140// Representation uses MSB-first binary representation.
141func (b BitSet) PaddedString() string {
142	wordCount := len(b.words)
143	if wordCount == 0 {
144		return ""
145	}
146
147	buf := make([]byte, wordCount*bitsPerWord)
148	for i, w := range b.words {
149		base := (wordCount - 1 - i) * bitsPerWord
150		for j := 0; j < bitsPerWord; j++ {
151			if w>>uint64(bitsPerWord-1-j)&1 == 1 {
152				buf[base+j] = '1'
153			} else {
154				buf[base+j] = '0'
155			}
156		}
157	}
158	return string(buf)
159}
160
161// String returns a binary representation of the bitset.
162// Representation uses MSB-first binary representation.
163func (b BitSet) String() string {
164	s := strings.TrimLeft(b.PaddedString(), "0")
165	if s == "" && len(b.words) > 0 {
166		return "0"
167	}
168	return s
169}
170
171func (b *BitSet) grow(length int) {
172	words := make([]uint64, length)
173	copy(words, b.words)
174	b.words = words
175}
176
177func wordIndex(i uint64) int {
178	return int(i / bitsPerWord)
179}
180
181func bitMask(i uint64) uint64 {
182	return 1 << (i % bitsPerWord)
183}
184
185func countSetBits(x uint64) int {
186	// Count the number of set bits using Kernighan's algorithm
187	var count int
188	for x != 0 {
189		x &= x - 1
190		count++
191	}
192	return count
193}