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  }