github.com/devwanda/aphelion-staking@v0.33.9/crypto/multisig/bitarray/compact_bit_array.go (about) 1 package bitarray 2 3 import ( 4 "bytes" 5 "encoding/binary" 6 "errors" 7 "fmt" 8 "regexp" 9 "strings" 10 ) 11 12 // CompactBitArray is an implementation of a space efficient bit array. 13 // This is used to ensure that the encoded data takes up a minimal amount of 14 // space after amino encoding. 15 // This is not thread safe, and is not intended for concurrent usage. 16 type CompactBitArray struct { 17 ExtraBitsStored byte `json:"extra_bits"` // The number of extra bits in elems. 18 Elems []byte `json:"bits"` 19 } 20 21 // NewCompactBitArray returns a new compact bit array. 22 // It returns nil if the number of bits is zero. 23 func NewCompactBitArray(bits int) *CompactBitArray { 24 if bits <= 0 { 25 return nil 26 } 27 return &CompactBitArray{ 28 ExtraBitsStored: byte(bits % 8), 29 Elems: make([]byte, (bits+7)/8), 30 } 31 } 32 33 // Size returns the number of bits in the bitarray 34 func (bA *CompactBitArray) Size() int { 35 if bA == nil { 36 return 0 37 } else if bA.ExtraBitsStored == byte(0) { 38 return len(bA.Elems) * 8 39 } 40 // num_bits = 8*num_full_bytes + overflow_in_last_byte 41 // num_full_bytes = (len(bA.Elems)-1) 42 return (len(bA.Elems)-1)*8 + int(bA.ExtraBitsStored) 43 } 44 45 // GetIndex returns the bit at index i within the bit array. 46 // The behavior is undefined if i >= bA.Size() 47 func (bA *CompactBitArray) GetIndex(i int) bool { 48 if bA == nil { 49 return false 50 } 51 if i >= bA.Size() { 52 return false 53 } 54 return bA.Elems[i>>3]&(uint8(1)<<uint8(7-(i%8))) > 0 55 } 56 57 // SetIndex sets the bit at index i within the bit array. 58 // The behavior is undefined if i >= bA.Size() 59 func (bA *CompactBitArray) SetIndex(i int, v bool) bool { 60 if bA == nil { 61 return false 62 } 63 if i >= bA.Size() { 64 return false 65 } 66 if v { 67 bA.Elems[i>>3] |= (uint8(1) << uint8(7-(i%8))) 68 } else { 69 bA.Elems[i>>3] &= ^(uint8(1) << uint8(7-(i%8))) 70 } 71 return true 72 } 73 74 // NumTrueBitsBefore returns the number of bits set to true before the 75 // given index. e.g. if bA = _XX__XX, NumOfTrueBitsBefore(4) = 2, since 76 // there are two bits set to true before index 4. 77 func (bA *CompactBitArray) NumTrueBitsBefore(index int) int { 78 numTrueValues := 0 79 for i := 0; i < index; i++ { 80 if bA.GetIndex(i) { 81 numTrueValues++ 82 } 83 } 84 return numTrueValues 85 } 86 87 // Copy returns a copy of the provided bit array. 88 func (bA *CompactBitArray) Copy() *CompactBitArray { 89 if bA == nil { 90 return nil 91 } 92 c := make([]byte, len(bA.Elems)) 93 copy(c, bA.Elems) 94 return &CompactBitArray{ 95 ExtraBitsStored: bA.ExtraBitsStored, 96 Elems: c, 97 } 98 } 99 100 // String returns a string representation of CompactBitArray: BA{<bit-string>}, 101 // where <bit-string> is a sequence of 'x' (1) and '_' (0). 102 // The <bit-string> includes spaces and newlines to help people. 103 // For a simple sequence of 'x' and '_' characters with no spaces or newlines, 104 // see the MarshalJSON() method. 105 // Example: "BA{_x_}" or "nil-BitArray" for nil. 106 func (bA *CompactBitArray) String() string { 107 return bA.StringIndented("") 108 } 109 110 // StringIndented returns the same thing as String(), but applies the indent 111 // at every 10th bit, and twice at every 50th bit. 112 func (bA *CompactBitArray) StringIndented(indent string) string { 113 if bA == nil { 114 return "nil-BitArray" 115 } 116 lines := []string{} 117 bits := "" 118 size := bA.Size() 119 for i := 0; i < size; i++ { 120 if bA.GetIndex(i) { 121 bits += "x" 122 } else { 123 bits += "_" 124 } 125 if i%100 == 99 { 126 lines = append(lines, bits) 127 bits = "" 128 } 129 if i%10 == 9 { 130 bits += indent 131 } 132 if i%50 == 49 { 133 bits += indent 134 } 135 } 136 if len(bits) > 0 { 137 lines = append(lines, bits) 138 } 139 return fmt.Sprintf("BA{%v:%v}", size, strings.Join(lines, indent)) 140 } 141 142 // MarshalJSON implements json.Marshaler interface by marshaling bit array 143 // using a custom format: a string of '-' or 'x' where 'x' denotes the 1 bit. 144 func (bA *CompactBitArray) MarshalJSON() ([]byte, error) { 145 if bA == nil { 146 return []byte("null"), nil 147 } 148 149 bits := `"` 150 size := bA.Size() 151 for i := 0; i < size; i++ { 152 if bA.GetIndex(i) { 153 bits += `x` 154 } else { 155 bits += `_` 156 } 157 } 158 bits += `"` 159 return []byte(bits), nil 160 } 161 162 var bitArrayJSONRegexp = regexp.MustCompile(`\A"([_x]*)"\z`) 163 164 // UnmarshalJSON implements json.Unmarshaler interface by unmarshaling a custom 165 // JSON description. 166 func (bA *CompactBitArray) UnmarshalJSON(bz []byte) error { 167 b := string(bz) 168 if b == "null" { 169 // This is required e.g. for encoding/json when decoding 170 // into a pointer with pre-allocated BitArray. 171 bA.ExtraBitsStored = 0 172 bA.Elems = nil 173 return nil 174 } 175 176 // Validate 'b'. 177 match := bitArrayJSONRegexp.FindStringSubmatch(b) 178 if match == nil { 179 return fmt.Errorf("bitArray in JSON should be a string of format %q but got %s", bitArrayJSONRegexp.String(), b) 180 } 181 bits := match[1] 182 183 // Construct new CompactBitArray and copy over. 184 numBits := len(bits) 185 bA2 := NewCompactBitArray(numBits) 186 for i := 0; i < numBits; i++ { 187 if bits[i] == 'x' { 188 bA2.SetIndex(i, true) 189 } 190 } 191 *bA = *bA2 192 return nil 193 } 194 195 // CompactMarshal is a space efficient encoding for CompactBitArray. 196 // It is not amino compatible. 197 func (bA *CompactBitArray) CompactMarshal() []byte { 198 size := bA.Size() 199 if size <= 0 { 200 return []byte("null") 201 } 202 bz := make([]byte, 0, size/8) 203 // length prefix number of bits, not number of bytes. This difference 204 // takes 3-4 bits in encoding, as opposed to instead encoding the number of 205 // bytes (saving 3-4 bits) and including the offset as a full byte. 206 bz = appendUvarint(bz, uint64(size)) 207 bz = append(bz, bA.Elems...) 208 return bz 209 } 210 211 // CompactUnmarshal is a space efficient decoding for CompactBitArray. 212 // It is not amino compatible. 213 func CompactUnmarshal(bz []byte) (*CompactBitArray, error) { 214 if len(bz) < 2 { 215 return nil, errors.New("compact bit array: invalid compact unmarshal size") 216 } else if bytes.Equal(bz, []byte("null")) { 217 return NewCompactBitArray(0), nil 218 } 219 size, n := binary.Uvarint(bz) 220 bz = bz[n:] 221 if len(bz) != int(size+7)/8 { 222 return nil, errors.New("compact bit array: invalid compact unmarshal size") 223 } 224 225 bA := &CompactBitArray{byte(int(size % 8)), bz} 226 return bA, nil 227 } 228 229 func appendUvarint(b []byte, x uint64) []byte { 230 var a [binary.MaxVarintLen64]byte 231 n := binary.PutUvarint(a[:], x) 232 return append(b, a[:n]...) 233 }