150 lines
2.7 KiB
Go
150 lines
2.7 KiB
Go
package streaming
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
)
|
|
|
|
type BitReader struct {
|
|
reader io.Reader
|
|
bitOffset int
|
|
tailByte byte
|
|
err error
|
|
}
|
|
|
|
func NewBitReader(reader io.Reader) *BitReader {
|
|
return &BitReader{reader: reader}
|
|
}
|
|
|
|
func (r *BitReader) ReadBool() bool {
|
|
return r.ReadUint(1) != 0
|
|
}
|
|
|
|
func (r *BitReader) ReadIntPacked(countPacked int) int64 {
|
|
if r.err != nil {
|
|
return 0
|
|
}
|
|
|
|
count, err := unpackSize(countPacked)
|
|
if err != nil {
|
|
r.err = err
|
|
return 0
|
|
}
|
|
|
|
return r.ReadInt(count)
|
|
}
|
|
|
|
func (r *BitReader) ReadInt(count int) int64 {
|
|
return twosComplement(r.ReadUint(count), count)
|
|
}
|
|
|
|
func (r *BitReader) ReadUintPacked(countPacked int) uint64 {
|
|
if r.err != nil {
|
|
return 0
|
|
}
|
|
|
|
count, err := unpackSize(countPacked)
|
|
if err != nil {
|
|
r.err = err
|
|
return 0
|
|
}
|
|
|
|
return r.ReadUint(count)
|
|
}
|
|
|
|
func (r *BitReader) ReadUint(count int) uint64 {
|
|
if r.err != nil {
|
|
return 0
|
|
}
|
|
|
|
buffer, bitOffset, err := r.readBytes(count)
|
|
if err != nil {
|
|
r.err = err
|
|
return 0
|
|
}
|
|
|
|
return readBits(buffer, bitOffset, count)
|
|
}
|
|
|
|
func (r *BitReader) Error() error {
|
|
return r.err
|
|
}
|
|
|
|
func (r *BitReader) readBytes(count int) ([]byte, int, error) {
|
|
if count == 0 {
|
|
return nil, 0, nil
|
|
}
|
|
|
|
var (
|
|
bitOffsetInByte = r.bitOffset % 8
|
|
bitsLeftInByte = 8 - bitOffsetInByte
|
|
bytesNeeded = 1
|
|
)
|
|
|
|
if bitsLeftInByte < count {
|
|
bitsOverrun := count - bitsLeftInByte
|
|
bytesNeeded += bitsOverrun / 8
|
|
if bitsOverrun%8 != 0 {
|
|
bytesNeeded++
|
|
}
|
|
}
|
|
|
|
buffer := make([]byte, bytesNeeded)
|
|
bufferToRead := buffer
|
|
if bitsLeftInByte < 8 {
|
|
buffer[0] = r.tailByte
|
|
bufferToRead = buffer[1:]
|
|
}
|
|
|
|
if _, err := io.ReadAtLeast(r.reader, bufferToRead, len(bufferToRead)); err != nil {
|
|
return nil, 0, err
|
|
}
|
|
|
|
r.bitOffset += count
|
|
r.tailByte = buffer[bytesNeeded-1]
|
|
|
|
return buffer, bitOffsetInByte, nil
|
|
}
|
|
|
|
func readBits(buffer []byte, bitOffset, count int) uint64 {
|
|
var result uint64
|
|
|
|
remainder := count
|
|
for byteOffset := 0; remainder > 0; byteOffset++ {
|
|
bitsRead := 8 - bitOffset
|
|
if bitsRead > remainder {
|
|
bitsRead = remainder
|
|
}
|
|
|
|
bufferByte := buffer[byteOffset]
|
|
bufferByte >>= uint(bitOffset)
|
|
bufferByte &= ^(0xff << uint(bitsRead))
|
|
|
|
result |= (uint64(bufferByte) << uint(count-remainder))
|
|
|
|
remainder -= bitsRead
|
|
bitOffset = 0
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
func twosComplement(value uint64, bits int) int64 {
|
|
signMask := uint64(1 << uint(bits-1))
|
|
if value&signMask == 0 {
|
|
return int64(value &^ signMask)
|
|
} else {
|
|
valueMask := ^(^uint64(0) << uint(bits-1))
|
|
return -int64(valueMask & (^value + 1))
|
|
}
|
|
}
|
|
|
|
func unpackSize(sizePacked int) (int, error) {
|
|
sizes := []int{0, 1, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 26, 28, 30, 32}
|
|
if sizePacked >= len(sizes) {
|
|
return 0, errors.New("invalid packed size")
|
|
}
|
|
|
|
return sizes[sizePacked], nil
|
|
}
|