StringSplitter: recreate State after copy

commit_hash:75f3a4a20b797f07c51c63b96f0f1e1374707f33
ytsaurus · Oct 16, 2024 · 4dbbb2c · 4dbbb2c
1 parent da0ee10
commit 4dbbb2c
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 0 deletions.
diff --git a/util/string/split.h b/util/string/split.h
@@ -660,6 +660,15 @@ namespace NStringSplitPrivate {
         using TIterator = TIteratorOf<String>;
         friend class TStringSplitter<String>;
 
+        template <typename S = String, std::enable_if_t<THasData<S>::value, int> = 0>
+        TIterState(const String& string) noexcept
+            : TStringBufType()
+            , DelimiterEnd_(string.data())
+            , OriginEnd_(string.data() + string.size())
+        {
+        }
+
+        template <typename S = String, std::enable_if_t<!THasData<S>::value, int> = 0>
         TIterState(const String& string) noexcept
             : TStringBufType()
             , DelimiterEnd_(std::begin(string))
@@ -700,6 +709,11 @@ namespace NStringSplitPrivate {
             return TokenDelim() == DelimiterEnd_;
         }
 
+        void MarkExhausted() noexcept {
+            UpdateParentBuf(OriginEnd_, OriginEnd_);
+            DelimiterEnd_ = OriginEnd_;
+        }
+
     private:
         TIterator DelimiterEnd_;
         const TIterator OriginEnd_;
@@ -845,6 +859,24 @@ namespace NStringSplitPrivate {
             {
             }
 
+            TSplitRangeBase(const TSplitRangeBase& other)
+                : String_(other.String_)
+                , State_(String_)
+                , Delimiter_(other.Delimiter_)
+            {
+            }
+
+            TSplitRangeBase(TSplitRangeBase&& other)
+                : String_(std::move(other.String_))
+                , State_(String_)
+                , Delimiter_(std::move(other.Delimiter_))
+            {
+                other.State_.MarkExhausted();
+            }
+
+            TSplitRangeBase& operator=(const TSplitRangeBase& other) = delete;
+            TSplitRangeBase& operator=(TSplitRangeBase&& other) = delete;
+
             inline TIteratorState* Next() {
                 if (State_.DelimiterIsEmpty()) {
                     return nullptr;

diff --git a/util/string/split_ut.cpp b/util/string/split_ut.cpp
@@ -443,6 +443,13 @@ Y_UNIT_TEST_SUITE(StringSplitter) {
             sum += FromString<int>(it.Token());
         }
         UNIT_ASSERT_VALUES_EQUAL(sum, 6);
+
+        TString ssum;
+        for (const auto& it : StringSplitter("  1 2 3   " + std::string(100, ' ')).Split(' ').SkipEmpty()) {
+            ssum += FromString<TString>(it.Token());
+            ssum += ';';
+        }
+        UNIT_ASSERT_VALUES_EQUAL(ssum, "1;2;3;");
     }
 
     Y_UNIT_TEST(TestTake) {
@@ -746,6 +753,25 @@ Y_UNIT_TEST_SUITE(StringSplitter) {
         UNIT_ASSERT_VALUES_EQUAL(expected1, actual3);
     }
 
+    Y_UNIT_TEST(TesIterationAfterMove) {
+        const TString src = TString::Join(
+            "aaa",
+            TString(250, 'c'),
+            "bbb",
+            "aaa",
+            TString(250, 'c'),
+            "bbb");
+        auto s1 = StringSplitter(std::string(src)).SplitByString("c").SkipEmpty();
+        {
+            auto s2 = std::move(s1);
+            const TVector<TString> expected2 = {"aaa", "bbbaaa", "bbb"};
+            const auto result2 = s2.ToList<TString>();
+            UNIT_ASSERT_VALUES_EQUAL(result2, expected2);
+        }
+        const auto result1 = s1.ToList<TString>();
+        Y_UNUSED(result1); // valid but unspecified value
+    }
+
     Y_UNIT_TEST(TestConstCString) {
         const char* b = "a;b";
         const char* e = b + 3;