vendor/github.com/tendermint/tmlibs/common/bit_array.go

   1 package common
   2
   3 import (
   4         "encoding/binary"
   5         "fmt"
   6         "math/rand"
   7         "strings"
   8         "sync"
   9 )
  10
  11 type BitArray struct {
  12         mtx   sync.Mutex
  13         Bits  int      `json:"bits"`  // NOTE: persisted via reflect, must be exported
  14         Elems []uint64 `json:"elems"` // NOTE: persisted via reflect, must be exported
  15 }
  16
  17 // There is no BitArray whose Size is 0.  Use nil instead.
  18 func NewBitArray(bits int) *BitArray {
  19         if bits == 0 {
  20                 return nil
  21         }
  22         return &BitArray{
  23                 Bits:  bits,
  24                 Elems: make([]uint64, (bits+63)/64),
  25         }
  26 }
  27
  28 func (bA *BitArray) Size() int {
  29         if bA == nil {
  30                 return 0
  31         }
  32         return bA.Bits
  33 }
  34
  35 // NOTE: behavior is undefined if i >= bA.Bits
  36 func (bA *BitArray) GetIndex(i int) bool {
  37         if bA == nil {
  38                 return false
  39         }
  40         bA.mtx.Lock()
  41         defer bA.mtx.Unlock()
  42         return bA.getIndex(i)
  43 }
  44
  45 func (bA *BitArray) getIndex(i int) bool {
  46         if i >= bA.Bits {
  47                 return false
  48         }
  49         return bA.Elems[i/64]&(uint64(1)<<uint(i%64)) > 0
  50 }
  51
  52 // NOTE: behavior is undefined if i >= bA.Bits
  53 func (bA *BitArray) SetIndex(i int, v bool) bool {
  54         if bA == nil {
  55                 return false
  56         }
  57         bA.mtx.Lock()
  58         defer bA.mtx.Unlock()
  59         return bA.setIndex(i, v)
  60 }
  61
  62 func (bA *BitArray) setIndex(i int, v bool) bool {
  63         if i >= bA.Bits {
  64                 return false
  65         }
  66         if v {
  67                 bA.Elems[i/64] |= (uint64(1) << uint(i%64))
  68         } else {
  69                 bA.Elems[i/64] &= ^(uint64(1) << uint(i%64))
  70         }
  71         return true
  72 }
  73
  74 func (bA *BitArray) Copy() *BitArray {
  75         if bA == nil {
  76                 return nil
  77         }
  78         bA.mtx.Lock()
  79         defer bA.mtx.Unlock()
  80         return bA.copy()
  81 }
  82
  83 func (bA *BitArray) copy() *BitArray {
  84         c := make([]uint64, len(bA.Elems))
  85         copy(c, bA.Elems)
  86         return &BitArray{
  87                 Bits:  bA.Bits,
  88                 Elems: c,
  89         }
  90 }
  91
  92 func (bA *BitArray) copyBits(bits int) *BitArray {
  93         c := make([]uint64, (bits+63)/64)
  94         copy(c, bA.Elems)
  95         return &BitArray{
  96                 Bits:  bits,
  97                 Elems: c,
  98         }
  99 }
 100
 101 // Returns a BitArray of larger bits size.
 102 func (bA *BitArray) Or(o *BitArray) *BitArray {
 103         if bA == nil {
 104                 o.Copy()
 105         }
 106         bA.mtx.Lock()
 107         defer bA.mtx.Unlock()
 108         c := bA.copyBits(MaxInt(bA.Bits, o.Bits))
 109         for i := 0; i < len(c.Elems); i++ {
 110                 c.Elems[i] |= o.Elems[i]
 111         }
 112         return c
 113 }
 114
 115 // Returns a BitArray of smaller bit size.
 116 func (bA *BitArray) And(o *BitArray) *BitArray {
 117         if bA == nil {
 118                 return nil
 119         }
 120         bA.mtx.Lock()
 121         defer bA.mtx.Unlock()
 122         return bA.and(o)
 123 }
 124
 125 func (bA *BitArray) and(o *BitArray) *BitArray {
 126         c := bA.copyBits(MinInt(bA.Bits, o.Bits))
 127         for i := 0; i < len(c.Elems); i++ {
 128                 c.Elems[i] &= o.Elems[i]
 129         }
 130         return c
 131 }
 132
 133 func (bA *BitArray) Not() *BitArray {
 134         if bA == nil {
 135                 return nil // Degenerate
 136         }
 137         bA.mtx.Lock()
 138         defer bA.mtx.Unlock()
 139         c := bA.copy()
 140         for i := 0; i < len(c.Elems); i++ {
 141                 c.Elems[i] = ^c.Elems[i]
 142         }
 143         return c
 144 }
 145
 146 func (bA *BitArray) Sub(o *BitArray) *BitArray {
 147         if bA == nil {
 148                 return nil
 149         }
 150         bA.mtx.Lock()
 151         defer bA.mtx.Unlock()
 152         if bA.Bits > o.Bits {
 153                 c := bA.copy()
 154                 for i := 0; i < len(o.Elems)-1; i++ {
 155                         c.Elems[i] &= ^c.Elems[i]
 156                 }
 157                 i := len(o.Elems) - 1
 158                 if i >= 0 {
 159                         for idx := i * 64; idx < o.Bits; idx++ {
 160                                 // NOTE: each individual GetIndex() call to o is safe.
 161                                 c.setIndex(idx, c.getIndex(idx) && !o.GetIndex(idx))
 162                         }
 163                 }
 164                 return c
 165         } else {
 166                 return bA.and(o.Not()) // Note degenerate case where o == nil
 167         }
 168 }
 169
 170 func (bA *BitArray) IsEmpty() bool {
 171         if bA == nil {
 172                 return true // should this be opposite?
 173         }
 174         bA.mtx.Lock()
 175         defer bA.mtx.Unlock()
 176         for _, e := range bA.Elems {
 177                 if e > 0 {
 178                         return false
 179                 }
 180         }
 181         return true
 182 }
 183
 184 func (bA *BitArray) IsFull() bool {
 185         if bA == nil {
 186                 return true
 187         }
 188         bA.mtx.Lock()
 189         defer bA.mtx.Unlock()
 190
 191         // Check all elements except the last
 192         for _, elem := range bA.Elems[:len(bA.Elems)-1] {
 193                 if (^elem) != 0 {
 194                         return false
 195                 }
 196         }
 197
 198         // Check that the last element has (lastElemBits) 1's
 199         lastElemBits := (bA.Bits+63)%64 + 1
 200         lastElem := bA.Elems[len(bA.Elems)-1]
 201         return (lastElem+1)&((uint64(1)<<uint(lastElemBits))-1) == 0
 202 }
 203
 204 func (bA *BitArray) PickRandom() (int, bool) {
 205         if bA == nil {
 206                 return 0, false
 207         }
 208         bA.mtx.Lock()
 209         defer bA.mtx.Unlock()
 210
 211         length := len(bA.Elems)
 212         if length == 0 {
 213                 return 0, false
 214         }
 215         randElemStart := rand.Intn(length)
 216         for i := 0; i < length; i++ {
 217                 elemIdx := ((i + randElemStart) % length)
 218                 if elemIdx < length-1 {
 219                         if bA.Elems[elemIdx] > 0 {
 220                                 randBitStart := rand.Intn(64)
 221                                 for j := 0; j < 64; j++ {
 222                                         bitIdx := ((j + randBitStart) % 64)
 223                                         if (bA.Elems[elemIdx] & (uint64(1) << uint(bitIdx))) > 0 {
 224                                                 return 64*elemIdx + bitIdx, true
 225                                         }
 226                                 }
 227                                 PanicSanity("should not happen")
 228                         }
 229                 } else {
 230                         // Special case for last elem, to ignore straggler bits
 231                         elemBits := bA.Bits % 64
 232                         if elemBits == 0 {
 233                                 elemBits = 64
 234                         }
 235                         randBitStart := rand.Intn(elemBits)
 236                         for j := 0; j < elemBits; j++ {
 237                                 bitIdx := ((j + randBitStart) % elemBits)
 238                                 if (bA.Elems[elemIdx] & (uint64(1) << uint(bitIdx))) > 0 {
 239                                         return 64*elemIdx + bitIdx, true
 240                                 }
 241                         }
 242                 }
 243         }
 244         return 0, false
 245 }
 246
 247 func (bA *BitArray) String() string {
 248         if bA == nil {
 249                 return "nil-BitArray"
 250         }
 251         bA.mtx.Lock()
 252         defer bA.mtx.Unlock()
 253         return bA.stringIndented("")
 254 }
 255
 256 func (bA *BitArray) StringIndented(indent string) string {
 257         if bA == nil {
 258                 return "nil-BitArray"
 259         }
 260         bA.mtx.Lock()
 261         defer bA.mtx.Unlock()
 262         return bA.stringIndented(indent)
 263 }
 264
 265 func (bA *BitArray) stringIndented(indent string) string {
 266
 267         lines := []string{}
 268         bits := ""
 269         for i := 0; i < bA.Bits; i++ {
 270                 if bA.getIndex(i) {
 271                         bits += "X"
 272                 } else {
 273                         bits += "_"
 274                 }
 275                 if i%100 == 99 {
 276                         lines = append(lines, bits)
 277                         bits = ""
 278                 }
 279                 if i%10 == 9 {
 280                         bits += " "
 281                 }
 282                 if i%50 == 49 {
 283                         bits += " "
 284                 }
 285         }
 286         if len(bits) > 0 {
 287                 lines = append(lines, bits)
 288         }
 289         return fmt.Sprintf("BA{%v:%v}", bA.Bits, strings.Join(lines, indent))
 290 }
 291
 292 func (bA *BitArray) Bytes() []byte {
 293         bA.mtx.Lock()
 294         defer bA.mtx.Unlock()
 295
 296         numBytes := (bA.Bits + 7) / 8
 297         bytes := make([]byte, numBytes)
 298         for i := 0; i < len(bA.Elems); i++ {
 299                 elemBytes := [8]byte{}
 300                 binary.LittleEndian.PutUint64(elemBytes[:], bA.Elems[i])
 301                 copy(bytes[i*8:], elemBytes[:])
 302         }
 303         return bytes
 304 }
 305
 306 // NOTE: other bitarray o is not locked when reading,
 307 // so if necessary, caller must copy or lock o prior to calling Update.
 308 // If bA is nil, does nothing.
 309 func (bA *BitArray) Update(o *BitArray) {
 310         if bA == nil {
 311                 return
 312         }
 313         bA.mtx.Lock()
 314         defer bA.mtx.Unlock()
 315
 316         copy(bA.Elems, o.Elems)
 317 }