Improve bits reader and writer

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

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
 }