fq/pkg/bitio/buffer.go

package bitio

// TODO: NewBuffer with []byte arg to save on alloc

import (
	"io"
)

// Buffer is a bitio.Reader and bitio.Writer providing a bit buffer.
// Similar to bytes.Buffer.
type Buffer struct {
	buf     []byte
	bufBits int64
	bitsOff int64
}

func (b *Buffer) Len() int64 { return b.bufBits - b.bitsOff }

func (b *Buffer) Reset() {
	b.bufBits = 0
	b.bitsOff = 0
}

// Bits return unread bits in buffer
func (b *Buffer) Bits() ([]byte, int64) {
	l := b.Len()
	buf := make([]byte, BitsByteCount(l))
	copyBufBits(buf, 0, b.buf, b.bitsOff, l, true)
	return buf, b.bufBits
}

func (b *Buffer) WriteBits(p []byte, nBits int64) (n int64, err error) {
	tBytes := BitsByteCount(b.bufBits + nBits)

	if tBytes > int64(len(b.buf)) {
		if tBytes <= int64(cap(b.buf)) {
			b.buf = b.buf[:tBytes]
		} else {
			buf := make([]byte, tBytes, tBytes*2)
			copy(buf, b.buf)
			b.buf = buf
		}
	}

	copyBufBits(b.buf, b.bufBits, p, 0, nBits, true)
	b.bufBits += nBits

	return nBits, nil
}

func (b *Buffer) empty() bool { return b.bufBits <= b.bitsOff }

func (b *Buffer) ReadBits(p []byte, nBits int64) (n int64, err error) {
	if b.empty() {
		b.Reset()
		if nBits == 0 {
			return 0, nil
		}
		return 0, io.EOF
	}

	c := nBits
	left := b.Len()
	if c > left {
		c = left
	}

	copyBufBits(p, 0, b.buf, b.bitsOff, c, true)
	b.bitsOff += c

	return c, nil
}