Support IZeroCopyInput over TRope

initial

2 files changed, 77 insertions, 0 deletions

diff --git a/library/cpp/actors/util/rope.h b/library/cpp/actors/util/rope.h
index 4aacb41089..3c65588181 100644
--- a/library/cpp/actors/util/rope.h
+++ b/library/cpp/actors/util/rope.h
@@ -4,6 +4,7 @@
 #include <util/generic/string.h>
 #include <util/generic/hash_set.h>
 #include <util/generic/scope.h>
+#include <util/stream/zerocopy.h>
 #include <util/stream/str.h>
 #include <util/system/sanitizers.h>
 #include <util/system/valgrind.h>
@@ -1081,6 +1082,40 @@ public:
     }
 };
 
+class TRopeZeroCopyInput : public IZeroCopyInput {
+    TRope::TConstIterator Iter;
+    const char* Data = nullptr;
+    size_t Len = 0;
+
+private:
+    size_t DoNext(const void** ptr, size_t len) override {
+        Y_VERIFY_DEBUG(ptr);
+        if (Len == 0) {
+            if (Iter.Valid()) {
+                Data = Iter.ContiguousData();
+                Len = Iter.ContiguousSize();
+                Y_VERIFY_DEBUG(Len);
+                Y_VERIFY_DEBUG(Data);
+                ++Iter;
+            } else {
+                Data = nullptr;
+            }
+        }
+
+        size_t chunk = std::min(Len, len);
+        *ptr = Data;
+        Data += chunk;
+        Len -= chunk;
+        return chunk;
+    }
+
+public:
+    explicit TRopeZeroCopyInput(TRope::TConstIterator iter)
+        : Iter(iter)
+    {
+    }
+};
+
 inline TRope TRope::CopySpaceOptimized(TRope&& origin, size_t worstRatioPer1k, TRopeArena& arena) {
     TRope res;
     for (TRcBuf& chunk : origin.Chain) {
diff --git a/library/cpp/actors/util/rope_ut.cpp b/library/cpp/actors/util/rope_ut.cpp
index ba7c38044f..0ff85d6c59 100644
--- a/library/cpp/actors/util/rope_ut.cpp
+++ b/library/cpp/actors/util/rope_ut.cpp
@@ -373,4 +373,46 @@ Y_UNIT_TEST_SUITE(TRope) {
         }
     }
 
+    Y_UNIT_TEST(RopeZeroCopyInputBasic) {
+        TRope rope = CreateRope(Text, 3);
+        TRopeZeroCopyInput input(rope.Begin());
+
+        TString result;
+        TStringOutput output(result);
+        TransferData(&input, &output);
+        UNIT_ASSERT_EQUAL(result, Text);
+    }
+
+    Y_UNIT_TEST(RopeZeroCopyInput) {
+        TRope rope;
+        rope.Insert(rope.End(), TRope{"abc"});
+        rope.Insert(rope.End(), TRope{TString{}});
+        rope.Insert(rope.End(), TRope{"de"});
+        rope.Insert(rope.End(), TRope{TString{}});
+        rope.Insert(rope.End(), TRope{TString{}});
+        rope.Insert(rope.End(), TRope{"fghi"});
+
+        TRopeZeroCopyInput input(rope.Begin());
+
+        const char* data = nullptr;
+        size_t len;
+
+        len = input.Next(&data, 2);
+        UNIT_ASSERT_EQUAL("ab", TStringBuf(data, len));
+
+        len = input.Next(&data, 3);
+        UNIT_ASSERT_EQUAL("c", TStringBuf(data, len));
+
+        len = input.Next(&data, 3);
+        UNIT_ASSERT_EQUAL("de", TStringBuf(data, len));
+
+        len = input.Next(&data);
+        UNIT_ASSERT_EQUAL("fghi", TStringBuf(data, len));
+
+        len = input.Next(&data);
+        UNIT_ASSERT_EQUAL(len, 0);
+
+        len = input.Next(&data);
+        UNIT_ASSERT_EQUAL(len, 0);
+    }
 }