add DecompressInto for decompression of payloads with known sizes

In some systems, compressed payloads are prefixed or tagged with the size of
the original payload. This allows allocation of a perfectly-sized buffer to
hold the decompressed payload. The current Decompress interface prohibits this
type of usage, because Decompress allocates if the destination buffer is
smaller than the worst-case bound. In practice, a perfectly-sized buffer will
always be smaller than what's required for the worst-case payload.

This commit introduces an additional DecompressInto entrypoint. DecompressInto
accepts a destination buffer and requires that the decompressed payload fit
into the buffer. If it does not, it returns the original error returned by zstd.

Author: jbowens

zstd.go

@@ -143,12 +143,7 @@ func Decompress(dst, src []byte) ([]byte, error) {
 		dst = make([]byte, bound)
 	}
 
-	written := int(C.ZSTD_decompress(
-		unsafe.Pointer(&dst[0]),
-		C.size_t(len(dst)),
-		unsafe.Pointer(&src[0]),
-		C.size_t(len(src))))
-	err := getError(written)
+	written, err := DecompressInto(dst, src)
 	if err == nil {
 		return dst[:written], nil
 	}
@@ -161,3 +156,19 @@ func Decompress(dst, src []byte) ([]byte, error) {
 	defer r.Close()
 	return ioutil.ReadAll(r)
 }
+
+// DecompressInto decompresses src into dst. Unlike Decompress, DecompressInto
+// requires that dst be sufficiently large to hold the decompressed payload.
+// DecompressInto may be used when the caller knows the size of the decompressed
+// payload before attempting decompression.
+//
+// It returns the number of bytes copied and an error if any is encountered. If
+// dst is too small, DecompressInto errors.
+func DecompressInto(dst, src []byte) (int, error) {
+	written := int(C.ZSTD_decompress(
+		unsafe.Pointer(&dst[0]),
+		C.size_t(len(dst)),
+		unsafe.Pointer(&src[0]),
+		C.size_t(len(src))))
+	return written, getError(written)
+}

zstd_test.go

@@ -87,6 +87,35 @@ func TestCompressDecompress(t *testing.T) {
 	}
 }
 
+func TestCompressDecompressInto(t *testing.T) {
+	payload := []byte("Hello World!")
+	compressed, err := Compress(make([]byte, CompressBound(len(payload))), payload)
+	if err != nil {
+		t.Fatalf("Error while compressing: %v", err)
+	}
+	t.Logf("Compressed: %v", compressed)
+
+	// We know the size of the payload; construct a buffer that perfectly fits
+	// the payload and use DecompressInto.
+	decompressed := make([]byte, len(payload))
+	if n, err := DecompressInto(decompressed, compressed); err != nil {
+		t.Fatalf("error while decompressing into buffer of size %d: %v",
+			len(decompressed), err)
+	} else if n != len(decompressed) {
+		t.Errorf("DecompressedInto = (%d, nil), want (%d, nil)", n, len(decompressed))
+	}
+	if !bytes.Equal(payload, decompressed) {
+		t.Fatalf("DecompressInto(_, Compress(_, %q)) yielded %q, want %q", payload, decompressed, payload)
+	}
+
+	// Ensure that decompressing into a buffer too small errors appropriately.
+	smallBuffer := make([]byte, len(payload)-1)
+	if _, err := DecompressInto(smallBuffer, compressed); !IsDstSizeTooSmallError(err) {
+		t.Fatalf("DecompressInto(<%d-sized buffer>, Compress(_, %q)) = %v, want 'Destination buffer is too small'",
+			len(smallBuffer), payload, err)
+	}
+}
+
 func TestCompressLevel(t *testing.T) {
 	inputs := [][]byte{
 		nil, {}, {0}, []byte("Hello World!"),