bitset.gno
4.08 Kb · 193 lines
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}
