From f9d4e395976d3a6238a6fcd0d9e642f6f84cd1a4 Mon Sep 17 00:00:00 2001
From: Jon Maltiel Swenson <jmswen@fb.com>
Date: Fri, 16 Oct 2015 09:42:43 -0700
Subject: [PATCH] Bring ZSTD support into folly/io/Compression.h

Summary: Bring ZSTD support into folly/io/Compression.h

Reviewed By: chipturner

Differential Revision: D2503151

fb-gh-sync-id: c1450cecc8197db654b7a080a63753caa666f686
---
 folly/configure.ac                |  1 +
 folly/io/Compression.cpp          | 87 +++++++++++++++++++++++++++++++
 folly/io/Compression.h            |  9 +++-
 folly/io/test/CompressionTest.cpp |  4 +-
 4 files changed, 99 insertions(+), 2 deletions(-)

diff --git a/folly/configure.ac b/folly/configure.ac
index bcfe9c94..9ab499aa 100644
--- a/folly/configure.ac
+++ b/folly/configure.ac
@@ -446,6 +446,7 @@ AC_CHECK_HEADER([lz4.h], AC_CHECK_LIB([lz4], [LZ4_decompress_safe]))
 AC_CHECK_HEADER([snappy.h], AC_CHECK_LIB([snappy], [main]))
 AC_CHECK_HEADER([zlib.h], AC_CHECK_LIB([z], [main]))
 AC_CHECK_HEADER([lzma.h], AC_CHECK_LIB([lzma], [main]))
+AC_CHECK_HEADER([zstd.h], AC_CHECK_LIB([zstd], [main]))
 
 # Include directory that contains "folly" so #include <folly/Foo.h> works
 AM_CPPFLAGS='-I$(top_srcdir)/..'
diff --git a/folly/io/Compression.cpp b/folly/io/Compression.cpp
index 97c62f07..53d36c5f 100644
--- a/folly/io/Compression.cpp
+++ b/folly/io/Compression.cpp
@@ -36,6 +36,10 @@
 #include <lzma.h>
 #endif
 
+#if FOLLY_HAVE_LIBZSTD
+#include <zstd.h>
+#endif
+
 #include <folly/Conv.h>
 #include <folly/Memory.h>
 #include <folly/Portability.h>
@@ -919,6 +923,83 @@ std::unique_ptr<IOBuf> LZMA2Codec::doUncompress(const IOBuf* data,
 
 #endif  // FOLLY_HAVE_LIBLZMA
 
+#ifdef FOLLY_HAVE_LIBZSTD
+
+/**
+ * ZSTD_BETA compression
+ */
+class ZSTDCodec final : public Codec {
+ public:
+  static std::unique_ptr<Codec> create(int level, CodecType);
+  explicit ZSTDCodec(int level, CodecType type);
+
+ private:
+  bool doNeedsUncompressedLength() const override;
+  std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
+  std::unique_ptr<IOBuf> doUncompress(
+      const IOBuf* data,
+      uint64_t uncompressedLength) override;
+};
+
+std::unique_ptr<Codec> ZSTDCodec::create(int level, CodecType type) {
+  return make_unique<ZSTDCodec>(level, type);
+}
+
+ZSTDCodec::ZSTDCodec(int level, CodecType type) : Codec(type) {
+  DCHECK(type == CodecType::ZSTD_BETA);
+}
+
+bool ZSTDCodec::doNeedsUncompressedLength() const {
+  return true;
+}
+
+std::unique_ptr<IOBuf> ZSTDCodec::doCompress(const IOBuf* data) {
+  size_t rc;
+  size_t maxCompressedLength = ZSTD_compressBound(data->length());
+  auto out = IOBuf::createCombined(maxCompressedLength);
+
+  CHECK_EQ(out->length(), 0);
+
+  rc = ZSTD_compress(
+      out->writableTail(), out->capacity(), data->data(), data->length());
+
+  if (ZSTD_isError(rc)) {
+    throw std::runtime_error(to<std::string>(
+          "ZSTD compression returned an error: ",
+          ZSTD_getErrorName(rc)));
+  }
+
+  out->append(rc);
+  CHECK_EQ(out->length(), rc);
+
+  return out;
+}
+
+std::unique_ptr<IOBuf> ZSTDCodec::doUncompress(const IOBuf* data,
+                                               uint64_t uncompressedLength) {
+  size_t rc;
+  auto out = IOBuf::createCombined(uncompressedLength);
+
+  CHECK_GE(out->capacity(), uncompressedLength);
+  CHECK_EQ(out->length(), 0);
+
+  rc = ZSTD_decompress(
+      out->writableTail(), out->capacity(), data->data(), data->length());
+
+  if (ZSTD_isError(rc)) {
+    throw std::runtime_error(to<std::string>(
+          "ZSTD decompression returned an error: ",
+          ZSTD_getErrorName(rc)));
+  }
+
+  out->append(rc);
+  CHECK_EQ(out->length(), rc);
+
+  return out;
+}
+
+#endif  // FOLLY_HAVE_LIBZSTD
+
 }  // namespace
 
 std::unique_ptr<Codec> getCodec(CodecType type, int level) {
@@ -960,6 +1041,12 @@ std::unique_ptr<Codec> getCodec(CodecType type, int level) {
     nullptr,
     nullptr,
 #endif
+
+#if FOLLY_HAVE_LIBZSTD
+    ZSTDCodec::create,
+#else
+    nullptr,
+#endif
   };
 
   size_t idx = static_cast<size_t>(type);
diff --git a/folly/io/Compression.h b/folly/io/Compression.h
index 7bbbb1c8..5bcb5177 100644
--- a/folly/io/Compression.h
+++ b/folly/io/Compression.h
@@ -73,7 +73,14 @@ enum class CodecType {
   LZMA2 = 6,
   LZMA2_VARINT_SIZE = 7,
 
-  NUM_CODEC_TYPES = 8,
+  /**
+   * Use ZSTD_BETA compression.
+   * This format is not yet final; please do not rely on it for anything other
+   * than testing purposes yet.
+   */
+  ZSTD_BETA = 8,
+
+  NUM_CODEC_TYPES = 9,
 };
 
 class Codec {
diff --git a/folly/io/test/CompressionTest.cpp b/folly/io/test/CompressionTest.cpp
index 314eba5e..1a177f25 100644
--- a/folly/io/test/CompressionTest.cpp
+++ b/folly/io/test/CompressionTest.cpp
@@ -128,6 +128,7 @@ TEST(CompressionTestNeedsUncompressedLength, Simple) {
   EXPECT_TRUE(getCodec(CodecType::LZMA2)->needsUncompressedLength());
   EXPECT_FALSE(getCodec(CodecType::LZMA2_VARINT_SIZE)
     ->needsUncompressedLength());
+  EXPECT_TRUE(getCodec(CodecType::ZSTD_BETA)->needsUncompressedLength());
 }
 
 class CompressionTest
@@ -180,7 +181,8 @@ INSTANTIATE_TEST_CASE_P(
                                      CodecType::ZLIB,
                                      CodecType::LZ4_VARINT_SIZE,
                                      CodecType::LZMA2,
-                                     CodecType::LZMA2_VARINT_SIZE)));
+                                     CodecType::LZMA2_VARINT_SIZE,
+                                     CodecType::ZSTD_BETA)));
 
 class CompressionVarintTest
     : public testing::TestWithParam<std::tr1::tuple<int, CodecType>> {
-- 
2.34.1