add assign() method + tests for TCompactVector

add assign() method for better compatibility with std::vector
commit_hash:7a216f3a025414493f3e6dc676b5c203208e462b

2 files changed, 42 insertions, 0 deletions

diff --git a/library/cpp/containers/compact_vector/compact_vector.h b/library/cpp/containers/compact_vector/compact_vector.h
index f795b40eb41..580d7823815 100644
--- a/library/cpp/containers/compact_vector/compact_vector.h
+++ b/library/cpp/containers/compact_vector/compact_vector.h
@@ -314,6 +314,14 @@ public:
         Clear();
     }
 
+    void assign(size_t count, const T& value) {
+        Clear();
+        Reserve(count);
+        for (size_t i = 0; i < count; ++i) {
+            PushBack(value);
+        }
+    }
+
     void erase(iterator position) {
         Y_ENSURE(position >= begin() && position < end());
         std::move(position + 1, end(), position);
diff --git a/library/cpp/containers/compact_vector/compact_vector_ut.cpp b/library/cpp/containers/compact_vector/compact_vector_ut.cpp
index f9fd43d0c79..593e61c3687 100644
--- a/library/cpp/containers/compact_vector/compact_vector_ut.cpp
+++ b/library/cpp/containers/compact_vector/compact_vector_ut.cpp
@@ -210,4 +210,38 @@ Y_UNIT_TEST_SUITE(TCompactVectorTest) {
         UNIT_ASSERT_VALUES_EQUAL(ptr1, vector[0].data());
         UNIT_ASSERT_VALUES_EQUAL(ptr2, vector[1].data());
     }
+
+    Y_UNIT_TEST(TestAssignWithCount) {
+        TCompactVector<ui32> vector = {1, 2, 3, 4, 5};
+        UNIT_ASSERT_VALUES_EQUAL(5u, vector.Size());
+
+        vector.assign(3, 42);
+        UNIT_ASSERT_VALUES_EQUAL(3u, vector.Size());
+
+        UNIT_ASSERT_VALUES_EQUAL(42u, vector[0]);
+        UNIT_ASSERT_VALUES_EQUAL(42u, vector[1]);
+        UNIT_ASSERT_VALUES_EQUAL(42u, vector[2]);
+    }
+
+    Y_UNIT_TEST(TestAssignWithZeroCount) {
+        TCompactVector<ui32> vector = {1, 2, 3, 4, 5};
+        UNIT_ASSERT_VALUES_EQUAL(5u, vector.Size());
+
+        vector.assign(0, 42);
+        UNIT_ASSERT_VALUES_EQUAL(0u, vector.Size());
+        UNIT_ASSERT(vector.empty());
+    }
+
+    Y_UNIT_TEST(TestAssignWithString) {
+        TCompactVector<std::string> vector = {"hello", "world"};
+        UNIT_ASSERT_VALUES_EQUAL(2u, vector.Size());
+
+        vector.assign(4, "test");
+        UNIT_ASSERT_VALUES_EQUAL(4u, vector.Size());
+
+        UNIT_ASSERT_VALUES_EQUAL("test", vector[0]);
+        UNIT_ASSERT_VALUES_EQUAL("test", vector[1]);
+        UNIT_ASSERT_VALUES_EQUAL("test", vector[2]);
+        UNIT_ASSERT_VALUES_EQUAL("test", vector[3]);
+    }
 } // Y_UNIT_TEST_SUITE(TCompactVectorTest)