forked from cerc-io/ipld-eth-server
115 lines
2.4 KiB
Go
115 lines
2.4 KiB
Go
package bitfield
|
|
|
|
// NOTE: Don't bother replacing the divisions/modulo with shifts/ands, go is smart.
|
|
|
|
import (
|
|
"math/bits"
|
|
)
|
|
|
|
// NewBitfield creates a new fixed-sized Bitfield (allocated up-front).
|
|
//
|
|
// Panics if size is not a multiple of 8.
|
|
func NewBitfield(size int) Bitfield {
|
|
if size%8 != 0 {
|
|
panic("Bitfield size must be a multiple of 8")
|
|
}
|
|
return make([]byte, size/8)
|
|
}
|
|
|
|
// FromBytes constructs a new bitfield from a serialized bitfield.
|
|
func FromBytes(size int, bits []byte) Bitfield {
|
|
bf := NewBitfield(size)
|
|
start := len(bf) - len(bits)
|
|
if start < 0 {
|
|
panic("bitfield too small")
|
|
}
|
|
copy(bf[start:], bits)
|
|
return bf
|
|
}
|
|
|
|
func (bf Bitfield) offset(i int) (uint, uint8) {
|
|
return uint(len(bf)) - (uint(i) / 8) - 1, uint8(i) % 8
|
|
}
|
|
|
|
// Bitfield is, well, a bitfield.
|
|
type Bitfield []byte
|
|
|
|
// Bytes returns the Bitfield as a byte string.
|
|
//
|
|
// This function *does not* copy.
|
|
func (bf Bitfield) Bytes() []byte {
|
|
for i, b := range bf {
|
|
if b != 0 {
|
|
return bf[i:]
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Bit returns the ith bit.
|
|
//
|
|
// Panics if the bit is out of bounds.
|
|
func (bf Bitfield) Bit(i int) bool {
|
|
idx, off := bf.offset(i)
|
|
return (bf[idx]>>off)&0x1 != 0
|
|
}
|
|
|
|
// SetBit sets the ith bit.
|
|
//
|
|
// Panics if the bit is out of bounds.
|
|
func (bf Bitfield) SetBit(i int) {
|
|
idx, off := bf.offset(i)
|
|
bf[idx] |= 1 << off
|
|
}
|
|
|
|
// UnsetBit unsets the ith bit.
|
|
//
|
|
// Panics if the bit is out of bounds.
|
|
func (bf Bitfield) UnsetBit(i int) {
|
|
idx, off := bf.offset(i)
|
|
bf[idx] &= 0xFF ^ (1 << off)
|
|
}
|
|
|
|
// SetBytes sets the bits to the given byte array.
|
|
//
|
|
// Panics if 'b' is larger than the bitfield.
|
|
func (bf Bitfield) SetBytes(b []byte) {
|
|
start := len(bf) - len(b)
|
|
if start < 0 {
|
|
panic("bitfield too small")
|
|
}
|
|
for i := range bf[:start] {
|
|
bf[i] = 0
|
|
}
|
|
copy(bf[start:], b)
|
|
}
|
|
|
|
// Ones returns the number of bits set.
|
|
func (bf Bitfield) Ones() int {
|
|
cnt := 0
|
|
for _, b := range bf {
|
|
cnt += bits.OnesCount8(b)
|
|
}
|
|
return cnt
|
|
}
|
|
|
|
// OnesBefore returns the number of bits set *before* this bit.
|
|
func (bf Bitfield) OnesBefore(i int) int {
|
|
idx, off := bf.offset(i)
|
|
cnt := bits.OnesCount8(bf[idx] << (8 - off))
|
|
for _, b := range bf[idx+1:] {
|
|
cnt += bits.OnesCount8(b)
|
|
}
|
|
return cnt
|
|
}
|
|
|
|
// OnesAfter returns the number of bits set *after* this bit.
|
|
func (bf Bitfield) OnesAfter(i int) int {
|
|
idx, off := bf.offset(i)
|
|
cnt := bits.OnesCount8(bf[idx] >> off)
|
|
for _, b := range bf[:idx] {
|
|
cnt += bits.OnesCount8(b)
|
|
}
|
|
return cnt
|
|
}
|