Add reserve to TContiguousData

add reserve to TContiguousData

2 files changed, 79 insertions, 0 deletions

diff --git a/library/cpp/actors/util/contiguous_data.h b/library/cpp/actors/util/contiguous_data.h
index d7123ca457..35f9d25b33 100644
--- a/library/cpp/actors/util/contiguous_data.h
+++ b/library/cpp/actors/util/contiguous_data.h
@@ -1024,6 +1024,54 @@ public:
         return End;
     }
 
+    char& operator[](size_t pos) {
+        return UnsafeGetDataMut()[pos];
+    }
+
+    const char& operator[](size_t pos) const {
+        return GetData()[pos];
+    }
+
+    void reserve(size_t size) {
+        ReserveTailroom(size);
+    }
+
+    void ReserveHeadroom(size_t size) {
+        if (Headroom() >= size) {
+            return;
+        }
+        auto newData = TContiguousData::Uninitialized(GetSize(), size, UnsafeTailroom());
+        if (auto data = GetData(); data) {
+            std::memcpy(newData.UnsafeGetDataMut(), GetData(), GetSize());
+        }
+        *this = std::move(newData);
+    }
+
+    void ReserveTailroom(size_t size) {
+        if (Tailroom() >= size) {
+            return;
+        }
+        auto newData = TContiguousData::Uninitialized(GetSize(), UnsafeHeadroom(), size);
+        if (auto data = GetData(); data) {
+            std::memcpy(newData.UnsafeGetDataMut(), GetData(), GetSize());
+        }
+        *this = std::move(newData);
+    }
+
+    void ReserveBidi(size_t headroom, size_t tailroom) {
+        if (Headroom() >= headroom && Tailroom() >= tailroom) {
+            return;
+        }
+        auto newData = TContiguousData::Uninitialized(
+                GetSize(),
+                std::max(UnsafeHeadroom(), headroom),
+                std::max(UnsafeTailroom(), tailroom));
+        if (auto data = GetData(); data) {
+            std::memcpy(newData.UnsafeGetDataMut(), GetData(), GetSize());
+        }
+        *this = std::move(newData);
+    }
+
     EResizeResult GrowFront(size_t size, EResizeStrategy strategy = EResizeStrategy::KeepRooms) {
         if (Headroom() >= size && Backend.UpdateCookiesBegin(Begin, Begin - size)) {
             Begin -= size;
diff --git a/library/cpp/actors/util/contiguous_data_ut.cpp b/library/cpp/actors/util/contiguous_data_ut.cpp
index 6452e748b8..176b5d302f 100644
--- a/library/cpp/actors/util/contiguous_data_ut.cpp
+++ b/library/cpp/actors/util/contiguous_data_ut.cpp
@@ -183,4 +183,35 @@ Y_UNIT_TEST_SUITE(TContiguousData) {
         UNIT_ASSERT_EQUAL(data2.data(), extracted.data());
         UNIT_ASSERT_EQUAL(data2.size(), extracted.size());
     }
+
+    Y_UNIT_TEST(Reserve) {
+        TContiguousData data = TContiguousData::Copy("test", 4, 5, 6);
+        TContiguousData data2 = data;
+        data.reserve(1);
+        data.ReserveTailroom(6);
+        UNIT_ASSERT_EQUAL(data.GetData(), data2.GetData());
+        UNIT_ASSERT_EQUAL(data.GetSize(), data2.GetSize());
+        UNIT_ASSERT_EQUAL(data.Tailroom(), 6);
+        data.ReserveHeadroom(5);
+        UNIT_ASSERT_EQUAL(data.GetData(), data2.GetData());
+        UNIT_ASSERT_EQUAL(data.GetSize(), data2.GetSize());
+        UNIT_ASSERT_EQUAL(data.Headroom(), 5);
+        data.ReserveBidi(5, 6);
+        UNIT_ASSERT_EQUAL(data.GetData(), data2.GetData());
+        UNIT_ASSERT_EQUAL(data.GetSize(), data2.GetSize());
+        UNIT_ASSERT_EQUAL(data.Headroom(), 5);
+        UNIT_ASSERT_EQUAL(data.Tailroom(), 6);
+        data.ReserveHeadroom(6);
+        UNIT_ASSERT_EQUAL(data.Headroom(), 6);
+        UNIT_ASSERT_EQUAL(data.Tailroom(), 6);
+        UNIT_ASSERT_EQUAL(::memcmp(data.GetData(), "test", 4), 0);
+        data.ReserveTailroom(7);
+        UNIT_ASSERT_EQUAL(data.Headroom(), 6);
+        UNIT_ASSERT_EQUAL(data.Tailroom(), 7);
+        UNIT_ASSERT_EQUAL(::memcmp(data.GetData(), "test", 4), 0);
+        data.ReserveBidi(7, 8);
+        UNIT_ASSERT_EQUAL(data.Headroom(), 7);
+        UNIT_ASSERT_EQUAL(data.Tailroom(), 8);
+        UNIT_ASSERT_EQUAL(::memcmp(data.GetData(), "test", 4), 0);
+    }
 }