From 7ce1dc1acb08714f2824adbb0bc5fbf6d4a4f83f Mon Sep 17 00:00:00 2001
From: Adam Tobey <adam.tobey@datadoghq.com>
Date: Wed, 10 Jan 2024 10:04:15 -0500
Subject: [PATCH 1/2] linkname skips cgo and compile with -o 3

---
 zstd_ctx.go  | 47 ++++++++++++++++++++++++++++++++++++++++-------
 zstd_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 7 deletions(-)

diff --git a/zstd_ctx.go b/zstd_ctx.go
index 46c1976..ce18d5a 100644
--- a/zstd_ctx.go
+++ b/zstd_ctx.go
@@ -1,6 +1,7 @@
 package zstd
 
 /*
+//#cgo CFLAGS: -O3
 #include "zstd.h"
 */
 import "C"
@@ -19,6 +20,7 @@ type Ctx interface {
 
 	// CompressLevel is the same as Compress but you can pass a compression level
 	CompressLevel(dst, src []byte, level int) ([]byte, error)
+	CompressLevelDirect(dst, src []byte, level int) ([]byte, error)
 
 	// Decompress src into dst.  If you have a buffer to use, you can pass it to
 	// prevent allocation.  If it is too small, or if nil is passed, a new buffer
@@ -32,14 +34,14 @@ type ctx struct {
 }
 
 // Create a new ZStd Context.
-//  When compressing/decompressing many times, it is recommended to allocate a
-//  context just once, and re-use it for each successive compression operation.
-//  This will make workload friendlier for system's memory.
-//  Note : re-using context is just a speed / resource optimization.
-//         It doesn't change the compression ratio, which remains identical.
-//  Note 2 : In multi-threaded environments,
-//         use one different context per thread for parallel execution.
 //
+//	When compressing/decompressing many times, it is recommended to allocate a
+//	context just once, and re-use it for each successive compression operation.
+//	This will make workload friendlier for system's memory.
+//	Note : re-using context is just a speed / resource optimization.
+//	       It doesn't change the compression ratio, which remains identical.
+//	Note 2 : In multi-threaded environments,
+//	       use one different context per thread for parallel execution.
 func NewCtx() Ctx {
 	c := &ctx{
 		cctx: C.ZSTD_createCCtx(),
@@ -54,6 +56,37 @@ func (c *ctx) Compress(dst, src []byte) ([]byte, error) {
 	return c.CompressLevel(dst, src, DefaultCompression)
 }
 
+func (c *ctx) CompressLevelDirect(dst, src []byte, level int) ([]byte, error) {
+	bound := CompressBound(len(src))
+	if cap(dst) >= bound {
+		dst = dst[0:bound] // Reuse dst buffer
+	} else {
+		dst = make([]byte, bound)
+	}
+
+	// We need unsafe.Pointer(&src[0]) in the Cgo call to avoid "Go pointer to Go pointer" panics.
+	// This means we need to special case empty input. See:
+	// https://github.com/golang/go/issues/14210#issuecomment-346402945
+	var cWritten C.size_t
+	if len(src) == 0 {
+		cWritten = CompressCtx(c.cctx, unsafe.Pointer(&dst[0]), C.size_t(len(dst)), unsafe.Pointer(nil), C.size_t(0), C.int(level))
+	} else {
+		cWritten = CompressCtx(c.cctx, unsafe.Pointer(&dst[0]), C.size_t(len(dst)), unsafe.Pointer(&src[0]), C.size_t(len(src)), C.int(level))
+	}
+	written := int(cWritten)
+	// Check if the return is an Error code
+	if err := getError(written); err != nil {
+		return nil, err
+	}
+	return dst[:written], nil
+}
+
+// ZSTD_CCtx* cctx, void* dst, size_t dstCapacity, const void* src, size_t srcSize, int compressionLevel
+//
+//go:linkname CompressCtx ZSTD_compressCCtx
+//go:noescape
+func CompressCtx(ctx *C.ZSTD_CCtx, dst unsafe.Pointer, dc C.size_t, src unsafe.Pointer, sc C.size_t, level C.int) C.size_t
+
 func (c *ctx) CompressLevel(dst, src []byte, level int) ([]byte, error) {
 	bound := CompressBound(len(src))
 	if cap(dst) >= bound {
diff --git a/zstd_test.go b/zstd_test.go
index 2a8455e..1793dfc 100644
--- a/zstd_test.go
+++ b/zstd_test.go
@@ -341,6 +341,34 @@ func TestSmallPayload(t *testing.T) {
 
 }
 
+func BenchmarkCompress(b *testing.B) {
+	data := make([]byte, 1024)
+	for i := range data {
+		data[i] = byte(i)
+	}
+	dest := make([]byte, CompressBound(len(data)))
+	b.Run("cgo", func(b *testing.B) {
+		b.SetBytes(int64(len(data)))
+		ctx := NewCtx()
+		for i := 0; i < b.N; i++ {
+			_, err := ctx.CompressLevel(dest, data, 3)
+			if err != nil {
+				b.Fail()
+			}
+		}
+	})
+	b.Run("direct", func(b *testing.B) {
+		b.SetBytes(int64(len(data)))
+		ctx := NewCtx()
+		for i := 0; i < b.N; i++ {
+			_, err := ctx.CompressLevelDirect(dest, data, 3)
+			if err != nil {
+				b.Fail()
+			}
+		}
+	})
+}
+
 func BenchmarkCompression(b *testing.B) {
 	if raw == nil {
 		b.Fatal(ErrNoPayloadEnv)

From 1ab59b19cd468f6cb0be8a25f99e4f71b81e49f3 Mon Sep 17 00:00:00 2001
From: Adam Tobey <adam.tobey@datadoghq.com>
Date: Wed, 10 Jan 2024 10:07:13 -0500
Subject: [PATCH 2/2] enable -o 3

---
 zstd_ctx.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zstd_ctx.go b/zstd_ctx.go
index ce18d5a..cf27d75 100644
--- a/zstd_ctx.go
+++ b/zstd_ctx.go
@@ -1,7 +1,7 @@
 package zstd
 
 /*
-//#cgo CFLAGS: -O3
+#cgo CFLAGS: -O3
 #include "zstd.h"
 */
 import "C"