Support ShrinkFit option on blocks.

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
+}

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))
+				}
+			}
+		})
+	}
+}

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)

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