You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
233 lines
6.0 KiB
233 lines
6.0 KiB
package bitarray
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/binary"
|
|
"errors"
|
|
"fmt"
|
|
"regexp"
|
|
"strings"
|
|
)
|
|
|
|
// CompactBitArray is an implementation of a space efficient bit array.
|
|
// This is used to ensure that the encoded data takes up a minimal amount of
|
|
// space after amino encoding.
|
|
// This is not thread safe, and is not intended for concurrent usage.
|
|
type CompactBitArray struct {
|
|
ExtraBitsStored byte `json:"extra_bits"` // The number of extra bits in elems.
|
|
Elems []byte `json:"bits"`
|
|
}
|
|
|
|
// NewCompactBitArray returns a new compact bit array.
|
|
// It returns nil if the number of bits is zero.
|
|
func NewCompactBitArray(bits int) *CompactBitArray {
|
|
if bits <= 0 {
|
|
return nil
|
|
}
|
|
return &CompactBitArray{
|
|
ExtraBitsStored: byte(bits % 8),
|
|
Elems: make([]byte, (bits+7)/8),
|
|
}
|
|
}
|
|
|
|
// Size returns the number of bits in the bitarray
|
|
func (bA *CompactBitArray) Size() int {
|
|
if bA == nil {
|
|
return 0
|
|
} else if bA.ExtraBitsStored == byte(0) {
|
|
return len(bA.Elems) * 8
|
|
}
|
|
// num_bits = 8*num_full_bytes + overflow_in_last_byte
|
|
// num_full_bytes = (len(bA.Elems)-1)
|
|
return (len(bA.Elems)-1)*8 + int(bA.ExtraBitsStored)
|
|
}
|
|
|
|
// GetIndex returns the bit at index i within the bit array.
|
|
// The behavior is undefined if i >= bA.Size()
|
|
func (bA *CompactBitArray) GetIndex(i int) bool {
|
|
if bA == nil {
|
|
return false
|
|
}
|
|
if i >= bA.Size() {
|
|
return false
|
|
}
|
|
return bA.Elems[i>>3]&(uint8(1)<<uint8(7-(i%8))) > 0
|
|
}
|
|
|
|
// SetIndex sets the bit at index i within the bit array.
|
|
// The behavior is undefined if i >= bA.Size()
|
|
func (bA *CompactBitArray) SetIndex(i int, v bool) bool {
|
|
if bA == nil {
|
|
return false
|
|
}
|
|
if i >= bA.Size() {
|
|
return false
|
|
}
|
|
if v {
|
|
bA.Elems[i>>3] |= (uint8(1) << uint8(7-(i%8)))
|
|
} else {
|
|
bA.Elems[i>>3] &= ^(uint8(1) << uint8(7-(i%8)))
|
|
}
|
|
return true
|
|
}
|
|
|
|
// NumTrueBitsBefore returns the number of bits set to true before the
|
|
// given index. e.g. if bA = _XX__XX, NumOfTrueBitsBefore(4) = 2, since
|
|
// there are two bits set to true before index 4.
|
|
func (bA *CompactBitArray) NumTrueBitsBefore(index int) int {
|
|
numTrueValues := 0
|
|
for i := 0; i < index; i++ {
|
|
if bA.GetIndex(i) {
|
|
numTrueValues++
|
|
}
|
|
}
|
|
return numTrueValues
|
|
}
|
|
|
|
// Copy returns a copy of the provided bit array.
|
|
func (bA *CompactBitArray) Copy() *CompactBitArray {
|
|
if bA == nil {
|
|
return nil
|
|
}
|
|
c := make([]byte, len(bA.Elems))
|
|
copy(c, bA.Elems)
|
|
return &CompactBitArray{
|
|
ExtraBitsStored: bA.ExtraBitsStored,
|
|
Elems: c,
|
|
}
|
|
}
|
|
|
|
// String returns a string representation of CompactBitArray: BA{<bit-string>},
|
|
// where <bit-string> is a sequence of 'x' (1) and '_' (0).
|
|
// The <bit-string> includes spaces and newlines to help people.
|
|
// For a simple sequence of 'x' and '_' characters with no spaces or newlines,
|
|
// see the MarshalJSON() method.
|
|
// Example: "BA{_x_}" or "nil-BitArray" for nil.
|
|
func (bA *CompactBitArray) String() string {
|
|
return bA.StringIndented("")
|
|
}
|
|
|
|
// StringIndented returns the same thing as String(), but applies the indent
|
|
// at every 10th bit, and twice at every 50th bit.
|
|
func (bA *CompactBitArray) StringIndented(indent string) string {
|
|
if bA == nil {
|
|
return "nil-BitArray"
|
|
}
|
|
lines := []string{}
|
|
bits := ""
|
|
size := bA.Size()
|
|
for i := 0; i < size; i++ {
|
|
if bA.GetIndex(i) {
|
|
bits += "x"
|
|
} else {
|
|
bits += "_"
|
|
}
|
|
if i%100 == 99 {
|
|
lines = append(lines, bits)
|
|
bits = ""
|
|
}
|
|
if i%10 == 9 {
|
|
bits += indent
|
|
}
|
|
if i%50 == 49 {
|
|
bits += indent
|
|
}
|
|
}
|
|
if len(bits) > 0 {
|
|
lines = append(lines, bits)
|
|
}
|
|
return fmt.Sprintf("BA{%v:%v}", size, strings.Join(lines, indent))
|
|
}
|
|
|
|
// MarshalJSON implements json.Marshaler interface by marshaling bit array
|
|
// using a custom format: a string of '-' or 'x' where 'x' denotes the 1 bit.
|
|
func (bA *CompactBitArray) MarshalJSON() ([]byte, error) {
|
|
if bA == nil {
|
|
return []byte("null"), nil
|
|
}
|
|
|
|
bits := `"`
|
|
size := bA.Size()
|
|
for i := 0; i < size; i++ {
|
|
if bA.GetIndex(i) {
|
|
bits += `x`
|
|
} else {
|
|
bits += `_`
|
|
}
|
|
}
|
|
bits += `"`
|
|
return []byte(bits), nil
|
|
}
|
|
|
|
var bitArrayJSONRegexp = regexp.MustCompile(`\A"([_x]*)"\z`)
|
|
|
|
// UnmarshalJSON implements json.Unmarshaler interface by unmarshaling a custom
|
|
// JSON description.
|
|
func (bA *CompactBitArray) UnmarshalJSON(bz []byte) error {
|
|
b := string(bz)
|
|
if b == "null" {
|
|
// This is required e.g. for encoding/json when decoding
|
|
// into a pointer with pre-allocated BitArray.
|
|
bA.ExtraBitsStored = 0
|
|
bA.Elems = nil
|
|
return nil
|
|
}
|
|
|
|
// Validate 'b'.
|
|
match := bitArrayJSONRegexp.FindStringSubmatch(b)
|
|
if match == nil {
|
|
return fmt.Errorf("bitArray in JSON should be a string of format %q but got %s", bitArrayJSONRegexp.String(), b)
|
|
}
|
|
bits := match[1]
|
|
|
|
// Construct new CompactBitArray and copy over.
|
|
numBits := len(bits)
|
|
bA2 := NewCompactBitArray(numBits)
|
|
for i := 0; i < numBits; i++ {
|
|
if bits[i] == 'x' {
|
|
bA2.SetIndex(i, true)
|
|
}
|
|
}
|
|
*bA = *bA2
|
|
return nil
|
|
}
|
|
|
|
// CompactMarshal is a space efficient encoding for CompactBitArray.
|
|
// It is not amino compatible.
|
|
func (bA *CompactBitArray) CompactMarshal() []byte {
|
|
size := bA.Size()
|
|
if size <= 0 {
|
|
return []byte("null")
|
|
}
|
|
bz := make([]byte, 0, size/8)
|
|
// length prefix number of bits, not number of bytes. This difference
|
|
// takes 3-4 bits in encoding, as opposed to instead encoding the number of
|
|
// bytes (saving 3-4 bits) and including the offset as a full byte.
|
|
bz = appendUvarint(bz, uint64(size))
|
|
bz = append(bz, bA.Elems...)
|
|
return bz
|
|
}
|
|
|
|
// CompactUnmarshal is a space efficient decoding for CompactBitArray.
|
|
// It is not amino compatible.
|
|
func CompactUnmarshal(bz []byte) (*CompactBitArray, error) {
|
|
if len(bz) < 2 {
|
|
return nil, errors.New("compact bit array: invalid compact unmarshal size")
|
|
} else if bytes.Equal(bz, []byte("null")) {
|
|
return NewCompactBitArray(0), nil
|
|
}
|
|
size, n := binary.Uvarint(bz)
|
|
bz = bz[n:]
|
|
if len(bz) != int(size+7)/8 {
|
|
return nil, errors.New("compact bit array: invalid compact unmarshal size")
|
|
}
|
|
|
|
bA := &CompactBitArray{byte(int(size % 8)), bz}
|
|
return bA, nil
|
|
}
|
|
|
|
func appendUvarint(b []byte, x uint64) []byte {
|
|
var a [binary.MaxVarintLen64]byte
|
|
n := binary.PutUvarint(a[:], x)
|
|
return append(b, a[:n]...)
|
|
}
|