add Ctx.DecompressInto for decompression of payloads with known sizes

Adds DecompressInto() to zstd_context to match the functionality previously implemented in the zstd file.

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. This method provides a way to decompress directly into a pre-allocated buffer.

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.

Relates to/Follows up on: DataDog#130

Author: EdwardX29

zstd_ctx.go

@@ -24,6 +24,15 @@ type Ctx interface {
 	// prevent allocation.  If it is too small, or if nil is passed, a new buffer
 	// will be allocated and returned.
 	Decompress(dst, src []byte) ([]byte, error)
+
+	// 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.
+	DecompressInto(dst, src []byte) (int, error)
 }
 
 type ctx struct {
@@ -104,14 +113,7 @@ func (c *ctx) Decompress(dst, src []byte) ([]byte, error) {
 		dst = make([]byte, bound)
 	}
 
-	written := int(C.ZSTD_decompressDCtx(
-		c.dctx,
-		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
 	}
@@ -125,6 +127,17 @@ func (c *ctx) Decompress(dst, src []byte) ([]byte, error) {
 	return ioutil.ReadAll(r)
 }
 
+func (c *ctx) DecompressInto(dst, src []byte) (int, error) {
+	written := int(C.ZSTD_decompressDCtx(
+		c.dctx,
+		unsafe.Pointer(&dst[0]),
+		C.size_t(len(dst)),
+		unsafe.Pointer(&src[0]),
+		C.size_t(len(src))))
+	err := getError(written)
+	return written, err
+}
+
 func finalizeCtx(c *ctx) {
 	C.ZSTD_freeCCtx(c.cctx)
 	C.ZSTD_freeDCtx(c.dctx)

zstd_ctx_test.go

@@ -118,6 +118,36 @@ func TestCtxDecompressZeroLengthBuf(t *testing.T) {
 	}
 }
 
+func TestCtxCompressDecompressInto(t *testing.T) {
+	ctx := NewCtx()
+	payload := []byte("Hello World!")
+	compressed, err := ctx.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 := ctx.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("Ctx.DecompressedInto = (%d, nil), want (%d, nil)", n, len(decompressed))
+	}
+	if !bytes.Equal(payload, decompressed) {
+		t.Fatalf("Ctx.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 := ctx.DecompressInto(smallBuffer, compressed); !IsDstSizeTooSmallError(err) {
+		t.Fatalf("Ctx.DecompressInto(<%d-sized buffer>, Compress(_, %q)) = %v, want 'Destination buffer is too small'",
+			len(smallBuffer), payload, err)
+	}
+}
+
 func TestCtxTooSmall(t *testing.T) {
 	ctx := NewCtx()