From 3dbdd15091ef164a44606cddf7048b9c97dbb1ad Mon Sep 17 00:00:00 2001 From: John Chadwick Date: Fri, 15 May 2026 01:34:00 -0400 Subject: [PATCH 1/5] Add more Writer functionality --- kaitai/seekbuf.go | 77 +++++++++ kaitai/seekbuf_test.go | 168 +++++++++++++++++++ kaitai/util.go | 15 ++ kaitai/util_test.go | 27 ++++ kaitai/writer.go | 187 +++++++++++++++++++++ kaitai/writer_test.go | 358 +++++++++++++++++++++++++++++++++++++++++ 6 files changed, 832 insertions(+) create mode 100644 kaitai/seekbuf.go create mode 100644 kaitai/seekbuf_test.go 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..a5a4c36 100644 --- a/kaitai/util.go +++ b/kaitai/util.go @@ -54,6 +54,21 @@ func ProcessZlib(in []byte) ([]byte, error) { return res, nil } +// ProcessZlibCompress compresses data using zlib. +func ProcessZlibCompress(in []byte) ([]byte, error) { + var buf bytes.Buffer + w := zlib.NewWriter(&buf) + _, err := w.Write(in) + if err != nil { + return nil, fmt.Errorf("ProcessZlibCompress: error writing zlib data: %w", err) + } + err = w.Close() + if err != nil { + return nil, fmt.Errorf("ProcessZlibCompress: 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..db091f2 100644 --- a/kaitai/util_test.go +++ b/kaitai/util_test.go @@ -101,6 +101,33 @@ func TestProcessZlib(t *testing.T) { } } +func TestProcessZlibCompress(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 := ProcessZlibCompress(tt.in) + if err != nil { + t.Fatalf("ProcessZlibCompress() 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..8355dad 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,9 @@ type Writer struct { io.Writer buf [8]byte + + bits uint64 + bitsLeft int } // NewWriter creates and initializes a new Writer using w. @@ -19,6 +23,38 @@ 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. +// +// Seek does not flush bits buffered by WriteBitsIntBe/Le. Call AlignToByte +// or AlignToByteLe first if a bit-level write is pending; otherwise the +// buffered bits could be emitted at the new position. +func (k *Writer) Seek(offset int64, whence int) (int64, error) { + 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 { k.buf[0] = v @@ -152,3 +188,154 @@ func (k *Writer) WriteBytes(b []byte) error { } 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.Write(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.Write(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 { + if n < 64 { + val &= (1 << uint(n)) - 1 + } + // Handle overflow: when bitsLeft + n > 64, we can't shift into w.bits. + // Flush existing bits first, then handle val directly. + if k.bitsLeft > 0 && k.bitsLeft+n > 64 { + // Flush existing bits by combining with the high bits of val + bitsNeeded := 8 - k.bitsLeft + if bitsNeeded <= n { + // Take bitsNeeded from the top of val + highBits := val >> uint(n-bitsNeeded) + b := byte((k.bits << uint(bitsNeeded)) | highBits) + err := k.WriteU1(b) + if err != nil { + return err + } + n -= bitsNeeded + if n < 64 { + val &= (1 << uint(n)) - 1 + } + k.bits = 0 + k.bitsLeft = 0 + } + } + // Now bitsLeft + n <= 64, safe to accumulate. + k.bits = (k.bits << uint(n)) | val + k.bitsLeft += n + for k.bitsLeft >= 8 { + k.bitsLeft -= 8 + b := byte(k.bits >> uint(k.bitsLeft)) + err := k.WriteU1(b) + if err != nil { + return fmt.Errorf("WriteBitsIntBe: failed to write full byte: %w", err) + } + } + if k.bitsLeft > 0 { + k.bits &= (1 << uint(k.bitsLeft)) - 1 + } else { + k.bits = 0 + } + return nil +} + +// WriteBitsIntLe writes n bits in little-endian bit order. +func (k *Writer) WriteBitsIntLe(n int, val uint64) error { + if n < 64 { + val &= (1 << uint(n)) - 1 + } + // If bitsLeft + n > 64, `val << bitsLeft` would lose the high bits of val. + // Combine the buffered bits with enough low bits of val to flush one byte, + // which makes room for the rest. bitsLeft is always in 0..7 here. + if k.bitsLeft > 0 && k.bitsLeft+n > 64 { + take := 8 - k.bitsLeft + b := byte(k.bits | (val&((1<>= uint(take) + n -= take + k.bits = 0 + k.bitsLeft = 0 + } + k.bits |= val << uint(k.bitsLeft) + k.bitsLeft += n + for k.bitsLeft >= 8 { + b := byte(k.bits & 0xff) + err := k.WriteU1(b) + if err != nil { + return fmt.Errorf("WriteBitsIntLe: failed to write full byte: %w", err) + } + k.bits >>= 8 + k.bitsLeft -= 8 + } + return nil +} + +// AlignToByte flushes any remaining bits, padding with zeros. +func (k *Writer) AlignToByte() error { + if k.bitsLeft > 0 { + b := byte(k.bits << uint(8-k.bitsLeft)) + err := k.WriteU1(b) + if err != nil { + return err + } + k.bits = 0 + k.bitsLeft = 0 + } + return nil +} + +// AlignToByteLe flushes any remaining bits in little-endian order. +func (k *Writer) AlignToByteLe() error { + if k.bitsLeft > 0 { + b := byte(k.bits & 0xff) + err := k.WriteU1(b) + if err != nil { + return err + } + k.bits = 0 + k.bitsLeft = 0 + } + return nil +} diff --git a/kaitai/writer_test.go b/kaitai/writer_test.go index 003e93f..a42b3ce 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,363 @@ 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.AlignToByteLe() + 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.AlignToByteLe() + 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.AlignToByteLe() + 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_AlignToByteLe(t *testing.T) { + t.Run("no bits buffered", func(t *testing.T) { + buf := &bytes.Buffer{} + w := NewWriter(buf) + err := w.AlignToByteLe() + if err != nil { + t.Fatalf("err = %v", err) + } + if buf.Len() != 0 { + t.Errorf("AlignToByteLe with no bits wrote %d bytes, want 0", buf.Len()) + } + }) + 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.AlignToByteLe() + 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 From b91fbdb4fc0189c9d1eb838a17a5ea90d1fbb6bd Mon Sep 17 00:00:00 2001 From: John Chadwick Date: Sun, 17 May 2026 15:59:42 -0400 Subject: [PATCH 2/5] Align bit writing impl to match Python --- kaitai/writer.go | 153 +++++++++++++++++++----------------------- kaitai/writer_test.go | 21 ++---- 2 files changed, 75 insertions(+), 99 deletions(-) diff --git a/kaitai/writer.go b/kaitai/writer.go index 8355dad..65c6bbc 100644 --- a/kaitai/writer.go +++ b/kaitai/writer.go @@ -16,6 +16,7 @@ type Writer struct { bits uint64 bitsLeft int + bitsLe bool } // NewWriter creates and initializes a new Writer using w. @@ -38,11 +39,8 @@ func (k *Writer) Pos() (int64, error) { } // Seek seeks to the given position, if the stream is seekable. -// -// Seek does not flush bits buffered by WriteBitsIntBe/Le. Call AlignToByte -// or AlignToByteLe first if a bit-level write is pending; otherwise the -// buffered bits could be emitted at the new position. func (k *Writer) Seek(offset int64, whence int) (int64, error) { + k.AlignToByte() switch w := k.Writer.(type) { case io.Seeker: n, err := w.Seek(offset, whence) @@ -57,6 +55,7 @@ func (k *Writer) Seek(offset int64, whence int) (int64, error) { // WriteU1 writes a uint8 to the underlying writer. func (k *Writer) WriteU1(v uint8) error { + k.AlignToByte() k.buf[0] = v _, err := k.Write(k.buf[:1]) if err != nil { @@ -67,6 +66,7 @@ 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 { + k.AlignToByte() binary.BigEndian.PutUint16(k.buf[:2], v) _, err := k.Write(k.buf[:2]) if err != nil { @@ -77,6 +77,7 @@ 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 { + k.AlignToByte() binary.BigEndian.PutUint32(k.buf[:4], v) _, err := k.Write(k.buf[:4]) if err != nil { @@ -87,6 +88,7 @@ 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 { + k.AlignToByte() binary.BigEndian.PutUint64(k.buf[:8], v) _, err := k.Write(k.buf[:8]) if err != nil { @@ -97,6 +99,7 @@ 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 { + k.AlignToByte() binary.LittleEndian.PutUint16(k.buf[:2], v) _, err := k.Write(k.buf[:2]) if err != nil { @@ -107,6 +110,7 @@ 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 { + k.AlignToByte() binary.LittleEndian.PutUint32(k.buf[:4], v) _, err := k.Write(k.buf[:4]) if err != nil { @@ -117,6 +121,7 @@ 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 { + k.AlignToByte() binary.LittleEndian.PutUint64(k.buf[:8], v) _, err := k.Write(k.buf[:8]) if err != nil { @@ -182,6 +187,7 @@ func (k *Writer) WriteF8le(v float64) error { // WriteBytes writes the byte slice b to the underlying writer. func (k *Writer) WriteBytes(b []byte) error { + k.AlignToByte() _, err := k.Write(b) if err != nil { return fmt.Errorf("WriteBytes: failed to write bytes: %w", err) @@ -196,7 +202,7 @@ func (k *Writer) WriteBytesLimit(data []byte, size int, term int, padRight int) if len(data) > size { data = data[:size] } - _, err := k.Write(data) + err := k.WriteBytes(data) if err != nil { return fmt.Errorf("WriteBytesLimit: failed to write bytes: %w", err) } @@ -225,7 +231,7 @@ func (k *Writer) WriteBytesLimit(data []byte, size int, term int, padRight int) for i := range padding { padding[i] = pad } - _, err := k.Write(padding) + err := k.WriteBytes(padding) if err != nil { return fmt.Errorf("WriteBytesLimit: failed to write padding: %w", err) } @@ -235,107 +241,88 @@ func (k *Writer) WriteBytesLimit(data []byte, size int, term int, padRight int) // WriteBitsIntBe writes n bits in big-endian bit order. func (k *Writer) WriteBitsIntBe(n int, val uint64) error { - if n < 64 { - val &= (1 << uint(n)) - 1 - } - // Handle overflow: when bitsLeft + n > 64, we can't shift into w.bits. - // Flush existing bits first, then handle val directly. - if k.bitsLeft > 0 && k.bitsLeft+n > 64 { - // Flush existing bits by combining with the high bits of val - bitsNeeded := 8 - k.bitsLeft - if bitsNeeded <= n { - // Take bitsNeeded from the top of val - highBits := val >> uint(n-bitsNeeded) - b := byte((k.bits << uint(bitsNeeded)) | highBits) - err := k.WriteU1(b) - if err != nil { - return err - } - n -= bitsNeeded - if n < 64 { - val &= (1 << uint(n)) - 1 - } - k.bits = 0 - k.bitsLeft = 0 + k.bitsLe = false + + mask := uint64((1 << uint(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 << uint(k.bitsLeft)) - 1) // `bitsLeft` is in range 0..7 + newBits := val & mask + val = val>>uint(k.bitsLeft) | k.bits<= 0; i-- { + buf[i] = byte(val & 0xFF) + val >>= 8 } - } - // Now bitsLeft + n <= 64, safe to accumulate. - k.bits = (k.bits << uint(n)) | val - k.bitsLeft += n - for k.bitsLeft >= 8 { - k.bitsLeft -= 8 - b := byte(k.bits >> uint(k.bitsLeft)) - err := k.WriteU1(b) + _, err := k.Write(buf) if err != nil { - return fmt.Errorf("WriteBitsIntBe: failed to write full byte: %w", err) + return fmt.Errorf("WriteBitsIntBe: %w", err) } - } - if k.bitsLeft > 0 { - k.bits &= (1 << uint(k.bitsLeft)) - 1 } else { - k.bits = 0 + k.bits = k.bits< 64, `val << bitsLeft` would lose the high bits of val. - // Combine the buffered bits with enough low bits of val to flush one byte, - // which makes room for the rest. bitsLeft is always in 0..7 here. - if k.bitsLeft > 0 && k.bitsLeft+n > 64 { - take := 8 - k.bitsLeft - b := byte(k.bits | (val&((1< 0 { + buf := make([]byte, bytesToWrite) + + newBits := val >> uint(n-k.bitsLeft) + val = val<>= 8 } - val >>= uint(take) - n -= take - k.bits = 0 - k.bitsLeft = 0 - } - k.bits |= val << uint(k.bitsLeft) - k.bitsLeft += n - for k.bitsLeft >= 8 { - b := byte(k.bits & 0xff) - err := k.WriteU1(b) + _, err := k.Write(buf) if err != nil { - return fmt.Errorf("WriteBitsIntLe: failed to write full byte: %w", err) + return fmt.Errorf("WriteBitsIntLe: %w", err) } - k.bits >>= 8 - k.bitsLeft -= 8 + } else { + k.bits |= val << uint(oldBitsLeft) } + + var mask uint64 = (1 << uint(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 := byte(k.bits << uint(8-k.bitsLeft)) - err := k.WriteU1(b) - if err != nil { - return err + b := k.bits + if !k.bitsLe { + b <<= uint(8 - k.bitsLeft) } - k.bits = 0 k.bitsLeft = 0 - } - return nil -} - -// AlignToByteLe flushes any remaining bits in little-endian order. -func (k *Writer) AlignToByteLe() error { - if k.bitsLeft > 0 { - b := byte(k.bits & 0xff) - err := k.WriteU1(b) + k.bits = 0 + k.buf[0] = byte(b) + _, err := k.Write(k.buf[:1]) if err != nil { - return err + return fmt.Errorf("AlignToByte: %w", err) } - k.bits = 0 - k.bitsLeft = 0 } return nil } diff --git a/kaitai/writer_test.go b/kaitai/writer_test.go index a42b3ce..b8a21c7 100644 --- a/kaitai/writer_test.go +++ b/kaitai/writer_test.go @@ -660,7 +660,7 @@ func TestWriter_WriteBitsIntLe(t *testing.T) { if err != nil { t.Fatalf("err = %v", err) } - err = w.AlignToByteLe() + err = w.AlignToByte() if err != nil { t.Fatalf("align err = %v", err) } @@ -676,7 +676,7 @@ func TestWriter_WriteBitsIntLe(t *testing.T) { if err != nil { t.Fatalf("err = %v", err) } - err = w.AlignToByteLe() + err = w.AlignToByte() if err != nil { t.Fatalf("align err = %v", err) } @@ -701,7 +701,7 @@ func TestWriter_WriteBitsIntLe(t *testing.T) { if err != nil { t.Fatalf("err = %v", err) } - err = w.AlignToByteLe() + err = w.AlignToByte() if err != nil { t.Fatalf("align err = %v", err) } @@ -752,18 +752,7 @@ func TestWriter_AlignToByte(t *testing.T) { }) } -func TestWriter_AlignToByteLe(t *testing.T) { - t.Run("no bits buffered", func(t *testing.T) { - buf := &bytes.Buffer{} - w := NewWriter(buf) - err := w.AlignToByteLe() - if err != nil { - t.Fatalf("err = %v", err) - } - if buf.Len() != 0 { - t.Errorf("AlignToByteLe with no bits wrote %d bytes, want 0", buf.Len()) - } - }) +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) @@ -771,7 +760,7 @@ func TestWriter_AlignToByteLe(t *testing.T) { if err != nil { t.Fatalf("err = %v", err) } - err = w.AlignToByteLe() + err = w.AlignToByte() if err != nil { t.Fatalf("err = %v", err) } From e6c2c5e8e87229febe437075fb36884366e6211f Mon Sep 17 00:00:00 2001 From: John Chadwick Date: Sun, 17 May 2026 16:12:01 -0400 Subject: [PATCH 3/5] ProcessZlibCompress -> UnprocessZlib --- kaitai/util.go | 8 ++++---- kaitai/util_test.go | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/kaitai/util.go b/kaitai/util.go index a5a4c36..3c19544 100644 --- a/kaitai/util.go +++ b/kaitai/util.go @@ -54,17 +54,17 @@ func ProcessZlib(in []byte) ([]byte, error) { return res, nil } -// ProcessZlibCompress compresses data using zlib. -func ProcessZlibCompress(in []byte) ([]byte, error) { +// 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("ProcessZlibCompress: error writing zlib data: %w", err) + return nil, fmt.Errorf("UnprocessZlib: error writing zlib data: %w", err) } err = w.Close() if err != nil { - return nil, fmt.Errorf("ProcessZlibCompress: error flushing zlib stream: %w", err) + return nil, fmt.Errorf("UnprocessZlib: error flushing zlib stream: %w", err) } return buf.Bytes(), nil } diff --git a/kaitai/util_test.go b/kaitai/util_test.go index db091f2..4d2b112 100644 --- a/kaitai/util_test.go +++ b/kaitai/util_test.go @@ -101,7 +101,7 @@ func TestProcessZlib(t *testing.T) { } } -func TestProcessZlibCompress(t *testing.T) { +func TestUnprocessZlib(t *testing.T) { tests := []struct { name string in []byte @@ -112,9 +112,9 @@ func TestProcessZlibCompress(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - compressed, err := ProcessZlibCompress(tt.in) + compressed, err := UnprocessZlib(tt.in) if err != nil { - t.Fatalf("ProcessZlibCompress() error = %v", err) + t.Fatalf("UnprocessZlib() error = %v", err) } // Round-trip through ProcessZlib to verify validity. got, err := ProcessZlib(compressed) From 8510c43e9e02ef1a12c1fd078b51121bbd805fe7 Mon Sep 17 00:00:00 2001 From: John Chadwick Date: Sun, 17 May 2026 16:39:36 -0400 Subject: [PATCH 4/5] Oops, should check error from AlignToByte --- kaitai/writer.go | 61 ++++++++++++++++++++++++++++++++++-------------- 1 file changed, 44 insertions(+), 17 deletions(-) diff --git a/kaitai/writer.go b/kaitai/writer.go index 65c6bbc..07dc4e9 100644 --- a/kaitai/writer.go +++ b/kaitai/writer.go @@ -40,7 +40,10 @@ func (k *Writer) Pos() (int64, error) { // Seek seeks to the given position, if the stream is seekable. func (k *Writer) Seek(offset int64, whence int) (int64, error) { - k.AlignToByte() + err := k.AlignToByte() + if err != nil { + return 0, err + } switch w := k.Writer.(type) { case io.Seeker: n, err := w.Seek(offset, whence) @@ -55,9 +58,12 @@ func (k *Writer) Seek(offset int64, whence int) (int64, error) { // WriteU1 writes a uint8 to the underlying writer. func (k *Writer) WriteU1(v uint8) error { - k.AlignToByte() + 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) } @@ -66,9 +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 { - k.AlignToByte() + 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) } @@ -77,9 +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 { - k.AlignToByte() + 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) } @@ -88,9 +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 { - k.AlignToByte() + 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) } @@ -99,9 +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 { - k.AlignToByte() + 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) } @@ -110,9 +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 { - k.AlignToByte() + 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) } @@ -121,9 +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 { - k.AlignToByte() + 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) } @@ -187,8 +211,11 @@ func (k *Writer) WriteF8le(v float64) error { // WriteBytes writes the byte slice b to the underlying writer. func (k *Writer) WriteBytes(b []byte) error { - k.AlignToByte() - _, 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) } From 2c896e49d7d1f1f2750fd5b8304e68f4b5eaf943 Mon Sep 17 00:00:00 2001 From: John Chadwick Date: Sun, 17 May 2026 18:48:13 -0400 Subject: [PATCH 5/5] Do not defensively cast shift operand to uint --- kaitai/writer.go | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/kaitai/writer.go b/kaitai/writer.go index 07dc4e9..49cbec9 100644 --- a/kaitai/writer.go +++ b/kaitai/writer.go @@ -270,7 +270,7 @@ func (k *Writer) WriteBytesLimit(data []byte, size int, term int, padRight int) func (k *Writer) WriteBitsIntBe(n int, val uint64) error { k.bitsLe = false - mask := uint64((1 << uint(n)) - 1) + mask := uint64((1 << n) - 1) val &= mask bitsToWrite := k.bitsLeft + n @@ -281,9 +281,9 @@ func (k *Writer) WriteBitsIntBe(n int, val uint64) error { if bytesToWrite > 0 { buf := make([]byte, bytesToWrite) - mask := uint64((1 << uint(k.bitsLeft)) - 1) // `bitsLeft` is in range 0..7 + mask := uint64((1 << k.bitsLeft) - 1) // `bitsLeft` is in range 0..7 newBits := val & mask - val = val>>uint(k.bitsLeft) | k.bits<>k.bitsLeft | k.bits<<(n-k.bitsLeft) k.bits = newBits for i := bytesToWrite - 1; i >= 0; i-- { @@ -295,7 +295,7 @@ func (k *Writer) WriteBitsIntBe(n int, val uint64) error { return fmt.Errorf("WriteBitsIntBe: %w", err) } } else { - k.bits = k.bits< 0 { buf := make([]byte, bytesToWrite) - newBits := val >> uint(n-k.bitsLeft) - val = val<> (n - k.bitsLeft) + val = val< 0 { b := k.bits if !k.bitsLe { - b <<= uint(8 - k.bitsLeft) + b <<= 8 - k.bitsLeft } k.bitsLeft = 0 k.bits = 0