support passing stream offset to `TBrotliCompress`

commit_hash:21e0759b151fbed8aa1664b62c42bd5ba144301b

3 files changed, 45 insertions, 14 deletions

diff --git a/library/cpp/streams/brotli/brotli.cpp b/library/cpp/streams/brotli/brotli.cpp
index b7aa89f8e31..8acd43044ff 100644
--- a/library/cpp/streams/brotli/brotli.cpp
+++ b/library/cpp/streams/brotli/brotli.cpp
@@ -23,7 +23,8 @@ class TBrotliCompress::TImpl {
 public:
     TImpl(IOutputStream* slave,
           int quality,
-          const TBrotliDictionary* dictionary)
+          const TBrotliDictionary* dictionary,
+          size_t offset = 0)
         : Slave_(slave)
         , EncoderState_(BrotliEncoderCreateInstance(&TAllocator::Allocate, &TAllocator::Deallocate, nullptr))
     {
@@ -35,19 +36,12 @@ public:
             BrotliEncoderAttachPreparedDictionary(EncoderState_, dictionary->GetPreparedDictionary());
         }
 
-        auto res = BrotliEncoderSetParameter(
-            EncoderState_,
-            BROTLI_PARAM_QUALITY,
-            quality);
-
-        if (!res) {
-            BrotliEncoderDestroyInstance(EncoderState_);
-            ythrow yexception() << "Failed to set brotli encoder quality to " << quality;
-        }
+        SetEncoderParameter(BROTLI_PARAM_QUALITY, quality, "quality");
+        SetEncoderParameter(BROTLI_PARAM_STREAM_OFFSET, offset, "stream offset");
     }
 
     ~TImpl() {
-        BrotliEncoderDestroyInstance(EncoderState_);
+        FreeEncoder();
     }
 
     void Write(const void* buffer, size_t size) {
@@ -68,6 +62,19 @@ private:
     IOutputStream* Slave_;
     BrotliEncoderState* EncoderState_;
 
+    void SetEncoderParameter(BrotliEncoderParameter parameter,
+                             ui32 value,
+                             TStringBuf parameterName) {
+        if (!BrotliEncoderSetParameter(EncoderState_, parameter, value)) {
+            FreeEncoder();
+            ythrow yexception() << "Failed to set brotli encoder " << parameterName << " to " << value;
+        }
+    }
+
+    void FreeEncoder() {
+        BrotliEncoderDestroyInstance(EncoderState_);
+    }
+
     void DoWrite(const void* buffer, size_t size, BrotliEncoderOperation operation) {
         size_t availableOut = 0;
         ui8* outputBuffer = nullptr;
@@ -99,8 +106,12 @@ private:
 
 TBrotliCompress::TBrotliCompress(IOutputStream* slave,
                                  int quality,
-                                 const TBrotliDictionary* dictionary) {
-    Impl_.Reset(new TImpl(slave, quality, dictionary));
+                                 const TBrotliDictionary* dictionary,
+                                 size_t offset) {
+    Impl_.Reset(new TImpl(slave,
+                          quality,
+                          dictionary,
+                          offset));
 }
 
 TBrotliCompress::~TBrotliCompress() {
diff --git a/library/cpp/streams/brotli/brotli.h b/library/cpp/streams/brotli/brotli.h
index bf1813d7fc8..ef0c9611acf 100644
--- a/library/cpp/streams/brotli/brotli.h
+++ b/library/cpp/streams/brotli/brotli.h
@@ -18,10 +18,12 @@ public:
       @param slave stream to write compressed data to
       @param quality the higher the quality, the slower and better the compression. Range is 0 to 11.
       @param dictionary custom brotli dictionary
+      @param offset number of bytes already processed by a different encoder instance
     */
     explicit TBrotliCompress(IOutputStream* slave,
                              int quality = NBrotli::BEST_BROTLI_QUALITY,
-                             const TBrotliDictionary* dictionary = nullptr);
+                             const TBrotliDictionary* dictionary = nullptr,
+                             size_t offset = 0);
     ~TBrotliCompress() override;
 
 private:
diff --git a/library/cpp/streams/brotli/brotli_ut.cpp b/library/cpp/streams/brotli/brotli_ut.cpp
index 8eed051956b..fdf6cb2702e 100644
--- a/library/cpp/streams/brotli/brotli_ut.cpp
+++ b/library/cpp/streams/brotli/brotli_ut.cpp
@@ -95,4 +95,22 @@ Y_UNIT_TEST_SUITE(TBrotliTestSuite) {
         TStringBuf dict = "Bond";
         UNIT_ASSERT_VALUES_EQUAL(str, Decompress(Compress(str, dict), dict));
     }
+
+    Y_UNIT_TEST(TestStreamOffset) {
+        TString first = "apple pen";
+        TString second = " pineapple pen";
+
+        TString compressed;
+        TStringOutput out1(compressed);
+        TBrotliCompress stream1(&out1, NBrotli::BEST_BROTLI_QUALITY);
+        stream1.Write(first);
+        stream1.Flush();
+
+        TStringOutput out2(compressed);
+        TBrotliCompress stream2(&out2, NBrotli::BEST_BROTLI_QUALITY, nullptr, first.size());
+        stream2.Write(second);
+        stream2.Finish();
+
+        UNIT_ASSERT_VALUES_EQUAL(first + second, Decompress(compressed));
+    }
 } // Y_UNIT_TEST_SUITE(TBrotliTestSuite)