AK/StringUtils: Ensure needle positions don't overlap in replace

Previously, `replace` used `find_all` to find all of the positions to
replace. But `find_all` finds all the *overlapping* instances of the
needle, while `replace` assumed that the next position was always at
least `needle.length()` away from the last one. This led to crashes like
https://github.com/SerenityOS/jakt/issues/1159.

AK/StringUtils.cpp

@@ -521,32 +521,53 @@ ByteString invert_case(StringView str)
     return builder.to_byte_string();
 }
 
-ByteString replace(StringView str, StringView needle, StringView replacement, ReplaceMode replace_mode)
+// Finishes the replacing algorithm once it is known that ita least one
+// replacemnet is going to be done. Otherwise the caller may want to follow a
+// different route to construct its output.
+static StringBuilder replace_into_builder(StringView str, StringView needle, StringView replacement, ReplaceMode replace_mode, size_t first_replacement_position)
+{
+    StringBuilder replaced_string;
+
+    replaced_string.append(str.substring_view(0, first_replacement_position));
+    replaced_string.append(replacement);
+
+    StringView remaining = str.substring_view(first_replacement_position + needle.length());
+
+    switch (replace_mode) {
+    case ReplaceMode::All:
+        while (!remaining.is_empty()) {
+            auto maybe_pos = remaining.find(needle);
+            if (!maybe_pos.has_value())
+                break;
+            replaced_string.append(remaining.substring_view(0, *maybe_pos));
+            replaced_string.append(replacement);
+            remaining = remaining.substring_view(*maybe_pos + needle.length());
+        }
+        break;
+    case ReplaceMode::FirstOnly:
+        // We already made the first replacement.
+        break;
+    }
+
+    // The remaining bits either don't contain the needle or are ignored due to
+    // `replace_mode` being `ReplaceMode::FirstOnly`.
+    replaced_string.append(remaining);
+
+    return replaced_string;
+}
+
+ByteString replace(StringView str, StringView needle, StringView replacement,
+    ReplaceMode replace_mode)
 {
     if (str.is_empty())
         return str;
 
-    Vector<size_t> positions;
+    auto maybe_first = str.find(needle);
-    if (replace_mode == ReplaceMode::All) {
+    if (!maybe_first.has_value())
-        positions = str.find_all(needle);
+        return str;
-        if (!positions.size())
-            return str;
-    } else {
-        auto pos = str.find(needle);
-        if (!pos.has_value())
-            return str;
-        positions.append(pos.value());
-    }
 
-    StringBuilder replaced_string;
+    auto resulting_builder = replace_into_builder(str, needle, replacement, replace_mode, *maybe_first);
-    size_t last_position = 0;
+    return resulting_builder.to_byte_string();
-    for (auto& position : positions) {
-        replaced_string.append(str.substring_view(last_position, position - last_position));
-        replaced_string.append(replacement);
-        last_position = position + needle.length();
-    }
-    replaced_string.append(str.substring_view(last_position, str.length() - last_position));
-    return replaced_string.to_byte_string();
 }
 
 ErrorOr<String> replace(String const& haystack, StringView needle, StringView replacement, ReplaceMode replace_mode)
@@ -554,28 +575,14 @@ ErrorOr<String> replace(String const& haystack, StringView needle, StringView re
     if (haystack.is_empty())
         return haystack;
 
-    // FIXME: Propagate Vector allocation failures (or do this without putting positions in a vector)
+    auto const source_bytes = haystack.bytes_as_string_view();
-    Vector<size_t> positions;
-    if (replace_mode == ReplaceMode::All) {
-        positions = haystack.bytes_as_string_view().find_all(needle);
-        if (!positions.size())
-            return haystack;
-    } else {
-        auto pos = haystack.bytes_as_string_view().find(needle);
-        if (!pos.has_value())
-            return haystack;
-        positions.append(pos.value());
-    }
 
-    StringBuilder replaced_string;
+    auto maybe_first = source_bytes.find(needle);
-    size_t last_position = 0;
+    if (!maybe_first.has_value())
-    for (auto& position : positions) {
+        return haystack;
-        replaced_string.append(haystack.bytes_as_string_view().substring_view(last_position, position - last_position));
+
-        replaced_string.append(replacement);
+    auto resulting_builder = replace_into_builder(source_bytes, needle, replacement, replace_mode, *maybe_first);
-        last_position = position + needle.length();
+    return resulting_builder.to_string();
-    }
-    replaced_string.append(haystack.bytes_as_string_view().substring_view(last_position, haystack.bytes_as_string_view().length() - last_position));
-    return replaced_string.to_string();
 }
 #endif
 

Tests/AK/TestStringUtils.cpp

@@ -407,6 +407,26 @@ TEST_CASE(find)
     EXPECT_EQ(AK::StringUtils::find(test_string, "78"sv).has_value(), false);
 }
 
+TEST_CASE(replace_all_overlapping)
+{
+    // Replace only should take into account non-overlapping instances of the
+    // needle, since it is looking to replace them.
+
+    // These samples were grabbed from ADKaster's sample code in
+    // https://github.com/SerenityOS/jakt/issues/1159. This is the equivalent
+    // C++ code that triggered the same bug from Jakt's code generator.
+
+    auto const replace_like_in_jakt = [](StringView source) -> ByteString {
+        ByteString replaced = AK::StringUtils::replace(source, "\\\""sv, "\""sv, ReplaceMode::All);
+        replaced = AK::StringUtils::replace(replaced.view(), "\\\\"sv, "\\"sv, ReplaceMode::All);
+        return replaced;
+    };
+
+    EXPECT_EQ(replace_like_in_jakt("\\\\\\\\\\\\\\\\"sv), "\\\\\\\\"sv);
+    EXPECT_EQ(replace_like_in_jakt(" auto str4 = \"\\\";"sv), " auto str4 = \"\";"sv);
+    EXPECT_EQ(replace_like_in_jakt(" auto str5 = \"\\\\\";"sv), " auto str5 = \"\\\";"sv);
+}
+
 TEST_CASE(to_snakecase)
 {
     EXPECT_EQ(AK::StringUtils::to_snakecase("foobar"sv), "foobar");