diff --git a/kaitai/seekbuf.go b/kaitai/seekbuf.go new file mode 100644 index 0000000..14fccaa --- /dev/null +++ b/kaitai/seekbuf.go @@ -0,0 +1,77 @@ +package kaitai + +import ( + "errors" + "io" +) + +// SeekableBuffer is an in-memory buffer that implements io.ReadWriteSeeker. +type SeekableBuffer struct { + data []byte + pos int64 +} + +// NewSeekableBuffer creates a SeekableBuffer with the given initial data. +func NewSeekableBuffer(data []byte) *SeekableBuffer { + return &SeekableBuffer{data: data} +} + +// NewSeekableBufferSize creates a SeekableBuffer pre-allocated to size bytes. +func NewSeekableBufferSize(size int) *SeekableBuffer { + return &SeekableBuffer{data: make([]byte, size)} +} + +func (sb *SeekableBuffer) Read(p []byte) (n int, err error) { + if sb.pos >= int64(len(sb.data)) { + return 0, io.EOF + } + n = copy(p, sb.data[sb.pos:]) + sb.pos += int64(n) + return n, nil +} + +func (sb *SeekableBuffer) Write(p []byte) (n int, err error) { + end := sb.pos + int64(len(p)) + if end > int64(len(sb.data)) { + // Grow buffer + if end > int64(cap(sb.data)) { + newData := make([]byte, end, end*2) + copy(newData, sb.data) + sb.data = newData + } else { + sb.data = sb.data[:end] + } + } + n = copy(sb.data[sb.pos:], p) + sb.pos += int64(n) + return n, nil +} + +func (sb *SeekableBuffer) Seek(offset int64, whence int) (int64, error) { + var newPos int64 + switch whence { + case io.SeekStart: + newPos = offset + case io.SeekCurrent: + newPos = sb.pos + offset + case io.SeekEnd: + newPos = int64(len(sb.data)) + offset + default: + return 0, errors.New("SeekableBuffer.Seek: invalid whence") + } + if newPos < 0 { + return 0, errors.New("SeekableBuffer.Seek: negative position") + } + sb.pos = newPos + return newPos, nil +} + +// Bytes returns the buffer contents. +func (sb *SeekableBuffer) Bytes() []byte { + return sb.data +} + +// Len returns the current length of the buffer. +func (sb *SeekableBuffer) Len() int { + return len(sb.data) +} diff --git a/kaitai/seekbuf_test.go b/kaitai/seekbuf_test.go new file mode 100644 index 0000000..428ab09 --- /dev/null +++ b/kaitai/seekbuf_test.go @@ -0,0 +1,168 @@ +package kaitai + +import ( + "bytes" + "errors" + "io" + "testing" +) + +func TestNewSeekableBuffer(t *testing.T) { + data := []byte{0x01, 0x02, 0x03} + sb := NewSeekableBuffer(data) + if !bytes.Equal(sb.Bytes(), data) { + t.Errorf("Bytes() = %v, want %v", sb.Bytes(), data) + } + if sb.Len() != len(data) { + t.Errorf("Len() = %d, want %d", sb.Len(), len(data)) + } + pos, err := sb.Seek(0, io.SeekCurrent) + if err != nil { + t.Fatalf("Seek returned error: %v", err) + } + if pos != 0 { + t.Errorf("initial position = %d, want 0", pos) + } +} + +func TestNewSeekableBufferSize(t *testing.T) { + sb := NewSeekableBufferSize(16) + if sb.Len() != 16 { + t.Errorf("Len() = %d, want 16", sb.Len()) + } + for i, b := range sb.Bytes() { + if b != 0 { + t.Errorf("Bytes()[%d] = %#x, want 0", i, b) + } + } +} + +func TestSeekableBuffer_Read(t *testing.T) { + sb := NewSeekableBuffer([]byte{0x10, 0x20, 0x30, 0x40}) + + out := make([]byte, 2) + n, err := sb.Read(out) + if err != nil { + t.Fatalf("Read error: %v", err) + } + if n != 2 || !bytes.Equal(out, []byte{0x10, 0x20}) { + t.Errorf("first Read = (%d, %v), want (2, [10 20])", n, out) + } + + out = make([]byte, 8) + n, err = sb.Read(out) + if err != nil { + t.Fatalf("second Read error: %v", err) + } + if n != 2 || !bytes.Equal(out[:n], []byte{0x30, 0x40}) { + t.Errorf("second Read = (%d, %v), want (2, [30 40])", n, out[:n]) + } + + n, err = sb.Read(out) + if !errors.Is(err, io.EOF) { + t.Errorf("expected io.EOF at end, got n=%d err=%v", n, err) + } +} + +func TestSeekableBuffer_Write(t *testing.T) { + sb := NewSeekableBuffer(nil) + n, err := sb.Write([]byte{0xAA, 0xBB}) + if err != nil { + t.Fatalf("Write error: %v", err) + } + if n != 2 { + t.Errorf("Write n = %d, want 2", n) + } + if !bytes.Equal(sb.Bytes(), []byte{0xAA, 0xBB}) { + t.Errorf("Bytes() = %v, want [AA BB]", sb.Bytes()) + } + + _, err = sb.Write([]byte{0xCC, 0xDD, 0xEE}) + if err != nil { + t.Fatalf("Write error: %v", err) + } + if !bytes.Equal(sb.Bytes(), []byte{0xAA, 0xBB, 0xCC, 0xDD, 0xEE}) { + t.Errorf("Bytes() = %v, want [AA BB CC DD EE]", sb.Bytes()) + } +} + +func TestSeekableBuffer_WriteOverwrite(t *testing.T) { + sb := NewSeekableBuffer([]byte{0x01, 0x02, 0x03, 0x04}) + _, err := sb.Seek(1, io.SeekStart) + if err != nil { + t.Fatalf("Seek error: %v", err) + } + _, err = sb.Write([]byte{0xFE, 0xFF}) + if err != nil { + t.Fatalf("Write error: %v", err) + } + want := []byte{0x01, 0xFE, 0xFF, 0x04} + if !bytes.Equal(sb.Bytes(), want) { + t.Errorf("Bytes() = %v, want %v", sb.Bytes(), want) + } +} + +func TestSeekableBuffer_WriteGrowAfterSeek(t *testing.T) { + sb := NewSeekableBuffer([]byte{0x01, 0x02}) + _, err := sb.Seek(1, io.SeekStart) + if err != nil { + t.Fatalf("Seek error: %v", err) + } + _, err = sb.Write([]byte{0xAA, 0xBB, 0xCC}) + if err != nil { + t.Fatalf("Write error: %v", err) + } + want := []byte{0x01, 0xAA, 0xBB, 0xCC} + if !bytes.Equal(sb.Bytes(), want) { + t.Errorf("Bytes() = %v, want %v", sb.Bytes(), want) + } +} + +func TestSeekableBuffer_Seek(t *testing.T) { + tests := []struct { + name string + offset int64 + whence int + wantPos int64 + wantErr bool + }{ + {"SeekStart", 2, io.SeekStart, 2, false}, + {"SeekCurrent forward", 1, io.SeekCurrent, 1, false}, + {"SeekEnd zero offset", 0, io.SeekEnd, 4, false}, + {"SeekStart zero", 0, io.SeekStart, 0, false}, + {"SeekStart negative", -1, io.SeekStart, 0, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sb := NewSeekableBuffer([]byte{0x01, 0x02, 0x03, 0x04}) + got, err := sb.Seek(tt.offset, tt.whence) + if (err != nil) != tt.wantErr { + t.Fatalf("Seek error = %v, wantErr %v", err, tt.wantErr) + } + if err == nil && got != tt.wantPos { + t.Errorf("Seek returned %d, want %d", got, tt.wantPos) + } + }) + } +} + +func TestSeekableBuffer_RoundTrip(t *testing.T) { + sb := NewSeekableBuffer(nil) + want := []byte{0xDE, 0xAD, 0xBE, 0xEF} + _, err := sb.Write(want) + if err != nil { + t.Fatalf("Write error: %v", err) + } + _, err = sb.Seek(0, io.SeekStart) + if err != nil { + t.Fatalf("Seek error: %v", err) + } + got := make([]byte, len(want)) + _, err = io.ReadFull(sb, got) + if err != nil { + t.Fatalf("ReadFull error: %v", err) + } + if !bytes.Equal(got, want) { + t.Errorf("round-tripped bytes = %v, want %v", got, want) + } +} diff --git a/kaitai/util.go b/kaitai/util.go index c44eaed..3c19544 100644 --- a/kaitai/util.go +++ b/kaitai/util.go @@ -54,6 +54,21 @@ func ProcessZlib(in []byte) ([]byte, error) { return res, nil } +// UnprocessZlib compresses data using zlib. +func UnprocessZlib(in []byte) ([]byte, error) { + var buf bytes.Buffer + w := zlib.NewWriter(&buf) + _, err := w.Write(in) + if err != nil { + return nil, fmt.Errorf("UnprocessZlib: error writing zlib data: %w", err) + } + err = w.Close() + if err != nil { + return nil, fmt.Errorf("UnprocessZlib: error flushing zlib stream: %w", err) + } + return buf.Bytes(), nil +} + // BytesToStr returns a string decoded by the given decoder. func BytesToStr(in []byte, decoder *encoding.Decoder) (string, error) { i := bytes.NewReader(in) diff --git a/kaitai/util_test.go b/kaitai/util_test.go index 667fcf0..4d2b112 100644 --- a/kaitai/util_test.go +++ b/kaitai/util_test.go @@ -101,6 +101,33 @@ func TestProcessZlib(t *testing.T) { } } +func TestUnprocessZlib(t *testing.T) { + tests := []struct { + name string + in []byte + }{ + {"Empty", []byte{}}, + {"Short string", []byte("goodbye, world")}, + {"Binary data", []byte{0x00, 0x01, 0x02, 0xff, 0xfe, 0xfd, 0x00, 0x00, 0x00}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + compressed, err := UnprocessZlib(tt.in) + if err != nil { + t.Fatalf("UnprocessZlib() error = %v", err) + } + // Round-trip through ProcessZlib to verify validity. + got, err := ProcessZlib(compressed) + if err != nil { + t.Fatalf("ProcessZlib() after compress error = %v", err) + } + if !reflect.DeepEqual(got, tt.in) && !(len(got) == 0 && len(tt.in) == 0) { + t.Errorf("round-trip mismatch: got %v, want %v", got, tt.in) + } + }) + } +} + func TestBytesToStr(t *testing.T) { utf16 := unicode.UTF16(unicode.BigEndian, unicode.ExpectBOM) type args struct { diff --git a/kaitai/writer.go b/kaitai/writer.go index 89e5656..49cbec9 100644 --- a/kaitai/writer.go +++ b/kaitai/writer.go @@ -2,6 +2,7 @@ package kaitai import ( "encoding/binary" + "errors" "fmt" "io" "math" @@ -12,6 +13,10 @@ type Writer struct { io.Writer buf [8]byte + + bits uint64 + bitsLeft int + bitsLe bool } // NewWriter creates and initializes a new Writer using w. @@ -19,10 +24,46 @@ func NewWriter(w io.Writer) *Writer { return &Writer{Writer: w} } +// Pos returns the current position in the stream, if the stream is seekable. +func (k *Writer) Pos() (int64, error) { + switch w := k.Writer.(type) { + case io.Seeker: + n, err := w.Seek(0, io.SeekCurrent) + if err != nil { + return 0, fmt.Errorf("Pos: failed to get pos: %w", err) + } + return n, nil + default: + return 0, errors.New("Pos: stream is not seekable") + } +} + +// Seek seeks to the given position, if the stream is seekable. +func (k *Writer) Seek(offset int64, whence int) (int64, error) { + err := k.AlignToByte() + if err != nil { + return 0, err + } + switch w := k.Writer.(type) { + case io.Seeker: + n, err := w.Seek(offset, whence) + if err != nil { + return 0, fmt.Errorf("Seek: failed to seek stream: %w", err) + } + return n, nil + default: + return 0, errors.New("Seek: stream is not seekable") + } +} + // WriteU1 writes a uint8 to the underlying writer. func (k *Writer) WriteU1(v uint8) error { + err := k.AlignToByte() + if err != nil { + return err + } k.buf[0] = v - _, err := k.Write(k.buf[:1]) + _, err = k.Write(k.buf[:1]) if err != nil { return fmt.Errorf("WriteU1: failed to write uint8: %w", err) } @@ -31,8 +72,12 @@ func (k *Writer) WriteU1(v uint8) error { // WriteU2be writes a uint16 in big-endian order to the underlying writer. func (k *Writer) WriteU2be(v uint16) error { + err := k.AlignToByte() + if err != nil { + return err + } binary.BigEndian.PutUint16(k.buf[:2], v) - _, err := k.Write(k.buf[:2]) + _, err = k.Write(k.buf[:2]) if err != nil { return fmt.Errorf("WriteU2be: failed to write uint16: %w", err) } @@ -41,8 +86,12 @@ func (k *Writer) WriteU2be(v uint16) error { // WriteU4be writes a uint32 in big-endian order to the underlying writer. func (k *Writer) WriteU4be(v uint32) error { + err := k.AlignToByte() + if err != nil { + return err + } binary.BigEndian.PutUint32(k.buf[:4], v) - _, err := k.Write(k.buf[:4]) + _, err = k.Write(k.buf[:4]) if err != nil { return fmt.Errorf("WriteU4be: failed to write uint32: %w", err) } @@ -51,8 +100,12 @@ func (k *Writer) WriteU4be(v uint32) error { // WriteU8be writes a uint64 in big-endian order to the underlying writer. func (k *Writer) WriteU8be(v uint64) error { + err := k.AlignToByte() + if err != nil { + return err + } binary.BigEndian.PutUint64(k.buf[:8], v) - _, err := k.Write(k.buf[:8]) + _, err = k.Write(k.buf[:8]) if err != nil { return fmt.Errorf("WriteU8be: failed to write uint64: %w", err) } @@ -61,8 +114,12 @@ func (k *Writer) WriteU8be(v uint64) error { // WriteU2le writes a uint16 in little-endian order to the underlying writer. func (k *Writer) WriteU2le(v uint16) error { + err := k.AlignToByte() + if err != nil { + return err + } binary.LittleEndian.PutUint16(k.buf[:2], v) - _, err := k.Write(k.buf[:2]) + _, err = k.Write(k.buf[:2]) if err != nil { return fmt.Errorf("WriteU2le: failed to write uint16: %w", err) } @@ -71,8 +128,12 @@ func (k *Writer) WriteU2le(v uint16) error { // WriteU4le writes a uint32 in little-endian order to the underlying writer. func (k *Writer) WriteU4le(v uint32) error { + err := k.AlignToByte() + if err != nil { + return err + } binary.LittleEndian.PutUint32(k.buf[:4], v) - _, err := k.Write(k.buf[:4]) + _, err = k.Write(k.buf[:4]) if err != nil { return fmt.Errorf("WriteU4le: failed to write uint32: %w", err) } @@ -81,8 +142,12 @@ func (k *Writer) WriteU4le(v uint32) error { // WriteU8le writes a uint64 in little-endian order to the underlying writer. func (k *Writer) WriteU8le(v uint64) error { + err := k.AlignToByte() + if err != nil { + return err + } binary.LittleEndian.PutUint64(k.buf[:8], v) - _, err := k.Write(k.buf[:8]) + _, err = k.Write(k.buf[:8]) if err != nil { return fmt.Errorf("WriteU8le: failed to write uint64: %w", err) } @@ -146,9 +211,145 @@ func (k *Writer) WriteF8le(v float64) error { // WriteBytes writes the byte slice b to the underlying writer. func (k *Writer) WriteBytes(b []byte) error { - _, err := k.Write(b) + err := k.AlignToByte() + if err != nil { + return err + } + _, err = k.Write(b) if err != nil { return fmt.Errorf("WriteBytes: failed to write bytes: %w", err) } return nil } + +// WriteBytesLimit writes fixed-size data with padding or terminator. +// term: terminator byte to write after data (-1 = no terminator). +// padRight: padding byte to fill remaining space (-1 = use term, or 0x00 if none) +func (k *Writer) WriteBytesLimit(data []byte, size int, term int, padRight int) error { + if len(data) > size { + data = data[:size] + } + err := k.WriteBytes(data) + if err != nil { + return fmt.Errorf("WriteBytesLimit: failed to write bytes: %w", err) + } + remaining := size - len(data) + if remaining <= 0 { + return nil + } + // Determine pad byte: if padRight is set use it, else if term is set use term, else 0 + pad := byte(0) + if padRight >= 0 { + pad = byte(padRight) + } else if term >= 0 { + pad = byte(term) + } + // Write terminator if specified + if term >= 0 && remaining > 0 { + err := k.WriteU1(byte(term)) + if err != nil { + return fmt.Errorf("WriteBytesLimit: failed to write terminator: %w", err) + } + remaining-- + } + // Fill remaining with pad byte + if remaining > 0 { + padding := make([]byte, remaining) + for i := range padding { + padding[i] = pad + } + err := k.WriteBytes(padding) + if err != nil { + return fmt.Errorf("WriteBytesLimit: failed to write padding: %w", err) + } + } + return nil +} + +// WriteBitsIntBe writes n bits in big-endian bit order. +func (k *Writer) WriteBitsIntBe(n int, val uint64) error { + k.bitsLe = false + + mask := uint64((1 << n) - 1) + val &= mask + + bitsToWrite := k.bitsLeft + n + + bytesToWrite := bitsToWrite / 8 + k.bitsLeft = bitsToWrite % 8 + + if bytesToWrite > 0 { + buf := make([]byte, bytesToWrite) + + mask := uint64((1 << k.bitsLeft) - 1) // `bitsLeft` is in range 0..7 + newBits := val & mask + val = val>>k.bitsLeft | k.bits<<(n-k.bitsLeft) + k.bits = newBits + + for i := bytesToWrite - 1; i >= 0; i-- { + buf[i] = byte(val & 0xFF) + val >>= 8 + } + _, err := k.Write(buf) + if err != nil { + return fmt.Errorf("WriteBitsIntBe: %w", err) + } + } else { + k.bits = k.bits< 0 { + buf := make([]byte, bytesToWrite) + + newBits := val >> (n - k.bitsLeft) + val = val<>= 8 + } + _, err := k.Write(buf) + if err != nil { + return fmt.Errorf("WriteBitsIntLe: %w", err) + } + } else { + k.bits |= val << oldBitsLeft + } + + var mask uint64 = (1 << k.bitsLeft) - 1 // `bitsLeft` is in range 0..7 + k.bits &= mask + + return nil +} + +// AlignToByte flushes any remaining bits, padding with zeros. +func (k *Writer) AlignToByte() error { + if k.bitsLeft > 0 { + b := k.bits + if !k.bitsLe { + b <<= 8 - k.bitsLeft + } + k.bitsLeft = 0 + k.bits = 0 + k.buf[0] = byte(b) + _, err := k.Write(k.buf[:1]) + if err != nil { + return fmt.Errorf("AlignToByte: %w", err) + } + } + return nil +} diff --git a/kaitai/writer_test.go b/kaitai/writer_test.go index 003e93f..b8a21c7 100644 --- a/kaitai/writer_test.go +++ b/kaitai/writer_test.go @@ -2,6 +2,7 @@ package kaitai import ( "bytes" + "io" "reflect" "testing" ) @@ -423,6 +424,352 @@ func TestWriter_WriteF8le(t *testing.T) { } } +func TestWriter_Pos(t *testing.T) { + t.Run("seekable", func(t *testing.T) { + sb := NewSeekableBuffer(nil) + w := NewWriter(sb) + pos, err := w.Pos() + if err != nil { + t.Fatalf("Pos() error = %v", err) + } + if pos != 0 { + t.Errorf("initial Pos() = %d, want 0", pos) + } + err = w.WriteU4be(0x12345678) + if err != nil { + t.Fatalf("WriteU4be error: %v", err) + } + pos, err = w.Pos() + if err != nil { + t.Fatalf("Pos() error = %v", err) + } + if pos != 4 { + t.Errorf("Pos() after 4-byte write = %d, want 4", pos) + } + }) + t.Run("non-seekable", func(t *testing.T) { + w := NewWriter(&bytes.Buffer{}) + _, err := w.Pos() + if err == nil { + t.Error("Pos() on non-seekable writer should return error") + } + }) +} + +func TestWriter_Seek(t *testing.T) { + t.Run("seekable", func(t *testing.T) { + sb := NewSeekableBuffer(nil) + w := NewWriter(sb) + err := w.WriteBytes([]byte{0x01, 0x02, 0x03, 0x04}) + if err != nil { + t.Fatalf("WriteBytes error: %v", err) + } + pos, err := w.Seek(1, io.SeekStart) + if err != nil { + t.Fatalf("Seek error: %v", err) + } + if pos != 1 { + t.Errorf("Seek returned %d, want 1", pos) + } + err = w.WriteU1(0xFF) + if err != nil { + t.Fatalf("WriteU1 error: %v", err) + } + want := []byte{0x01, 0xFF, 0x03, 0x04} + if !bytes.Equal(sb.Bytes(), want) { + t.Errorf("Bytes() = %v, want %v", sb.Bytes(), want) + } + }) + t.Run("non-seekable", func(t *testing.T) { + w := NewWriter(&bytes.Buffer{}) + _, err := w.Seek(0, io.SeekStart) + if err == nil { + t.Error("Seek() on non-seekable writer should return error") + } + }) +} + +func TestWriter_WriteBytesLimit(t *testing.T) { + tests := []struct { + name string + data []byte + size int + term int + padRight int + want []byte + }{ + {"shorter, zero-padded", []byte{0x01, 0x02}, 4, -1, -1, []byte{0x01, 0x02, 0x00, 0x00}}, + {"shorter, term fills", []byte{0x01}, 4, 0x55, -1, []byte{0x01, 0x55, 0x55, 0x55}}, + {"shorter, padRight only", []byte{0x01, 0x02}, 5, -1, 0xAA, []byte{0x01, 0x02, 0xAA, 0xAA, 0xAA}}, + {"shorter, term + padRight", []byte{0x01}, 4, 0x55, 0xAA, []byte{0x01, 0x55, 0xAA, 0xAA}}, + {"equal length", []byte{0x01, 0x02, 0x03}, 3, -1, -1, []byte{0x01, 0x02, 0x03}}, + {"longer, truncated", []byte{0x01, 0x02, 0x03, 0x04, 0x05}, 3, -1, -1, []byte{0x01, 0x02, 0x03}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + err := w.WriteBytesLimit(tt.data, tt.size, tt.term, tt.padRight) + if err != nil { + t.Fatalf("WriteBytesLimit error: %v", err) + } + if !bytes.Equal(buf.Bytes(), tt.want) { + t.Errorf("got %v, want %v", buf.Bytes(), tt.want) + } + }) + } +} + +func TestWriter_WriteBitsIntBe(t *testing.T) { + t.Run("two nibbles", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + err := w.WriteBitsIntBe(4, 0xA) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.WriteBitsIntBe(4, 0x5) + if err != nil { + t.Fatalf("err = %v", err) + } + if !bytes.Equal(buf.Bytes(), []byte{0xA5}) { + t.Errorf("got %v, want [A5]", buf.Bytes()) + } + }) + + t.Run("twelve bits then align", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + err := w.WriteBitsIntBe(12, 0xABC) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.AlignToByte() + if err != nil { + t.Fatalf("align err = %v", err) + } + if !bytes.Equal(buf.Bytes(), []byte{0xAB, 0xC0}) { + t.Errorf("got %v, want [AB C0]", buf.Bytes()) + } + }) + + t.Run("mask oversized value", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + // Only the low 4 bits should make it out (0xF, then 0x0 from 0x10's low nibble). + err := w.WriteBitsIntBe(4, 0xFF) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.WriteBitsIntBe(4, 0x10) + if err != nil { + t.Fatalf("err = %v", err) + } + if !bytes.Equal(buf.Bytes(), []byte{0xF0}) { + t.Errorf("got %v, want [F0]", buf.Bytes()) + } + }) + + t.Run("round trip via Stream", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + // 5 bits then 5 bits crossing a byte boundary. + err := w.WriteBitsIntBe(5, 0x15) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.WriteBitsIntBe(5, 0x0A) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.AlignToByte() + if err != nil { + t.Fatalf("align err = %v", err) + } + s := NewStream(bytes.NewReader(buf.Bytes())) + v1, err := s.ReadBitsIntBe(5) + if err != nil { + t.Fatalf("read err: %v", err) + } + v2, err := s.ReadBitsIntBe(5) + if err != nil { + t.Fatalf("read err: %v", err) + } + if v1 != 0x15 || v2 != 0x0A { + t.Errorf("round-trip: got (%#x, %#x), want (0x15, 0x0A)", v1, v2) + } + }) + + t.Run("full 64 bits from aligned state", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + val := uint64(0xDEADBEEF8BADF00D) + err := w.WriteBitsIntBe(64, val) + if err != nil { + t.Fatalf("err = %v", err) + } + want := []byte{0xDE, 0xAD, 0xBE, 0xEF, 0x8B, 0xAD, 0xF0, 0x0D} + if !bytes.Equal(buf.Bytes(), want) { + t.Errorf("got %v, want %v", buf.Bytes(), want) + } + }) + + t.Run("64 bits after partial byte", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + err := w.WriteBitsIntBe(4, 0xA) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.WriteBitsIntBe(64, 0xDEADBEEF8BADF00D) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.AlignToByte() + if err != nil { + t.Fatalf("align err = %v", err) + } + want := []byte{0xAD, 0xEA, 0xDB, 0xEE, 0xF8, 0xBA, 0xDF, 0x00, 0xD0} + if !bytes.Equal(buf.Bytes(), want) { + t.Errorf("got %v, want %v", buf.Bytes(), want) + } + }) +} + +func TestWriter_WriteBitsIntLe(t *testing.T) { + t.Run("two nibbles", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + err := w.WriteBitsIntLe(4, 0xA) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.WriteBitsIntLe(4, 0x5) + if err != nil { + t.Fatalf("err = %v", err) + } + if !bytes.Equal(buf.Bytes(), []byte{0x5A}) { + t.Errorf("got %v, want [5A]", buf.Bytes()) + } + }) + + t.Run("twelve bits then align", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + err := w.WriteBitsIntLe(12, 0xABC) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.AlignToByte() + if err != nil { + t.Fatalf("align err = %v", err) + } + if !bytes.Equal(buf.Bytes(), []byte{0xBC, 0x0A}) { + t.Errorf("got %v, want [BC 0A]", buf.Bytes()) + } + }) + + t.Run("round trip via Stream", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + err := w.WriteBitsIntLe(12, 0xABC) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.AlignToByte() + if err != nil { + t.Fatalf("align err = %v", err) + } + s := NewStream(bytes.NewReader(buf.Bytes())) + v, err := s.ReadBitsIntLe(12) + if err != nil { + t.Fatalf("read err: %v", err) + } + if v != 0xABC { + t.Errorf("round-trip got %#x, want 0xABC", v) + } + }) + + t.Run("64 bits after partial byte", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + err := w.WriteBitsIntLe(4, 0xA) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.WriteBitsIntLe(64, 0xDEADBEEF8BADF00D) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.AlignToByte() + if err != nil { + t.Fatalf("align err = %v", err) + } + s := NewStream(bytes.NewReader(buf.Bytes())) + v1, err := s.ReadBitsIntLe(4) + if err != nil { + t.Fatalf("read v1 err: %v", err) + } + v2, err := s.ReadBitsIntLe(64) + if err != nil { + t.Fatalf("read v2 err: %v", err) + } + if v1 != 0xA { + t.Errorf("v1 = %#x, want 0xA", v1) + } + if v2 != 0xDEADBEEF8BADF00D { + t.Errorf("v2 = %#x, want 0xDEADBEEF8BADF00D", v2) + } + }) +} + +func TestWriter_AlignToByte(t *testing.T) { + t.Run("no bits buffered", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + err := w.AlignToByte() + if err != nil { + t.Fatalf("err = %v", err) + } + if buf.Len() != 0 { + t.Errorf("AlignToByte with no bits wrote %d bytes, want 0", buf.Len()) + } + }) + t.Run("partial byte padded with zeros on the right", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + err := w.WriteBitsIntBe(3, 0x5) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.AlignToByte() + if err != nil { + t.Fatalf("err = %v", err) + } + if !bytes.Equal(buf.Bytes(), []byte{0xA0}) { + t.Errorf("got %v, want [A0]", buf.Bytes()) + } + }) +} + +func TestWriter_AlignToByte_LittleEndian(t *testing.T) { + t.Run("partial byte preserved in low bits", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + err := w.WriteBitsIntLe(3, 0x5) + if err != nil { + t.Fatalf("err = %v", err) + } + err = w.AlignToByte() + if err != nil { + t.Fatalf("err = %v", err) + } + if !bytes.Equal(buf.Bytes(), []byte{0x05}) { + t.Errorf("got %v, want [05]", buf.Bytes()) + } + }) +} + func TestWriter_WriteBytes(t *testing.T) { type args struct { b []byte