Allow compression level of GzipOutputStream to be configured.

3 files changed, 134 insertions, 13 deletions

diff --git a/src/google/protobuf/io/gzip_stream.cc b/src/google/protobuf/io/gzip_stream.cc
index 243b3e32..e1a35ea2 100644
--- a/src/google/protobuf/io/gzip_stream.cc
+++ b/src/google/protobuf/io/gzip_stream.cc
@@ -168,14 +168,38 @@ int64 GzipInputStream::ByteCount() const {
 
 // =========================================================================
 
+GzipOutputStream::Options::Options()
+    : format(GZIP),
+      buffer_size(kDefaultBufferSize),
+      compression_level(Z_DEFAULT_COMPRESSION),
+      compression_strategy(Z_DEFAULT_STRATEGY) {}
+
+GzipOutputStream::GzipOutputStream(ZeroCopyOutputStream* sub_stream) {
+  Init(sub_stream, Options());
+}
+
+GzipOutputStream::GzipOutputStream(ZeroCopyOutputStream* sub_stream,
+                                   const Options& options) {
+  Init(sub_stream, options);
+}
+
 GzipOutputStream::GzipOutputStream(
-    ZeroCopyOutputStream* sub_stream, Format format, int buffer_size)
-  : sub_stream_(sub_stream), sub_data_(NULL), sub_data_size_(0) {
-  if (buffer_size == -1) {
-    input_buffer_length_ = kDefaultBufferSize;
-  } else {
-    input_buffer_length_ = buffer_size;
+    ZeroCopyOutputStream* sub_stream, Format format, int buffer_size) {
+  Options options;
+  options.format = format;
+  if (buffer_size != -1) {
+    options.buffer_size = buffer_size;
   }
+  Init(sub_stream, options);
+}
+
+void GzipOutputStream::Init(ZeroCopyOutputStream* sub_stream,
+                            const Options& options) {
+  sub_stream_ = sub_stream;
+  sub_data_ = NULL;
+  sub_data_size_ = 0;
+
+  input_buffer_length_ = options.buffer_size;
   input_buffer_ = operator new(input_buffer_length_);
   GOOGLE_CHECK(input_buffer_ != NULL);
 
@@ -191,17 +215,18 @@ GzipOutputStream::GzipOutputStream(
   zcontext_.msg = NULL;
   // default to GZIP format
   int windowBitsFormat = 16;
-  if (format == ZLIB) {
+  if (options.format == ZLIB) {
     windowBitsFormat = 0;
   }
   zerror_ = deflateInit2(
       &zcontext_,
-      Z_BEST_COMPRESSION,
+      options.compression_level,
       Z_DEFLATED,
       /* windowBits */15 | windowBitsFormat,
       /* memLevel (default) */8,
-      Z_DEFAULT_STRATEGY);
+      options.compression_strategy);
 }
+
 GzipOutputStream::~GzipOutputStream() {
   Close();
   if (input_buffer_ != NULL) {
diff --git a/src/google/protobuf/io/gzip_stream.h b/src/google/protobuf/io/gzip_stream.h
index 50a2ad70..65dbc5b5 100644
--- a/src/google/protobuf/io/gzip_stream.h
+++ b/src/google/protobuf/io/gzip_stream.h
@@ -117,11 +117,39 @@ class LIBPROTOBUF_EXPORT GzipOutputStream : public ZeroCopyOutputStream {
     ZLIB = 2,
   };
 
-  // buffer_size and format may be -1 for default of 64kB and GZIP format
-  explicit GzipOutputStream(
+  struct Options {
+    // Defaults to GZIP.
+    Format format;
+
+    // What size buffer to use internally.  Defaults to 64kB.
+    int buffer_size;
+
+    // A number between 0 and 9, where 0 is no compression and 9 is best
+    // compression.  Defaults to Z_DEFAULT_COMPRESSION (see zlib.h).
+    int compression_level;
+
+    // Defaults to Z_DEFAULT_STRATEGY.  Can also be set to Z_FILTERED,
+    // Z_HUFFMAN_ONLY, or Z_RLE.  See the documentation for deflateInit2 in
+    // zlib.h for definitions of these constants.
+    int compression_strategy;
+
+    Options();  // Initializes with default values.
+  };
+
+  // Create a GzipOutputStream with default options.
+  explicit GzipOutputStream(ZeroCopyOutputStream* sub_stream);
+
+  // Create a GzipOutputStream with the given options.
+  GzipOutputStream(
       ZeroCopyOutputStream* sub_stream,
-      Format format = GZIP,
-      int buffer_size = -1);
+      const Options& options);
+
+  // DEPRECATED:  Use one of the above constructors instead.
+  GzipOutputStream(
+      ZeroCopyOutputStream* sub_stream,
+      Format format,
+      int buffer_size = -1) GOOGLE_ATTRIBUTE_DEPRECATED;
+
   virtual ~GzipOutputStream();
 
   // Return last error message or NULL if no error.
@@ -161,6 +189,9 @@ class LIBPROTOBUF_EXPORT GzipOutputStream : public ZeroCopyOutputStream {
   void* input_buffer_;
   size_t input_buffer_length_;
 
+  // Shared constructor code.
+  void Init(ZeroCopyOutputStream* sub_stream, const Options& options);
+
   // Do some compression.
   // Takes zlib flush mode.
   // Returns zlib error code.
diff --git a/src/google/protobuf/io/zero_copy_stream_unittest.cc b/src/google/protobuf/io/zero_copy_stream_unittest.cc
index 5e3310ac..c8f669a0 100644
--- a/src/google/protobuf/io/zero_copy_stream_unittest.cc
+++ b/src/google/protobuf/io/zero_copy_stream_unittest.cc
@@ -68,6 +68,7 @@
 
 #include <google/protobuf/stubs/common.h>
 #include <google/protobuf/testing/googletest.h>
+#include <google/protobuf/testing/file.h>
 #include <gtest/gtest.h>
 
 namespace google {
@@ -114,6 +115,11 @@ class IoTest : public testing::Test {
   // via WriteStuffLarge().
   void ReadStuffLarge(ZeroCopyInputStream* input);
 
+#if HAVE_ZLIB
+  string Compress(const string& data, const GzipOutputStream::Options& options);
+  string Uncompress(const string& data);
+#endif
+
   static const int kBlockSizes[];
   static const int kBlockSizeCount;
 };
@@ -366,6 +372,65 @@ TEST_F(IoTest, ZlibIoInputAutodetect) {
   }
   delete [] buffer;
 }
+
+string IoTest::Compress(const string& data,
+                        const GzipOutputStream::Options& options) {
+  string result;
+  {
+    StringOutputStream output(&result);
+    GzipOutputStream gzout(&output, options);
+    WriteToOutput(&gzout, data.data(), data.size());
+  }
+  return result;
+}
+
+string IoTest::Uncompress(const string& data) {
+  string result;
+  {
+    ArrayInputStream input(data.data(), data.size());
+    GzipInputStream gzin(&input);
+    const void* buffer;
+    int size;
+    while (gzin.Next(&buffer, &size)) {
+      result.append(reinterpret_cast<const char*>(buffer), size);
+    }
+  }
+  return result;
+}
+
+TEST_F(IoTest, CompressionOptions) {
+  // Some ad-hoc testing of compression options.
+
+  string golden;
+  File::ReadFileToStringOrDie(
+      TestSourceDir() + "/google/protobuf/testdata/golden_message", &golden);
+
+  GzipOutputStream::Options options;
+  string gzip_compressed = Compress(golden, options);
+
+  options.compression_level = 0;
+  string not_compressed = Compress(golden, options);
+
+  // Try zlib compression for fun.
+  options = GzipOutputStream::Options();
+  options.format = GzipOutputStream::ZLIB;
+  string zlib_compressed = Compress(golden, options);
+
+  // Uncompressed should be bigger than the original since it should have some
+  // sort of header.
+  EXPECT_GT(not_compressed.size(), golden.size());
+
+  // Higher compression levels should result in smaller sizes.
+  EXPECT_LT(zlib_compressed.size(), not_compressed.size());
+
+  // ZLIB format should differ from GZIP format.
+  EXPECT_TRUE(zlib_compressed != gzip_compressed);
+
+  // Everything should decompress correctly.
+  EXPECT_TRUE(Uncompress(not_compressed) == golden);
+  EXPECT_TRUE(Uncompress(gzip_compressed) == golden);
+  EXPECT_TRUE(Uncompress(zlib_compressed) == golden);
+}
 #endif
 
 // There is no string input, only string output.  Also, it doesn't support