TRope zero-copy extract

extend TRope interface with underlying container extraction

2 files changed, 79 insertions, 1 deletions

diff --git a/library/cpp/actors/util/rope.h b/library/cpp/actors/util/rope.h
index 9b234755f8c..4e5336dc637 100644
--- a/library/cpp/actors/util/rope.h
+++ b/library/cpp/actors/util/rope.h
@@ -117,6 +117,12 @@ namespace NRopeDetails {
         using TListIterator = typename TList::iterator;
     };
 
+    template<typename TContainer>
+    struct TContainerTraits {
+        static char* UnsafeGetDataMut(const TContainer& backend) {
+            return const_cast<char*>(backend.data());
+        }
+    };
 } // NRopeDetails
 
 class TRopeArena;
@@ -238,9 +244,29 @@ class TRope {
                 });
             }
 
+            template <class TType>
+            bool ContainsNativeType() const {
+                return Visit(Owner, [](EType, auto& value) {
+                    using T = std::decay_t<decltype(value)>;
+                    return std::is_same_v<T, TType>;
+                });
+            }
+
+            template <class TResult>
+            TResult GetRaw() const {
+                return Visit(Owner, [](EType, auto& value) {
+                    using T = std::decay_t<decltype(value)>;
+                    if constexpr (std::is_same_v<T, TResult>) {
+                        return value;
+                    } else {
+                        Y_FAIL();
+                        return TResult{}; // unreachable
+                    }
+                });
+            }
+
             explicit operator bool() const {
                 return Owner;
-
             }
 
         private:
@@ -387,6 +413,21 @@ class TRope {
         TChunk& operator =(const TChunk&) = default;
         TChunk& operator =(TChunk&&) = default;
 
+        template <class TType>
+        bool ContainsNativeType() const {
+            return Backend.ContainsNativeType<TType>();
+        }
+
+        template <class TResult>
+        TResult GetRaw() const {
+            return Backend.GetRaw<TResult>();
+        }
+
+        bool OwnsWholeContainer() const {
+            auto [data, size] = Backend.GetData();
+            return size == GetSize();
+        }
+
         size_t GetSize() const {
             return End - Begin;
         }
@@ -998,6 +1039,24 @@ public:
         return res;
     }
 
+    /**
+     * WARN: this method supports extracting only for natively supported types for any other type the data *will* be copied
+     */
+    template <class TResult>
+    TResult ExtractUnderlyingContainerOrCopy() const {
+        if (IsContiguous() && GetSize() != 0) {
+            const auto& chunk = Begin().GetChunk();
+            if (chunk.ContainsNativeType<TResult>() && chunk.OwnsWholeContainer()) {
+                return chunk.GetRaw<TResult>();
+            }
+        }
+
+        TResult res = TResult::Uninitialized(GetSize());
+        char* data = NRopeDetails::TContainerTraits<TResult>::UnsafeGetDataMut(res);
+        Begin().ExtractPlainDataAndAdvance(data, res.size());
+        return res;
+    }
+
     class TContiguousSpan;
     class TMutableContiguousSpan {
         friend class TContiguousSpan;
diff --git a/library/cpp/actors/util/rope_ut.cpp b/library/cpp/actors/util/rope_ut.cpp
index 55f995b66ab..1b2081ecedd 100644
--- a/library/cpp/actors/util/rope_ut.cpp
+++ b/library/cpp/actors/util/rope_ut.cpp
@@ -95,6 +95,25 @@ Y_UNIT_TEST_SUITE(TRope) {
     }
 
 #ifndef TSTRING_IS_STD_STRING
+    Y_UNIT_TEST(ExtractZeroCopy) {
+        TString str = Text;
+        TRope packed(str);
+        TString extracted = packed.ExtractUnderlyingContainerOrCopy<TString>();
+        UNIT_ASSERT_EQUAL(str.data(), extracted.data());
+    }
+
+    Y_UNIT_TEST(ExtractZeroCopySlice) {
+        TString str = Text;
+        TRope sliced(str);
+        sliced.EraseFront(1);
+        TString extracted = sliced.ExtractUnderlyingContainerOrCopy<TString>();
+        UNIT_ASSERT_UNEQUAL(str.data(), extracted.data());
+        TRope sliced2(str);
+        sliced2.EraseBack(1);
+        TString extracted2 = sliced2.ExtractUnderlyingContainerOrCopy<TString>();
+        UNIT_ASSERT_UNEQUAL(str.data(), extracted2.data());
+    }
+
     Y_UNIT_TEST(TStringDetach) {
         TRope pf;
         TRope rope;