diff --git a/internal/pkg/blk/shrink.go b/internal/pkg/blk/shrink.go new file mode 100644 index 0000000..c8249cc --- /dev/null +++ b/internal/pkg/blk/shrink.go @@ -0,0 +1,10 @@ +package blk + +func ShrinkToFit(buf []byte, n int) []byte { + if n >= len(buf) { + return buf + } + nBuf := make([]byte, n) + copy(nBuf, buf[:n]) + return nBuf +} diff --git a/internal/pkg/blk/shrink_test.go b/internal/pkg/blk/shrink_test.go new file mode 100644 index 0000000..9381550 --- /dev/null +++ b/internal/pkg/blk/shrink_test.go @@ -0,0 +1,68 @@ +package blk + +import ( + "bytes" + "testing" +) + +func TestShrinkToFit(t *testing.T) { + tests := []struct { + name string + buf []byte + n int + want []byte + sameUnderlying bool + }{ + { + name: "no shrink when n equals len", + buf: []byte{1, 2, 3}, + n: 3, + want: []byte{1, 2, 3}, + sameUnderlying: true, + }, + { + name: "shrink when n less than len", + buf: []byte{1, 2, 3, 4, 5}, + n: 3, + want: []byte{1, 2, 3}, + sameUnderlying: false, + }, + { + name: "shrink to zero length", + buf: []byte{1, 2, 3}, + n: 0, + want: []byte{}, + sameUnderlying: false, + }, + { + name: "n greater than len returns original", + buf: []byte{1, 2}, + n: 3, + want: []byte{1, 2}, + sameUnderlying: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := ShrinkToFit(tt.buf, tt.n) + + if !bytes.Equal(got, tt.want) { + t.Fatalf("unexpected contents: got %v, want %v", got, tt.want) + } + + if tt.sameUnderlying { + if len(got) > 0 && len(tt.buf) > 0 && &got[0] != &tt.buf[0] { + t.Fatalf("expected result to share underlying array with input") + } + } else { + if len(got) > 0 && len(tt.buf) > 0 && &got[0] == &tt.buf[0] { + t.Fatalf("expected result to have new underlying array") + } + if cap(got) != len(got) { + t.Fatalf("expected shrunk slice to have cap == len, got len=%d cap=%d", len(got), cap(got)) + } + } + }) + } +} diff --git a/internal/test/block_test.go b/internal/test/block_test.go index 6401b70..a6d50fb 100644 --- a/internal/test/block_test.go +++ b/internal/test/block_test.go @@ -51,6 +51,57 @@ func TestCompressDecompressBlockBasic(t *testing.T) { } } +// Verify that WithBlockShrinkToFit shrinks the returned compressed buffer +// while preserving the compressed contents. +func TestCompressBlockWithShrinkToFit(t *testing.T) { + defer testBorrowed(t) + + src := make([]byte, 4<<10) + if _, err := rand.Read(src); err != nil { + t.Fatalf("rand.Read failed: %v", err) + } + + // Use an oversized destination buffer so that the compressed size is + // guaranteed to be smaller than len(dst), regardless of whether the + // data is compressible. + dst := make([]byte, plz4.CompressBlockBound(len(src))*2) + + cmpNoShrink, err := plz4.CompressBlock(src, plz4.WithBlockDst(dst)) + if err != nil { + t.Fatalf("CompressBlock (no shrink) failed: %v", err) + } + + cmpShrink, err := plz4.CompressBlock(src, plz4.WithBlockDst(dst), plz4.WithBlockShrinkToFit(true)) + if err != nil { + t.Fatalf("CompressBlock (shrink) failed: %v", err) + } + + if !bytes.Equal(cmpNoShrink, cmpShrink) { + t.Fatalf("compressed data mismatch between shrink and no-shrink paths") + } + if len(cmpNoShrink) == 0 || len(cmpShrink) == 0 { + t.Fatalf("expected non-empty compressed buffers") + } + + // Without shrink, the result should reuse the provided destination buffer + // and retain its capacity. + if &cmpNoShrink[0] != &dst[0] { + t.Fatalf("expected no-shrink compressed buffer to share underlying array with dst") + } + if cap(cmpNoShrink) != cap(dst) { + t.Fatalf("expected no-shrink compressed buffer to retain dst capacity: got %d, want %d", cap(cmpNoShrink), cap(dst)) + } + + // With shrink, the result should use a new underlying buffer sized to + // the actual compressed length. + if &cmpShrink[0] == &dst[0] { + t.Fatalf("expected shrink compressed buffer to use a new underlying array") + } + if cap(cmpShrink) != len(cmpShrink) { + t.Fatalf("expected shrink compressed buffer to have cap == len, got len=%d cap=%d", len(cmpShrink), cap(cmpShrink)) + } +} + // Verify that compression level option does not break round-trip. func TestCompressDecompressBlockWithLevel(t *testing.T) { defer testBorrowed(t) @@ -81,6 +132,93 @@ func TestCompressDecompressBlockWithLevel(t *testing.T) { } } +// Verify that WithBlockShrinkToFit shrinks the decompressed buffer when a +// destination slice is provided, while still round-tripping the data. +func TestDecompressBlockWithShrinkToFit(t *testing.T) { + defer testBorrowed(t) + + src := make([]byte, 4<<10) + if _, err := rand.Read(src); err != nil { + t.Fatalf("rand.Read failed: %v", err) + } + + cmp, err := plz4.CompressBlock(src) + if err != nil { + t.Fatalf("CompressBlock failed: %v", err) + } + + // Use an oversized destination buffer to make shrink/no-shrink behavior visible. + dst := make([]byte, len(src)*2) + + decNoShrink, err := plz4.DecompressBlock(cmp, plz4.WithBlockDst(dst)) + if err != nil { + t.Fatalf("DecompressBlock (no shrink) failed: %v", err) + } + if !bytes.Equal(src, decNoShrink) { + t.Fatalf("no-shrink round-trip mismatch: got %d bytes, want %d", len(decNoShrink), len(src)) + } + if len(decNoShrink) == 0 { + t.Fatalf("expected non-empty decompressed buffer") + } + if &decNoShrink[0] != &dst[0] { + t.Fatalf("expected no-shrink result to share underlying array with provided dst") + } + if cap(decNoShrink) != cap(dst) { + t.Fatalf("expected no-shrink result to retain dst capacity: got %d, want %d", cap(decNoShrink), cap(dst)) + } + + decShrink, err := plz4.DecompressBlock(cmp, plz4.WithBlockDst(dst), plz4.WithBlockShrinkToFit(true)) + if err != nil { + t.Fatalf("DecompressBlock (shrink) failed: %v", err) + } + if !bytes.Equal(src, decShrink) { + t.Fatalf("shrink round-trip mismatch: got %d bytes, want %d", len(decShrink), len(src)) + } + if len(decShrink) == 0 { + t.Fatalf("expected non-empty decompressed buffer with shrink") + } + if &decShrink[0] == &dst[0] { + t.Fatalf("expected shrink result to use a new underlying array") + } + if cap(decShrink) != len(decShrink) { + t.Fatalf("expected shrink decompressed buffer to have cap == len, got len=%d cap=%d", len(decShrink), cap(decShrink)) + } +} + +// Verify that WithBlockShrinkToFit shrinks the decompressed buffer even when +// no destination buffer is provided (dst == nil). +func TestDecompressBlockWithShrinkToFitNoDst(t *testing.T) { + defer testBorrowed(t) + + src := make([]byte, 4<<10) + if _, err := rand.Read(src); err != nil { + t.Fatalf("rand.Read failed: %v", err) + } + + cmp, err := plz4.CompressBlock(src) + if err != nil { + t.Fatalf("CompressBlock failed: %v", err) + } + + decShrink, err := plz4.DecompressBlock(cmp, plz4.WithBlockShrinkToFit(true)) + if err != nil { + t.Fatalf("DecompressBlock (shrink, no dst) failed: %v", err) + } + + if !bytes.Equal(src, decShrink) { + t.Fatalf("shrink round-trip mismatch (no dst): got %d bytes, want %d", len(decShrink), len(src)) + } + if len(decShrink) == 0 { + t.Fatalf("expected non-empty decompressed buffer with shrink (no dst)") + } + // When no dst is provided, DecompressBlock allocates an internal buffer. + // WithBlockShrinkToFit(true) should return a slice that is tightly sized + // to the decompressed length. + if cap(decShrink) != len(decShrink) { + t.Fatalf("expected shrink decompressed buffer (no dst) to have cap == len, got len=%d cap=%d", len(decShrink), cap(decShrink)) + } +} + // Verify that providing a dictionary option is accepted and preserves round-trip. func TestCompressDecompressBlockWithDictionary(t *testing.T) { maybeSkip(t) diff --git a/plz4_block.go b/plz4_block.go index 810a1f5..def3286 100644 --- a/plz4_block.go +++ b/plz4_block.go @@ -1,6 +1,7 @@ package plz4 import ( + "github.com/prequel-dev/plz4/internal/pkg/blk" "github.com/prequel-dev/plz4/internal/pkg/compress" ) @@ -13,9 +14,10 @@ const ( type BlockOpt func(blockOpt) blockOpt type blockOpt struct { - lvl LevelT - dst []byte - dict *compress.DictT + lvl LevelT + dst []byte + shrink bool + dict *compress.DictT } func (o blockOpt) dictData() []byte { @@ -60,6 +62,18 @@ func WithBlockDst(dst []byte) BlockOpt { } } +// If true, will shrink the output buffer to the actual output size. +// This will allocate a new buffer if necessary. +// This applies only to block compression/decompression. +func WithBlockShrinkToFit(shrink bool) BlockOpt { + { + return func(o blockOpt) blockOpt { + o.shrink = shrink + return o + } + } +} + // Returns maximum compressed block size for input size sz. func CompressBlockBound(sz int) int { return compress.CompressBound(sz) @@ -97,6 +111,10 @@ func CompressBlock(src []byte, opts ...BlockOpt) ([]byte, error) { return nil, err } + if o.shrink { + return blk.ShrinkToFit(dst, n), nil + } + return dst[:n], nil } @@ -114,10 +132,14 @@ func DecompressBlock(src []byte, opts ...BlockOpt) ([]byte, error) { if dst != nil { n, err := d.Decompress(src, dst) - if err != nil { + switch { + case err != nil: return nil, err + case o.shrink: + return blk.ShrinkToFit(dst, n), nil + default: + return dst[:n], nil } - return dst[:n], nil } // No dst provided, allocate a buffer. @@ -136,6 +158,9 @@ func DecompressBlock(src []byte, opts ...BlockOpt) ([]byte, error) { switch { case err == nil: + if o.shrink { + return blk.ShrinkToFit(dst, n), nil + } return dst[:n], nil case nTry < maxTries: nTry += 1