github.com/pokt-network/tendermint@v0.32.11-0.20230426215212-59310158d3e9/libs/bits/bit_array.go (about) 1 package bits 2 3 import ( 4 "encoding/binary" 5 "fmt" 6 "regexp" 7 "strings" 8 "sync" 9 10 tmmath "github.com/tendermint/tendermint/libs/math" 11 tmrand "github.com/tendermint/tendermint/libs/rand" 12 tmprotobits "github.com/tendermint/tendermint/proto/libs/bits" 13 ) 14 15 // BitArray is a thread-safe implementation of a bit array. 16 type BitArray struct { 17 mtx sync.Mutex 18 Bits int `json:"bits"` // NOTE: persisted via reflect, must be exported 19 Elems []uint64 `json:"elems"` // NOTE: persisted via reflect, must be exported 20 trueIndicesSlice []int 21 } 22 23 // NewBitArray returns a new bit array. 24 // It returns nil if the number of bits is zero. 25 func NewBitArray(bits int) *BitArray { 26 if bits <= 0 { 27 return nil 28 } 29 return &BitArray{ 30 Bits: bits, 31 Elems: make([]uint64, (bits+63)/64), 32 } 33 } 34 35 // Size returns the number of bits in the bitarray 36 func (bA *BitArray) Size() int { 37 if bA == nil { 38 return 0 39 } 40 return bA.Bits 41 } 42 43 // GetIndex returns the bit at index i within the bit array. 44 // The behavior is undefined if i >= bA.Bits 45 func (bA *BitArray) GetIndex(i int) bool { 46 if bA == nil { 47 return false 48 } 49 bA.mtx.Lock() 50 defer bA.mtx.Unlock() 51 return bA.getIndex(i) 52 } 53 54 func (bA *BitArray) getIndex(i int) bool { 55 if i >= bA.Bits { 56 return false 57 } 58 return bA.Elems[i/64]&(uint64(1)<<uint(i%64)) > 0 59 } 60 61 // SetIndex sets the bit at index i within the bit array. 62 // The behavior is undefined if i >= bA.Bits 63 func (bA *BitArray) SetIndex(i int, v bool) bool { 64 if bA == nil { 65 return false 66 } 67 bA.mtx.Lock() 68 defer bA.mtx.Unlock() 69 return bA.setIndex(i, v) 70 } 71 72 func (bA *BitArray) setIndex(i int, v bool) bool { 73 if i >= bA.Bits { 74 return false 75 } 76 if v { 77 bA.Elems[i/64] |= (uint64(1) << uint(i%64)) 78 } else { 79 bA.Elems[i/64] &= ^(uint64(1) << uint(i%64)) 80 } 81 return true 82 } 83 84 // Copy returns a copy of the provided bit array. 85 func (bA *BitArray) Copy() *BitArray { 86 if bA == nil { 87 return nil 88 } 89 bA.mtx.Lock() 90 defer bA.mtx.Unlock() 91 return bA.copy() 92 } 93 94 func (bA *BitArray) copy() *BitArray { 95 c := make([]uint64, len(bA.Elems)) 96 copy(c, bA.Elems) 97 return &BitArray{ 98 Bits: bA.Bits, 99 Elems: c, 100 } 101 } 102 103 func (bA *BitArray) copyBits(bits int) *BitArray { 104 c := make([]uint64, (bits+63)/64) 105 copy(c, bA.Elems) 106 return &BitArray{ 107 Bits: bits, 108 Elems: c, 109 } 110 } 111 112 // Or returns a bit array resulting from a bitwise OR of the two bit arrays. 113 // If the two bit-arrys have different lengths, Or right-pads the smaller of the two bit-arrays with zeroes. 114 // Thus the size of the return value is the maximum of the two provided bit arrays. 115 func (bA *BitArray) Or(o *BitArray) *BitArray { 116 if bA == nil && o == nil { 117 return nil 118 } 119 if bA == nil && o != nil { 120 return o.Copy() 121 } 122 if o == nil { 123 return bA.Copy() 124 } 125 bA.mtx.Lock() 126 o.mtx.Lock() 127 c := bA.copyBits(tmmath.MaxInt(bA.Bits, o.Bits)) 128 smaller := tmmath.MinInt(len(bA.Elems), len(o.Elems)) 129 for i := 0; i < smaller; i++ { 130 c.Elems[i] |= o.Elems[i] 131 } 132 bA.mtx.Unlock() 133 o.mtx.Unlock() 134 return c 135 } 136 137 // And returns a bit array resulting from a bitwise AND of the two bit arrays. 138 // If the two bit-arrys have different lengths, this truncates the larger of the two bit-arrays from the right. 139 // Thus the size of the return value is the minimum of the two provided bit arrays. 140 func (bA *BitArray) And(o *BitArray) *BitArray { 141 if bA == nil || o == nil { 142 return nil 143 } 144 bA.mtx.Lock() 145 o.mtx.Lock() 146 defer func() { 147 bA.mtx.Unlock() 148 o.mtx.Unlock() 149 }() 150 return bA.and(o) 151 } 152 153 func (bA *BitArray) and(o *BitArray) *BitArray { 154 c := bA.copyBits(tmmath.MinInt(bA.Bits, o.Bits)) 155 for i := 0; i < len(c.Elems); i++ { 156 c.Elems[i] &= o.Elems[i] 157 } 158 return c 159 } 160 161 // Not returns a bit array resulting from a bitwise Not of the provided bit array. 162 func (bA *BitArray) Not() *BitArray { 163 if bA == nil { 164 return nil // Degenerate 165 } 166 bA.mtx.Lock() 167 defer bA.mtx.Unlock() 168 return bA.not() 169 } 170 171 func (bA *BitArray) not() *BitArray { 172 c := bA.copy() 173 for i := 0; i < len(c.Elems); i++ { 174 c.Elems[i] = ^c.Elems[i] 175 } 176 return c 177 } 178 179 // Sub subtracts the two bit-arrays bitwise, without carrying the bits. 180 // Note that carryless subtraction of a - b is (a and not b). 181 // The output is the same as bA, regardless of o's size. 182 // If bA is longer than o, o is right padded with zeroes 183 func (bA *BitArray) Sub(o *BitArray) *BitArray { 184 if bA == nil || o == nil { 185 // TODO: Decide if we should do 1's complement here? 186 return nil 187 } 188 bA.mtx.Lock() 189 o.mtx.Lock() 190 // output is the same size as bA 191 c := bA.copyBits(bA.Bits) 192 // Only iterate to the minimum size between the two. 193 // If o is longer, those bits are ignored. 194 // If bA is longer, then skipping those iterations is equivalent 195 // to right padding with 0's 196 smaller := tmmath.MinInt(len(bA.Elems), len(o.Elems)) 197 for i := 0; i < smaller; i++ { 198 // &^ is and not in golang 199 c.Elems[i] &^= o.Elems[i] 200 } 201 bA.mtx.Unlock() 202 o.mtx.Unlock() 203 return c 204 } 205 206 // IsEmpty returns true iff all bits in the bit array are 0 207 func (bA *BitArray) IsEmpty() bool { 208 if bA == nil { 209 return true // should this be opposite? 210 } 211 bA.mtx.Lock() 212 defer bA.mtx.Unlock() 213 for _, e := range bA.Elems { 214 if e > 0 { 215 return false 216 } 217 } 218 return true 219 } 220 221 // IsFull returns true iff all bits in the bit array are 1. 222 func (bA *BitArray) IsFull() bool { 223 if bA == nil { 224 return true 225 } 226 bA.mtx.Lock() 227 defer bA.mtx.Unlock() 228 229 // Check all elements except the last 230 for _, elem := range bA.Elems[:len(bA.Elems)-1] { 231 if (^elem) != 0 { 232 return false 233 } 234 } 235 236 // Check that the last element has (lastElemBits) 1's 237 lastElemBits := (bA.Bits+63)%64 + 1 238 lastElem := bA.Elems[len(bA.Elems)-1] 239 return (lastElem+1)&((uint64(1)<<uint(lastElemBits))-1) == 0 240 } 241 242 // PickRandom returns a random index for a set bit in the bit array. 243 // If there is no such value, it returns 0, false. 244 // It uses the global randomness in `random.go` to get this index. 245 func (bA *BitArray) PickRandom() (int, bool) { 246 if bA == nil { 247 return 0, false 248 } 249 250 bA.mtx.Lock() 251 defer bA.mtx.Unlock() 252 bA.getTrueIndices() 253 if len(bA.trueIndicesSlice) == 0 { // no bits set to true 254 return 0, false 255 } 256 randomIndex := bA.trueIndicesSlice[tmrand.Intn(len(bA.trueIndicesSlice))] 257 // reset the slice 258 bA.trueIndicesSlice = nil 259 return randomIndex, true 260 } 261 262 func (bA *BitArray) getTrueIndices() []int { 263 curBit := 0 264 numElems := len(bA.Elems) 265 // set all true indices 266 for i := 0; i < numElems-1; i++ { 267 elem := bA.Elems[i] 268 if elem == 0 { 269 curBit += 64 270 continue 271 } 272 for j := 0; j < 64; j++ { 273 if (elem & (uint64(1) << uint64(j))) > 0 { 274 bA.trueIndicesSlice = append(bA.trueIndicesSlice, curBit) 275 } 276 curBit++ 277 } 278 } 279 // handle last element 280 lastElem := bA.Elems[numElems-1] 281 numFinalBits := bA.Bits - curBit 282 for i := 0; i < numFinalBits; i++ { 283 if (lastElem & (uint64(1) << uint64(i))) > 0 { 284 bA.trueIndicesSlice = append(bA.trueIndicesSlice, curBit) 285 } 286 curBit++ 287 } 288 return bA.trueIndicesSlice 289 } 290 291 // String returns a string representation of BitArray: BA{<bit-string>}, 292 // where <bit-string> is a sequence of 'x' (1) and '_' (0). 293 // The <bit-string> includes spaces and newlines to help people. 294 // For a simple sequence of 'x' and '_' characters with no spaces or newlines, 295 // see the MarshalJSON() method. 296 // Example: "BA{_x_}" or "nil-BitArray" for nil. 297 func (bA *BitArray) String() string { 298 return bA.StringIndented("") 299 } 300 301 // StringIndented returns the same thing as String(), but applies the indent 302 // at every 10th bit, and twice at every 50th bit. 303 func (bA *BitArray) StringIndented(indent string) string { 304 if bA == nil { 305 return "nil-BitArray" 306 } 307 bA.mtx.Lock() 308 defer bA.mtx.Unlock() 309 return bA.stringIndented(indent) 310 } 311 312 func (bA *BitArray) stringIndented(indent string) string { 313 lines := []string{} 314 bits := "" 315 for i := 0; i < bA.Bits; i++ { 316 if bA.getIndex(i) { 317 bits += "x" 318 } else { 319 bits += "_" 320 } 321 if i%100 == 99 { 322 lines = append(lines, bits) 323 bits = "" 324 } 325 if i%10 == 9 { 326 bits += indent 327 } 328 if i%50 == 49 { 329 bits += indent 330 } 331 } 332 if len(bits) > 0 { 333 lines = append(lines, bits) 334 } 335 return fmt.Sprintf("BA{%v:%v}", bA.Bits, strings.Join(lines, indent)) 336 } 337 338 // Bytes returns the byte representation of the bits within the bitarray. 339 func (bA *BitArray) Bytes() []byte { 340 bA.mtx.Lock() 341 defer bA.mtx.Unlock() 342 343 numBytes := (bA.Bits + 7) / 8 344 bytes := make([]byte, numBytes) 345 for i := 0; i < len(bA.Elems); i++ { 346 elemBytes := [8]byte{} 347 binary.LittleEndian.PutUint64(elemBytes[:], bA.Elems[i]) 348 copy(bytes[i*8:], elemBytes[:]) 349 } 350 return bytes 351 } 352 353 // Update sets the bA's bits to be that of the other bit array. 354 // The copying begins from the begin of both bit arrays. 355 func (bA *BitArray) Update(o *BitArray) { 356 if bA == nil || o == nil { 357 return 358 } 359 bA.mtx.Lock() 360 o.mtx.Lock() 361 defer func() { 362 bA.mtx.Unlock() 363 o.mtx.Unlock() 364 }() 365 366 copy(bA.Elems, o.Elems) 367 } 368 369 // MarshalJSON implements json.Marshaler interface by marshaling bit array 370 // using a custom format: a string of '-' or 'x' where 'x' denotes the 1 bit. 371 func (bA *BitArray) MarshalJSON() ([]byte, error) { 372 if bA == nil { 373 return []byte("null"), nil 374 } 375 376 bA.mtx.Lock() 377 defer bA.mtx.Unlock() 378 379 bits := `"` 380 for i := 0; i < bA.Bits; i++ { 381 if bA.getIndex(i) { 382 bits += `x` 383 } else { 384 bits += `_` 385 } 386 } 387 bits += `"` 388 return []byte(bits), nil 389 } 390 391 var bitArrayJSONRegexp = regexp.MustCompile(`\A"([_x]*)"\z`) 392 393 // UnmarshalJSON implements json.Unmarshaler interface by unmarshaling a custom 394 // JSON description. 395 func (bA *BitArray) UnmarshalJSON(bz []byte) error { 396 b := string(bz) 397 if b == "null" { 398 // This is required e.g. for encoding/json when decoding 399 // into a pointer with pre-allocated BitArray. 400 bA.Bits = 0 401 bA.Elems = nil 402 return nil 403 } 404 405 // Validate 'b'. 406 match := bitArrayJSONRegexp.FindStringSubmatch(b) 407 if match == nil { 408 return fmt.Errorf("bitArray in JSON should be a string of format %q but got %s", bitArrayJSONRegexp.String(), b) 409 } 410 bits := match[1] 411 412 // Construct new BitArray and copy over. 413 numBits := len(bits) 414 bA2 := NewBitArray(numBits) 415 for i := 0; i < numBits; i++ { 416 if bits[i] == 'x' { 417 bA2.SetIndex(i, true) 418 } 419 } 420 *bA = *bA2 //nolint:govet 421 return nil 422 } 423 424 // ToProto converts BitArray to protobuf 425 func (bA *BitArray) ToProto() *tmprotobits.BitArray { 426 if bA == nil || len(bA.Elems) == 0 { 427 return nil 428 } 429 430 return &tmprotobits.BitArray{ 431 Bits: int64(bA.Bits), 432 Elems: bA.Elems, 433 } 434 } 435 436 // FromProto sets a protobuf BitArray to the given pointer. 437 func (bA *BitArray) FromProto(protoBitArray *tmprotobits.BitArray) { 438 if protoBitArray == nil { 439 bA = nil 440 return 441 } 442 443 bA.Bits = int(protoBitArray.Bits) 444 if len(protoBitArray.Elems) > 0 { 445 bA.Elems = protoBitArray.Elems 446 } 447 }