diff --git a/pkg/bits/reader.go b/pkg/bits/reader.go index cb58d789..7832e0f2 100644 --- a/pkg/bits/reader.go +++ b/pkg/bits/reader.go @@ -89,6 +89,25 @@ func (r *Reader) ReadBits64(n byte) (res uint64) { return } +func (r *Reader) ReadBytes(n int) (b []byte) { + if r.bits == 0 { + if r.pos+n > len(r.buf) { + r.EOF = true + return nil + } + + b = r.buf[r.pos : r.pos+n] + r.pos += n + } else { + b = make([]byte, n) + for i := 0; i < n; i++ { + b[i] = r.ReadByte() + } + } + + return +} + // ReadUEGolomb - ReadExponentialGolomb (unsigned) func (r *Reader) ReadUEGolomb() uint32 { var size byte diff --git a/pkg/bits/writer.go b/pkg/bits/writer.go index 887da3ae..a54abe55 100644 --- a/pkg/bits/writer.go +++ b/pkg/bits/writer.go @@ -2,29 +2,33 @@ package bits type Writer struct { buf []byte // total buf - byte byte // current byte + byte *byte // pointer to current byte bits byte // bits left in byte - len int // current len of buf } -func NewWriter() *Writer { - return &Writer{} +func NewWriter(buf []byte) *Writer { + return &Writer{buf: buf} +} + +//goland:noinspection GoStandardMethods +func (w *Writer) WriteByte(b byte) { + if w.bits != 0 { + w.WriteBits8(b, 8) + } + + w.buf = append(w.buf, b) } func (w *Writer) WriteBit(b byte) { if w.bits == 0 { - if w.len != 0 { - w.buf = append(w.buf, w.byte) - } - - w.byte = 0 + w.buf = append(w.buf, 0) + w.byte = &w.buf[len(w.buf)-1] w.bits = 7 - w.len++ } else { w.bits-- } - w.byte |= b << w.bits + *w.byte |= (b & 1) << w.bits } func (w *Writer) WriteBits(v uint32, n byte) { @@ -45,9 +49,39 @@ func (w *Writer) WriteBits8(v, n byte) { } } -func (w *Writer) Bytes() []byte { - if w.bits == 0 { - return w.buf +func (w *Writer) WriteAllBits(bit, n byte) { + for i := byte(0); i < n; i++ { + w.WriteBit(bit) } - return append(w.buf, w.byte) +} + +func (w *Writer) WriteUint16(v uint16) { + if w.bits != 0 { + w.WriteBits16(v, 16) + } + + w.buf = append(w.buf, byte(v>>8), byte(v)) +} + +func (w *Writer) WriteBytes(bytes ...byte) { + if w.bits != 0 { + for _, b := range bytes { + w.WriteByte(b) + } + } + + w.buf = append(w.buf, bytes...) +} + +func (w *Writer) Bytes() []byte { + return w.buf +} + +func (w *Writer) Len() int { + return len(w.buf) +} + +func (w *Writer) Reset() { + w.buf = w.buf[:0] + w.bits = 0 }