Merge pull request #151 from EdwardX29/ctx-decompress-into

add Ctx.DecompressInto for decompression of payloads with known sizes

Author: Viq111

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 := c.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(_, Ctx.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>, Ctx.Compress(_, %q)) = %v, want 'Destination buffer is too small'",
+			len(smallBuffer), payload, err)
+	}
+}
+
 func TestCtxTooSmall(t *testing.T) {
 	ctx := NewCtx()