2020-06-08 03:29:51 +03:00
|
|
|
package bitio
|
|
|
|
|
|
|
|
import (
|
2022-01-24 23:21:48 +03:00
|
|
|
"bytes"
|
2020-06-08 03:29:51 +03:00
|
|
|
"errors"
|
2022-01-24 23:21:48 +03:00
|
|
|
"fmt"
|
2020-06-08 03:29:51 +03:00
|
|
|
"io"
|
2022-01-24 23:21:48 +03:00
|
|
|
"strings"
|
2020-06-08 03:29:51 +03:00
|
|
|
)
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// ErrOffset means seek positions is invalid
|
2021-09-29 12:51:07 +03:00
|
|
|
var ErrOffset = errors.New("invalid seek offset")
|
2022-02-09 03:22:51 +03:00
|
|
|
|
|
|
|
// ErrNegativeNBits means read tried to read negative number of bits
|
2021-09-29 12:51:07 +03:00
|
|
|
var ErrNegativeNBits = errors.New("negative number of bits")
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// Reader is something that reads bits.
|
|
|
|
// Similar to io.Reader.
|
2022-01-24 23:21:48 +03:00
|
|
|
type Reader interface {
|
|
|
|
ReadBits(p []byte, nBits int64) (n int64, err error)
|
2020-06-08 03:29:51 +03:00
|
|
|
}
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// Writer is something that writs bits.
|
|
|
|
// Similar to io.Writer.
|
2022-02-05 18:52:22 +03:00
|
|
|
type Writer interface {
|
|
|
|
WriteBits(p []byte, nBits int64) (n int64, err error)
|
|
|
|
}
|
|
|
|
|
|
|
|
// Seeker is something that seeks bits
|
2022-02-09 03:22:51 +03:00
|
|
|
// Similar to io.Seeker.
|
2022-01-24 23:21:48 +03:00
|
|
|
type Seeker interface {
|
2020-06-08 03:29:51 +03:00
|
|
|
SeekBits(bitOffset int64, whence int) (int64, error)
|
|
|
|
}
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// ReaderAt is something that reads bits at an offset.
|
|
|
|
// Similar to io.ReaderAt.
|
2022-02-05 18:52:22 +03:00
|
|
|
type ReaderAt interface {
|
|
|
|
ReadBitsAt(p []byte, nBits int64, bitOff int64) (n int64, err error)
|
|
|
|
}
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// ReadSeeker is bitio.Reader and bitio.Seeker.
|
2022-01-24 23:21:48 +03:00
|
|
|
type ReadSeeker interface {
|
|
|
|
Reader
|
|
|
|
Seeker
|
2020-06-08 03:29:51 +03:00
|
|
|
}
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// ReadAtSeeker is bitio.ReaderAt and bitio.Seeker.
|
2022-01-24 23:21:48 +03:00
|
|
|
type ReadAtSeeker interface {
|
|
|
|
ReaderAt
|
|
|
|
Seeker
|
2020-06-08 03:29:51 +03:00
|
|
|
}
|
|
|
|
|
2022-02-05 18:52:22 +03:00
|
|
|
// ReaderAtSeeker is bitio.Reader, bitio.ReaderAt and bitio.Seeker
|
2022-01-24 23:21:48 +03:00
|
|
|
type ReaderAtSeeker interface {
|
|
|
|
Reader
|
|
|
|
ReaderAt
|
|
|
|
Seeker
|
2020-06-08 03:29:51 +03:00
|
|
|
}
|
|
|
|
|
2022-05-28 14:04:04 +03:00
|
|
|
// ReadCloner is a cloneable bitio.Reader, resetting read position to start
|
|
|
|
type ReadCloner interface {
|
|
|
|
CloneReader() (Reader, error)
|
|
|
|
}
|
|
|
|
|
|
|
|
// ReadSeekerCloner is a cloneable bitio.ReadSeeker, resetting read position to start
|
|
|
|
type ReadSeekerCloner interface {
|
|
|
|
CloneReadSeeker() (ReadSeeker, error)
|
|
|
|
}
|
|
|
|
|
|
|
|
// ReadAtSeekerCloner is a cloneable bitio.ReadAtSeeker, resetting read position to start
|
|
|
|
type ReadAtSeekerCloner interface {
|
|
|
|
CloneReadAtSeeker() (ReadAtSeeker, error)
|
|
|
|
}
|
|
|
|
|
|
|
|
// CloneReaderAtSeeker is a cloneable bitio.ReaderAtSeeker, resetting read position to start
|
|
|
|
type ReaderAtSeekerCloner interface {
|
|
|
|
CloneReaderAtSeeker() (ReaderAtSeeker, error)
|
|
|
|
}
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// NewBitReader reader reading nBits bits from a []byte.
|
2022-01-24 23:21:48 +03:00
|
|
|
// If nBits is -1 all bits will be used.
|
2022-02-09 03:22:51 +03:00
|
|
|
// Similar to bytes.NewReader.
|
2022-01-24 23:21:48 +03:00
|
|
|
func NewBitReader(buf []byte, nBits int64) *SectionReader {
|
|
|
|
if nBits < 0 {
|
|
|
|
nBits = int64(len(buf)) * 8
|
|
|
|
}
|
|
|
|
return NewSectionReader(
|
|
|
|
NewIOBitReadSeeker(bytes.NewReader(buf)),
|
|
|
|
0,
|
|
|
|
nBits,
|
|
|
|
)
|
|
|
|
}
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// BitsByteCount returns smallest amount of bytes to fit nBits bits.
|
2022-01-24 23:21:48 +03:00
|
|
|
func BitsByteCount(nBits int64) int64 {
|
|
|
|
n := nBits / 8
|
|
|
|
if nBits%8 != 0 {
|
|
|
|
n++
|
|
|
|
}
|
|
|
|
return n
|
|
|
|
}
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// BytesFromBitString from []byte to bit string, ex: "0101" -> ([]byte{0x50}, 4)
|
2022-01-24 23:21:48 +03:00
|
|
|
func BytesFromBitString(s string) ([]byte, int64) {
|
|
|
|
r := len(s) % 8
|
|
|
|
bufLen := len(s) / 8
|
|
|
|
if r > 0 {
|
|
|
|
bufLen++
|
|
|
|
}
|
|
|
|
buf := make([]byte, bufLen)
|
|
|
|
|
|
|
|
for i := 0; i < len(s); i++ {
|
|
|
|
d := s[i] - '0'
|
|
|
|
if d != 0 && d != 1 {
|
|
|
|
panic(fmt.Sprintf("invalid bit string %q at index %d %q", s, i, s[i]))
|
|
|
|
}
|
|
|
|
buf[i/8] |= d << (7 - i%8)
|
|
|
|
}
|
|
|
|
|
|
|
|
return buf, int64(len(s))
|
|
|
|
}
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// BitStringFromBytes from string to []byte, ex: ([]byte{0x50}, 4) -> "0101"
|
2022-01-24 23:21:48 +03:00
|
|
|
func BitStringFromBytes(buf []byte, nBits int64) string {
|
|
|
|
sb := &strings.Builder{}
|
|
|
|
for i := int64(0); i < nBits; i++ {
|
|
|
|
if buf[i/8]&(1<<(7-i%8)) > 0 {
|
|
|
|
sb.WriteString("1")
|
|
|
|
} else {
|
|
|
|
sb.WriteString("0")
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return sb.String()
|
|
|
|
}
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// CopyBuffer bits from src to dst using provided byte buffer.
|
|
|
|
// Similar to io.CopyBuffer.
|
2022-01-24 23:21:48 +03:00
|
|
|
func CopyBuffer(dst Writer, src Reader, buf []byte) (n int64, err error) {
|
2021-09-16 13:24:53 +03:00
|
|
|
// same default size as io.Copy
|
|
|
|
if buf == nil {
|
|
|
|
buf = make([]byte, 32*1024)
|
|
|
|
}
|
2020-06-08 03:29:51 +03:00
|
|
|
var written int64
|
|
|
|
|
|
|
|
for {
|
2022-01-24 23:21:48 +03:00
|
|
|
rBits, rErr := src.ReadBits(buf, int64(len(buf))*8)
|
2020-06-08 03:29:51 +03:00
|
|
|
if rBits > 0 {
|
|
|
|
wBits, wErr := dst.WriteBits(buf, rBits)
|
2022-01-24 23:21:48 +03:00
|
|
|
written += wBits
|
2020-06-08 03:29:51 +03:00
|
|
|
if wErr != nil {
|
|
|
|
err = wErr
|
|
|
|
break
|
|
|
|
}
|
|
|
|
if rBits != wBits {
|
|
|
|
err = io.ErrShortWrite
|
|
|
|
break
|
|
|
|
}
|
|
|
|
}
|
|
|
|
if rErr != nil {
|
|
|
|
if !errors.Is(rErr, io.EOF) {
|
|
|
|
err = rErr
|
|
|
|
}
|
|
|
|
break
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
return written, err
|
|
|
|
}
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// Copy bits from src to dst.
|
|
|
|
// Similar to io.Copy.
|
2022-01-24 23:21:48 +03:00
|
|
|
func Copy(dst Writer, src Reader) (n int64, err error) {
|
2021-09-16 13:24:53 +03:00
|
|
|
return CopyBuffer(dst, src, nil)
|
|
|
|
}
|
|
|
|
|
2022-01-24 23:21:48 +03:00
|
|
|
// TODO: make faster, align and use copy()
|
2022-02-05 18:52:22 +03:00
|
|
|
func copyBufBits(dst []byte, dstStart int64, src []byte, srcStart int64, n int64, zero bool) {
|
2022-01-24 23:21:48 +03:00
|
|
|
l := n
|
|
|
|
off := int64(0)
|
|
|
|
for l > 0 {
|
|
|
|
c := int64(64)
|
|
|
|
if l < c {
|
|
|
|
c = l
|
|
|
|
}
|
|
|
|
u := Read64(src, srcStart+off, c)
|
|
|
|
Write64(u, c, dst, dstStart+off)
|
|
|
|
off += c
|
|
|
|
l -= c
|
|
|
|
}
|
|
|
|
|
|
|
|
// zero fill last bits if not aligned
|
|
|
|
e := dstStart + n
|
|
|
|
if zero && e%8 != 0 {
|
|
|
|
Write64(0, 8-(e%8), dst, e)
|
2020-06-08 03:29:51 +03:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-01-24 23:21:48 +03:00
|
|
|
// TODO: redo?
|
|
|
|
func readFull(p []byte, nBits int64, bitOff int64, fn func(p []byte, nBits int64, bitOff int64) (int64, error)) (int64, error) {
|
2021-09-29 12:51:07 +03:00
|
|
|
if nBits < 0 {
|
|
|
|
return 0, ErrNegativeNBits
|
|
|
|
}
|
|
|
|
|
2022-01-24 23:21:48 +03:00
|
|
|
readBitOffset := int64(0)
|
2020-06-08 03:29:51 +03:00
|
|
|
for readBitOffset < nBits {
|
|
|
|
byteOffset := readBitOffset / 8
|
|
|
|
byteBitsOffset := readBitOffset % 8
|
|
|
|
partialByteBitsLeft := (8 - byteBitsOffset) % 8
|
|
|
|
leftBits := nBits - readBitOffset
|
|
|
|
|
|
|
|
if partialByteBitsLeft != 0 || leftBits < 8 {
|
|
|
|
readBits := partialByteBitsLeft
|
|
|
|
if partialByteBitsLeft == 0 || leftBits < readBits {
|
|
|
|
readBits = leftBits
|
|
|
|
}
|
|
|
|
|
|
|
|
var pb [1]byte
|
2022-01-24 23:21:48 +03:00
|
|
|
rBits, err := fn(pb[:], readBits, bitOff+readBitOffset)
|
2020-06-08 03:29:51 +03:00
|
|
|
Write64(uint64(pb[0]>>(8-rBits)), rBits, p, readBitOffset)
|
|
|
|
readBitOffset += rBits
|
|
|
|
|
|
|
|
if err != nil {
|
|
|
|
return nBits - readBitOffset, err
|
|
|
|
}
|
|
|
|
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
|
2022-01-24 23:21:48 +03:00
|
|
|
rBits, err := fn(p[byteOffset:], nBits-readBitOffset, bitOff+readBitOffset)
|
2020-06-08 03:29:51 +03:00
|
|
|
|
|
|
|
readBitOffset += rBits
|
|
|
|
if err != nil {
|
|
|
|
return nBits - readBitOffset, err
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
return nBits, nil
|
|
|
|
}
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// ReadAtFull expects to read nBits from r at bitOff.
|
|
|
|
// Similar to io.ReadFull.
|
2022-01-24 23:21:48 +03:00
|
|
|
func ReadAtFull(r ReaderAt, p []byte, nBits int64, bitOff int64) (int64, error) {
|
|
|
|
return readFull(p, nBits, bitOff, func(p []byte, nBits int64, bitOff int64) (int64, error) {
|
2020-06-08 03:29:51 +03:00
|
|
|
return r.ReadBitsAt(p, nBits, bitOff)
|
|
|
|
})
|
|
|
|
}
|
|
|
|
|
2022-02-09 03:22:51 +03:00
|
|
|
// ReadFull expects to read nBits from r.
|
|
|
|
// Similar to io.ReadFull.
|
2022-01-24 23:21:48 +03:00
|
|
|
func ReadFull(r Reader, p []byte, nBits int64) (int64, error) {
|
|
|
|
return readFull(p, nBits, 0, func(p []byte, nBits int64, bitOff int64) (int64, error) {
|
2020-06-08 03:29:51 +03:00
|
|
|
return r.ReadBits(p, nBits)
|
|
|
|
})
|
|
|
|
}
|
2022-05-28 14:04:04 +03:00
|
|
|
|
|
|
|
// CloneReader clones a bitio.Reader if possible, resetting read position to start
|
|
|
|
func CloneReader(r Reader) (Reader, error) {
|
|
|
|
switch r := r.(type) {
|
|
|
|
case ReadCloner:
|
|
|
|
return r.CloneReader()
|
|
|
|
}
|
|
|
|
return nil, fmt.Errorf("can't be cloned")
|
|
|
|
}
|
|
|
|
|
|
|
|
// CloneReadSeeker clones a bitio.ReadSeeker if possible, resetting read position to start
|
|
|
|
func CloneReadSeeker(r ReadSeeker) (ReadSeeker, error) {
|
|
|
|
switch r := r.(type) {
|
|
|
|
case ReadSeekerCloner:
|
|
|
|
return r.CloneReadSeeker()
|
|
|
|
}
|
|
|
|
return nil, fmt.Errorf("can't be cloned")
|
|
|
|
}
|
|
|
|
|
|
|
|
// CloneReadAtSeeker clones a bitio.ReadAtSeeker if possible, resetting read position to start
|
|
|
|
func CloneReadAtSeeker(r ReadAtSeeker) (ReadAtSeeker, error) {
|
|
|
|
switch r := r.(type) {
|
|
|
|
case ReadAtSeekerCloner:
|
|
|
|
return r.CloneReadAtSeeker()
|
|
|
|
}
|
|
|
|
return nil, fmt.Errorf("can't be cloned")
|
|
|
|
}
|
|
|
|
|
|
|
|
// CloneReaderAtSeeker clones a bitio.ReadAtSeeker if possible, resetting read position to start
|
|
|
|
func CloneReaderAtSeeker(r ReadAtSeeker) (ReaderAtSeeker, error) {
|
|
|
|
switch r := r.(type) {
|
|
|
|
case ReaderAtSeekerCloner:
|
|
|
|
return r.CloneReaderAtSeeker()
|
|
|
|
}
|
|
|
|
return nil, fmt.Errorf("can't be cloned")
|
|
|
|
}
|