Merge "Add the brotli compressor for bsdiff" am: c2b68e2059

am: 43f212fc8b

Change-Id: I6229007c72c2c0c2ef3b9deb86b51b6be7a7c032

Android.bp

@@ -15,7 +15,7 @@
 cc_defaults {
     name: "bsdiff_defaults",
     host_supported: true,
-    static_libs: ["libbz"],
+    static_libs: ["libbz", "libbrotli"],
     // Allow internal includes to be referenced with the "bsdiff/" prefix in the
     // path.
     include_dirs: ["external"],
@@ -58,8 +58,11 @@ cc_library_static {
     defaults: ["bsdiff_defaults"],
 
     srcs: [
+        "brotli_compressor.cc",
         "bsdiff.cc",
         "bz2_compressor.cc",
+        "compressor_buffer.cc",
+        "compressor_interface.cc",
         "diff_encoder.cc",
         "patch_writer.cc",
         "patch_writer_factory.cc",
@@ -69,6 +72,7 @@ cc_library_static {
     static_libs: [
         "libdivsufsort64",
         "libdivsufsort",
+        "libbrotli",
     ],
 }
 
@@ -81,6 +85,7 @@ cc_binary {
     static_libs: [
         "libbspatch",
         "libbz",
+        "libbrotli",
     ],
 }
 
@@ -95,6 +100,7 @@ cc_binary_host {
         "libdivsufsort64",
         "libdivsufsort",
         "libbz",
+        "libbrotli",
     ],
 }
 
@@ -104,6 +110,7 @@ cc_test {
     defaults: ["bsdiff_defaults"],
     test_suites: ["device-tests"],
     srcs: [
+        "brotli_compressor_unittest.cc",
         "bsdiff_unittest.cc",
         "bspatch_unittest.cc",
         "diff_encoder_unittest.cc",
@@ -123,6 +130,7 @@ cc_test {
         "libdivsufsort64",
         "libdivsufsort",
         "libbz",
+        "libbrotli",
     ],
     target: {
         android: {

Makefile

@@ -33,8 +33,11 @@ clean:
 ### List of source files for each project. Keep in sync with the Android.mk.
 # "bsdiff" program.
 bsdiff_src_files := \
+    brotli_compressor.cc \
     bsdiff.cc \
     bz2_compressor.cc \
+    compressor_buffer.cc \
+    compressor_interface.cc \
     diff_encoder.cc \
     patch_writer.cc \
     patch_writer_factory.cc \
@@ -56,6 +59,7 @@ bspatch_src_files := \
 
 # Unit test files.
 bsdiff_common_unittests := \
+    brotli_compressor_unittest.cc \
     bsdiff_unittest.cc \
     bspatch_unittest.cc \
     diff_encoder_unittest.cc \
@@ -69,10 +73,10 @@ bsdiff_common_unittests := \
     testrunner.cc
 
 
-BSDIFF_LIBS := -lbz2 -ldivsufsort -ldivsufsort64
+BSDIFF_LIBS := -lbz2 -lbrotlienc -ldivsufsort -ldivsufsort64
 BSDIFF_OBJS := $(bsdiff_src_files:.cc=.o)
 
-BSPATCH_LIBS := -lbz2
+BSPATCH_LIBS := -lbz2 -lbrotlidec
 BSPATCH_OBJS := $(bspatch_src_files:.cc=.o)
 
 UNITTEST_LIBS = -lgmock -lgtest -lpthread

brotli_compressor.cc

@@ -0,0 +1,93 @@
+// Copyright 2017 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "bsdiff/brotli_compressor.h"
+
+#include "bsdiff/logging.h"
+
+using std::endl;
+
+namespace {
+// TODO(xunchang) set brotli compression parameters based on input options.
+const size_t kBufferSize = 1024 * 1024;
+const uint32_t kBrotliQuality = 11;
+const uint32_t kBrotliLgwin = 24;
+
+}  // namespace
+
+namespace bsdiff {
+
+BrotliCompressor::BrotliCompressor() : comp_buffer_(kBufferSize) {
+  brotli_encoder_state_ =
+      BrotliEncoderCreateInstance(nullptr, nullptr, nullptr);
+  if (!brotli_encoder_state_) {
+    LOG(ERROR) << "Failed to initialize brotli decoder state" << endl;
+  } else {
+    BrotliEncoderSetParameter(brotli_encoder_state_, BROTLI_PARAM_QUALITY,
+                              kBrotliQuality);
+    BrotliEncoderSetParameter(brotli_encoder_state_, BROTLI_PARAM_LGWIN,
+                              kBrotliLgwin);
+  }
+}
+
+BrotliCompressor::~BrotliCompressor() {
+  if (brotli_encoder_state_) {
+    BrotliEncoderDestroyInstance(brotli_encoder_state_);
+  }
+}
+
+bool BrotliCompressor::Write(const uint8_t* buf, size_t size) {
+  if (!brotli_encoder_state_)
+    return false;
+
+  const uint8_t* next_in = buf;
+  size_t avail_in = size;
+  while (avail_in > 0) {
+    size_t avail_out = kBufferSize;
+    uint8_t* next_out = comp_buffer_.buffer_data();
+    if (!BrotliEncoderCompressStream(
+            brotli_encoder_state_, BROTLI_OPERATION_PROCESS, &avail_in,
+            &next_in, &avail_out, &next_out, nullptr)) {
+      LOG(ERROR) << "BrotliCompressor failed to compress " << avail_in
+                 << " bytes of data." << endl;
+      return false;
+    }
+
+    uint64_t output_bytes = comp_buffer_.buffer_size() - avail_out;
+    if (output_bytes > 0) {
+      comp_buffer_.AddDataToChunks(output_bytes);
+    }
+  }
+  return true;
+}
+
+bool BrotliCompressor::Finish() {
+  if (!brotli_encoder_state_)
+    return false;
+
+  const uint8_t* next_in = nullptr;
+  size_t avail_in = 0;
+  while (!BrotliEncoderIsFinished(brotli_encoder_state_)) {
+    size_t avail_out = kBufferSize;
+    uint8_t* next_out = comp_buffer_.buffer_data();
+    if (!BrotliEncoderCompressStream(
+            brotli_encoder_state_, BROTLI_OPERATION_FINISH, &avail_in, &next_in,
+            &avail_out, &next_out, nullptr)) {
+      LOG(ERROR) << "BrotliCompressor failed to finish compression" << endl;
+      return false;
+    }
+
+    uint64_t output_bytes = comp_buffer_.buffer_size() - avail_out;
+    if (output_bytes > 0) {
+      comp_buffer_.AddDataToChunks(output_bytes);
+    }
+  }
+  return true;
+}
+
+const std::vector<uint8_t>& BrotliCompressor::GetCompressedData() {
+  return comp_buffer_.GetCompressedData();
+}
+
+}  // namespace bsdiff

brotli_compressor.h

@@ -0,0 +1,37 @@
+// Copyright 2017 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef _BSDIFF_BROTLI_COMPRESSOR_H_
+#define _BSDIFF_BROTLI_COMPRESSOR_H_
+
+#include <stdint.h>
+
+#include <vector>
+
+#include <brotli/encode.h>
+
+#include "bsdiff/compressor_buffer.h"
+#include "bsdiff/compressor_interface.h"
+
+namespace bsdiff {
+
+class BrotliCompressor : public CompressorInterface {
+ public:
+  BrotliCompressor();
+  ~BrotliCompressor() override;
+
+  // CompressorInterface overrides.
+  bool Write(const uint8_t* buf, size_t size) override;
+  bool Finish() override;
+  const std::vector<uint8_t>& GetCompressedData() override;
+
+ private:
+  BrotliEncoderState* brotli_encoder_state_;
+
+  CompressorBuffer comp_buffer_;
+};
+
+}  // namespace bsdiff
+
+#endif  // _BSDIFF_BROTLI_COMPRESSOR_H_

brotli_compressor_unittest.cc

@@ -0,0 +1,29 @@
+// Copyright 2017 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "bsdiff/brotli_compressor.h"
+
+#include <gtest/gtest.h>
+
+namespace {
+
+const uint8_t kHelloWorld[] = {
+    0x48, 0x65, 0x6c, 0x6c, 0x6f, 0x20, 0x57, 0x6f, 0x72, 0x6c, 0x64, 0x0a,
+};
+}  // namespace
+
+namespace bsdiff {
+
+TEST(BrotliCompressorTest, BrotliCompressorSmoke) {
+  BrotliCompressor brotli_compressor;
+  EXPECT_TRUE(brotli_compressor.Write(kHelloWorld, sizeof(kHelloWorld)));
+  EXPECT_TRUE(brotli_compressor.Finish());
+  std::vector<uint8_t> compressed_data = brotli_compressor.GetCompressedData();
+  EXPECT_GT(compressed_data.size(), static_cast<size_t>(0));
+
+  // TODO(xunchang) run brotli decompressor and check we can get back
+  // kHelloWorld.
+}
+
+}  // namespace bsdiff

bsdiff.cc

@@ -49,7 +49,7 @@ namespace bsdiff {
 int bsdiff(const uint8_t* old_buf, size_t oldsize, const uint8_t* new_buf,
            size_t newsize, const char* patch_filename,
            SuffixArrayIndexInterface** sai_cache) {
-	BsdiffPatchWriter patch(patch_filename);
+	BsdiffPatchWriter patch(patch_filename, CompressorType::kBZ2);
 	return bsdiff(old_buf, oldsize, new_buf, newsize, &patch, sai_cache);
 }
 

bsdiff.gyp

@@ -50,13 +50,17 @@
       'link_settings': {
         'libraries': [
           '-lbz2',
+          '-lbrotlienc',
           '-ldivsufsort',
           '-ldivsufsort64',
         ],
       },
       'sources': [
+        'brotli_compressor.cc',
         'bsdiff.cc',
         'bz2_compressor.cc',
+        'compressor_buffer.cc',
+        'compressor_interface.cc',
         'diff_encoder.cc',
         'patch_writer.cc',
         'patch_writer_factory.cc',
@@ -82,6 +86,7 @@
       'link_settings': {
         'libraries': [
           '-lbz2',
+          '-lbrotlidec',
         ],
       },
       'sources': [
@@ -131,6 +136,7 @@
           },
           'includes': ['../common-mk/common_test.gypi'],
           'sources': [
+            'brotli_compressor_unittest.cc',
             'bsdiff_unittest.cc',
             'bspatch_unittest.cc',
             'diff_encoder_unittest.cc',

bz2_compressor.cc

@@ -20,7 +20,7 @@ const int kCompressionLevel = 9;
 
 namespace bsdiff {
 
-BZ2Compressor::BZ2Compressor() {
+BZ2Compressor::BZ2Compressor() : comp_buffer_(1024 * 1024) {
   memset(&bz_strm_, 0, sizeof(bz_strm_));
   int bz_err = BZ2_bzCompressInit(&bz_strm_, kCompressionLevel,
                                   0 /* verbosity */, 0 /* workFactor */);
@@ -29,7 +29,6 @@ BZ2Compressor::BZ2Compressor() {
   } else {
     bz_strm_initialized_ = true;
   }
-  comp_buffer_.resize(1024 * 1024);
 }
 
 BZ2Compressor::~BZ2Compressor() {
@@ -51,18 +50,17 @@ bool BZ2Compressor::Write(const uint8_t* buf, size_t size) {
   bz_strm_.avail_in = size;
 
   while (bz_strm_.avail_in) {
-    bz_strm_.next_out = reinterpret_cast<char*>(comp_buffer_.data());
-    bz_strm_.avail_out = comp_buffer_.size();
+    bz_strm_.next_out = reinterpret_cast<char*>(comp_buffer_.buffer_data());
+    bz_strm_.avail_out = comp_buffer_.buffer_size();
     int bz_err = BZ2_bzCompress(&bz_strm_, BZ_RUN);
     if (bz_err != BZ_RUN_OK) {
       LOG(ERROR) << "Compressing data, bz_error=" << bz_err << endl;
       return false;
     }
 
-    uint64_t output_bytes = comp_buffer_.size() - bz_strm_.avail_out;
+    uint64_t output_bytes = comp_buffer_.buffer_size() - bz_strm_.avail_out;
     if (output_bytes) {
-      comp_chunks_.emplace_back(comp_buffer_.data(),
-                                comp_buffer_.data() + output_bytes);
+      comp_buffer_.AddDataToChunks(output_bytes);
     }
   }
   return true;
@@ -76,14 +74,13 @@ bool BZ2Compressor::Finish() {
 
   int bz_err = BZ_FINISH_OK;
   while (bz_err == BZ_FINISH_OK) {
-    bz_strm_.next_out = reinterpret_cast<char*>(comp_buffer_.data());
-    bz_strm_.avail_out = comp_buffer_.size();
+    bz_strm_.next_out = reinterpret_cast<char*>(comp_buffer_.buffer_data());
+    bz_strm_.avail_out = comp_buffer_.buffer_size();
     bz_err = BZ2_bzCompress(&bz_strm_, BZ_FINISH);
 
-    uint64_t output_bytes = comp_buffer_.size() - bz_strm_.avail_out;
+    uint64_t output_bytes = comp_buffer_.buffer_size() - bz_strm_.avail_out;
     if (output_bytes) {
-      comp_chunks_.emplace_back(comp_buffer_.data(),
-                                comp_buffer_.data() + output_bytes);
+      comp_buffer_.AddDataToChunks(output_bytes);
     }
   }
   if (bz_err != BZ_STREAM_END) {
@@ -100,17 +97,7 @@ bool BZ2Compressor::Finish() {
 }
 
 const std::vector<uint8_t>& BZ2Compressor::GetCompressedData() {
-  if (!comp_chunks_.empty()) {
-    size_t chunks_size = 0;
-    for (const auto& chunk : comp_chunks_)
-      chunks_size += chunk.size();
-    comp_data_.reserve(comp_data_.size() + chunks_size);
-    for (const auto& chunk : comp_chunks_) {
-      comp_data_.insert(comp_data_.end(), chunk.begin(), chunk.end());
-    }
-    comp_chunks_.clear();
-  }
-  return comp_data_;
+  return comp_buffer_.GetCompressedData();
 }
 
 }  // namespace bsdiff

bz2_compressor.h

@@ -10,6 +10,9 @@
 
 #include <vector>
 
+#include "bsdiff/compressor_buffer.h"
+#include "bsdiff/compressor_interface.h"
+
 namespace bsdiff {
 
 // An in-memory class to wrap the low-level bzip2 compress functions. This class
@@ -17,19 +20,15 @@ namespace bsdiff {
 // data at the end of the compression step. For that, all the compressed data
 // is stored in memory.
 
-class BZ2Compressor {
+class BZ2Compressor : public CompressorInterface {
  public:
   BZ2Compressor();
-  ~BZ2Compressor();
+  ~BZ2Compressor() override;
 
-  // Compress and write the data at |buf| of size |size|.
-  bool Write(const uint8_t* buf, size_t size);
-
-  // Finish the compression step.
-  bool Finish();
-
-  // Return the compressed data. This method must be only called after Finish().
-  const std::vector<uint8_t>& GetCompressedData();
+  // CompressorInterface overrides.
+  bool Write(const uint8_t* buf, size_t size) override;
+  bool Finish() override;
+  const std::vector<uint8_t>& GetCompressedData() override;
 
  private:
   // The low-level bzip2 stream.
@@ -38,16 +37,7 @@ class BZ2Compressor {
   // Whether the bz_strm_ is initialized.
   bool bz_strm_initialized_{false};
 
-  // A list of chunks of compressed data. The final compressed representation is
-  // the concatenation of all the compressed data.
-  std::vector<std::vector<uint8_t>> comp_chunks_;
-
-  // A concatenated version of the |comp_chunks_|, used to store the compressed
-  // memory after Finish() is called.
-  std::vector<uint8_t> comp_data_;
-
-  // A temporary compression buffer for multiple calls to Write().
-  std::vector<uint8_t> comp_buffer_;
+  CompressorBuffer comp_buffer_;
 };
 
 }  // namespace bsdiff

compressor_buffer.cc

@@ -0,0 +1,36 @@
+// Copyright 2017 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "bsdiff/compressor_buffer.h"
+
+#include "bsdiff/logging.h"
+
+namespace bsdiff {
+
+const std::vector<uint8_t>& CompressorBuffer::GetCompressedData() {
+  if (!comp_chunks_.empty()) {
+    size_t chunks_size = 0;
+    for (const auto& chunk : comp_chunks_)
+      chunks_size += chunk.size();
+    comp_data_.reserve(comp_data_.size() + chunks_size);
+    for (const auto& chunk : comp_chunks_) {
+      comp_data_.insert(comp_data_.end(), chunk.begin(), chunk.end());
+    }
+    comp_chunks_.clear();
+  }
+  return comp_data_;
+}
+
+void CompressorBuffer::AddDataToChunks(size_t data_size) {
+  if (data_size > comp_buffer_.size()) {
+    LOG(ERROR) << "data size: " << data_size
+               << " is larger than buffer size: " << comp_buffer_.size()
+               << std::endl;
+    return;
+  }
+  comp_chunks_.emplace_back(comp_buffer_.data(),
+                            comp_buffer_.data() + data_size);
+}
+
+}  // namespace bsdiff

compressor_buffer.h

@@ -0,0 +1,53 @@
+// Copyright 2017 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef _BSDIFF_COMPRESSOR_BUFFER_H_
+#define _BSDIFF_COMPRESSOR_BUFFER_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <vector>
+
+namespace bsdiff {
+
+// An utility class to store the output from a compressor. It has a small
+// temporary buffer that holds the compressed data every time the compressed
+// stream has some output. That buffer is then appended to a list of vectors
+// stored in memory. This approach is more efficient than a simple vector
+// because it saves some resize / reallocation of memory. And the compressed
+// data is copied once no matter how many times we fetch output from the
+// compressor.
+class CompressorBuffer {
+ public:
+  explicit CompressorBuffer(size_t buffer_size) {
+    comp_buffer_.resize(buffer_size);
+  }
+
+  size_t buffer_size() { return comp_buffer_.size(); }
+  uint8_t* buffer_data() { return comp_buffer_.data(); }
+
+  // Concatenate the data in |comp_chunks_| and return the data in the form of
+  // a single vector.
+  const std::vector<uint8_t>& GetCompressedData();
+
+  // Add the data in |comp_buffer_| to the end of |comp_chunks_|.
+  void AddDataToChunks(size_t data_size);
+
+ private:
+  // A list of chunks of compressed data. The final compressed representation is
+  // the concatenation of all the compressed data.
+  std::vector<std::vector<uint8_t>> comp_chunks_;
+
+  // A concatenated version of the |comp_chunks_|, used to store the compressed
+  // memory after Finish() is called.
+  std::vector<uint8_t> comp_data_;
+
+  // A temporary compression buffer for multiple calls to Write().
+  std::vector<uint8_t> comp_buffer_;
+};
+
+}  // namespace bsdiff
+
+#endif  // _BSDIFF_COMPRESSOR_BUFFER_H_

compressor_interface.cc

@@ -0,0 +1,20 @@
+// Copyright 2017 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include <bsdiff/compressor_interface.h>
+
+#include <bsdiff/brotli_compressor.h>
+#include <bsdiff/bz2_compressor.h>
+
+namespace bsdiff {
+
+std::unique_ptr<CompressorInterface> CreateCompressor(CompressorType type) {
+  if (type == CompressorType::kBrotli) {
+    return std::unique_ptr<CompressorInterface>(new BrotliCompressor());
+  }
+  // Use BZ2 as a default compressor.
+  return std::unique_ptr<CompressorInterface>(new BZ2Compressor());
+}
+
+}  // namespace bsdiff

compressor_interface.h

@@ -0,0 +1,42 @@
+// Copyright 2017 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef _BSDIFF_COMPRESSOR_INTERFACE_H_
+#define _BSDIFF_COMPRESSOR_INTERFACE_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <memory>
+#include <vector>
+
+namespace bsdiff {
+
+enum class CompressorType {
+  kBrotli,
+  kBZ2,
+};
+
+class CompressorInterface {
+ public:
+  virtual ~CompressorInterface() = default;
+
+  // Compress and write |size| bytes of input data starting from |buf|.
+  virtual bool Write(const uint8_t* buf, size_t size) = 0;
+
+  // Finish the compression step.
+  virtual bool Finish() = 0;
+
+  // Return the compressed data. This method must be only called after Finish().
+  virtual const std::vector<uint8_t>& GetCompressedData() = 0;
+
+ protected:
+  CompressorInterface() = default;
+};
+
+std::unique_ptr<CompressorInterface> CreateCompressor(CompressorType type);
+
+}  // namespace bsdiff
+
+#endif  // _BSDIFF_COMPRESSOR_INTERFACE_H_

patch_writer.cc

@@ -6,6 +6,8 @@
 
 #include <string.h>
 
+#include "bsdiff/brotli_compressor.h"
+#include "bsdiff/bz2_compressor.h"
 #include "bsdiff/control_entry.h"
 #include "bsdiff/logging.h"
 
@@ -27,6 +29,14 @@ void EncodeInt64(int64_t x, uint8_t* buf) {
 
 namespace bsdiff {
 
+BsdiffPatchWriter::BsdiffPatchWriter(const std::string& patch_filename,
+                                     CompressorType type)
+    : patch_filename_(patch_filename) {
+  ctrl_stream_ = CreateCompressor(type);
+  diff_stream_ = CreateCompressor(type);
+  extra_stream_ = CreateCompressor(type);
+}
+
 bool BsdiffPatchWriter::Init(size_t /* new_size */) {
   fp_ = fopen(patch_filename_.c_str(), "w");
   if (!fp_) {
@@ -37,11 +47,11 @@ bool BsdiffPatchWriter::Init(size_t /* new_size */) {
 }
 
 bool BsdiffPatchWriter::WriteDiffStream(const uint8_t* data, size_t size) {
-  return diff_stream_.Write(data, size);
+  return diff_stream_->Write(data, size);
 }
 
 bool BsdiffPatchWriter::WriteExtraStream(const uint8_t* data, size_t size) {
-  return extra_stream_.Write(data, size);
+  return extra_stream_->Write(data, size);
 }
 
 bool BsdiffPatchWriter::AddControlEntry(const ControlEntry& entry) {
@@ -50,7 +60,7 @@ bool BsdiffPatchWriter::AddControlEntry(const ControlEntry& entry) {
   EncodeInt64(entry.diff_size, buf);
   EncodeInt64(entry.extra_size, buf + 8);
   EncodeInt64(entry.offset_increment, buf + 16);
-  if (!ctrl_stream_.Write(buf, sizeof(buf)))
+  if (!ctrl_stream_->Write(buf, sizeof(buf)))
     return false;
   written_output_ += entry.diff_size + entry.extra_size;
   return true;
@@ -62,15 +72,15 @@ bool BsdiffPatchWriter::Close() {
     return false;
   }
 
-  if (!ctrl_stream_.Finish() || !diff_stream_.Finish() ||
-      !extra_stream_.Finish()) {
+  if (!ctrl_stream_->Finish() || !diff_stream_->Finish() ||
+      !extra_stream_->Finish()) {
     LOG(ERROR) << "Finalizing compressed streams." << endl;
     return false;
   }
 
-  auto ctrl_data = ctrl_stream_.GetCompressedData();
-  auto diff_data = diff_stream_.GetCompressedData();
-  auto extra_data = extra_stream_.GetCompressedData();
+  auto ctrl_data = ctrl_stream_->GetCompressedData();
+  auto diff_data = diff_stream_->GetCompressedData();
+  auto extra_data = extra_stream_->GetCompressedData();
 
   if (!WriteHeader(ctrl_data.size(), diff_data.size()))
     return false;

patch_writer.h

@@ -5,10 +5,11 @@
 #ifndef _BSDIFF_PATCH_WRITER_H_
 #define _BSDIFF_PATCH_WRITER_H_
 
+#include <memory>
 #include <string>
 #include <vector>
 
-#include "bsdiff/bz2_compressor.h"
+#include "bsdiff/compressor_interface.h"
 #include "bsdiff/patch_writer_interface.h"
 
 namespace bsdiff {
@@ -17,10 +18,9 @@ namespace bsdiff {
 // BZ2-compressors and a 32-byte header.
 class BsdiffPatchWriter : public PatchWriterInterface {
  public:
-  // Create the patch writer using the file |patch_filename| where the patch
-  // data will be written to.
-  explicit BsdiffPatchWriter(const std::string& patch_filename)
-      : patch_filename_(patch_filename) {}
+  // Create the patch writer using |type| as the compression algorithm and the
+  // file |patch_filename| to write the patch data.
+  BsdiffPatchWriter(const std::string& patch_filename, CompressorType type);
 
   // PatchWriterInterface overrides.
   bool Init(size_t new_size) override;
@@ -43,9 +43,9 @@ class BsdiffPatchWriter : public PatchWriterInterface {
   std::string patch_filename_;
 
   // The three internal compressed streams.
-  BZ2Compressor ctrl_stream_;
-  BZ2Compressor diff_stream_;
-  BZ2Compressor extra_stream_;
+  std::unique_ptr<CompressorInterface> ctrl_stream_{nullptr};
+  std::unique_ptr<CompressorInterface> diff_stream_{nullptr};
+  std::unique_ptr<CompressorInterface> extra_stream_{nullptr};
 };
 
 }  // namespace bsdiff

patch_writer_factory.cc

@@ -8,10 +8,11 @@
 
 namespace bsdiff {
 
+// TODO(xunchang) choose the compression type based on the input parameter.
 std::unique_ptr<PatchWriterInterface> CreateBsdiffPatchWriter(
     const std::string& patch_filename) {
   return std::unique_ptr<PatchWriterInterface>(
-      new BsdiffPatchWriter(patch_filename));
+      new BsdiffPatchWriter(patch_filename, CompressorType::kBZ2));
 }
 
 }  // namespace bsdiff

patch_writer_unittest.cc

@@ -53,7 +53,7 @@ class BsdiffPatchWriterTest : public testing::Test {
   }
 
   test_utils::ScopedTempFile patch_file_{"bsdiff_newfile.XXXXXX"};
-  BsdiffPatchWriter patch_writer_{patch_file_.filename()};
+  BsdiffPatchWriter patch_writer_{patch_file_.filename(), CompressorType::kBZ2};
 };
 
 TEST_F(BsdiffPatchWriterTest, CreateEmptyPatchTest) {