From a0b545cc68c310bdd5ed2bcd745b1773097e3ae2 Mon Sep 17 00:00:00 2001
From: Runji Wang <wangrunji0408@163.com>
Date: Mon, 6 Nov 2023 19:34:58 +0800
Subject: [PATCH] fix(expr): fix `overlay` function with large range input
 (#13268)

Signed-off-by: Runji Wang <wangrunji0408@163.com>
---
 src/expr/impl/src/scalar/overlay.rs | 57 ++++++++++++++++-------------
 1 file changed, 31 insertions(+), 26 deletions(-)

diff --git a/src/expr/impl/src/scalar/overlay.rs b/src/expr/impl/src/scalar/overlay.rs
index 509740892ffd0..329485596eb2e 100644
--- a/src/expr/impl/src/scalar/overlay.rs
+++ b/src/expr/impl/src/scalar/overlay.rs
@@ -95,13 +95,16 @@ pub fn overlay_for(
         });
     }
 
-    let mut chars = s.chars();
-    for _ in 1..start {
-        if let Some(c) = chars.next() {
-            writer.write_char(c).unwrap();
-        }
-    }
+    let mut chars = s.char_indices().skip(start as usize - 1).peekable();
+
+    // write the substring before the overlay.
+    let leading = match chars.peek() {
+        Some((i, _)) => &s[..*i],
+        None => s,
+    };
+    writer.write_str(leading).unwrap();
 
+    // write the new substring.
     writer.write_str(new_sub_str).unwrap();
 
     let Ok(count) = count.try_into() else {
@@ -112,8 +115,9 @@ pub fn overlay_for(
         return super::substr::substr_start(s, start_right, writer);
     };
 
-    for c in chars.skip(count) {
-        writer.write_char(c).unwrap();
+    // write the substring after the overlay.
+    if let Some((i, _)) = chars.nth(count) {
+        writer.write_str(&s[i..]).unwrap();
     }
 
     Ok(())
@@ -125,25 +129,26 @@ mod tests {
 
     #[test]
     fn test_overlay() {
-        let cases = vec![
-            ("aaa__aaa", "XY", 4, None, "aaaXYaaa"),
-            // Place at end.
-            ("aaa", "XY", 4, None, "aaaXY"),
-            // Place at start.
-            ("aaa", "XY", 1, Some(0), "XYaaa"),
-            // Replace shorter string.
-            ("aaa_aaa", "XYZ", 4, Some(1), "aaaXYZaaa"),
-            ("aaaaaa", "XYZ", 4, Some(0), "aaaXYZaaa"),
-            // Replace longer string.
-            ("aaa___aaa", "X", 4, Some(3), "aaaXaaa"),
-            // start too large.
-            ("aaa", "XY", 123, None, "aaaXY"),
-            // count too small or large.
-            ("aaa", "X", 4, Some(-123), "aaaXaaa"),
-            ("aaa_", "X", 4, Some(123), "aaaX"),
-        ];
+        case("aaa__aaa", "XY", 4, None, "aaaXYaaa");
+        // Place at end.
+        case("aaa", "XY", 4, None, "aaaXY");
+        // Place at start.
+        case("aaa", "XY", 1, Some(0), "XYaaa");
+        // Replace shorter string.
+        case("aaa_aaa", "XYZ", 4, Some(1), "aaaXYZaaa");
+        case("aaaaaa", "XYZ", 4, Some(0), "aaaXYZaaa");
+        // Replace longer string.
+        case("aaa___aaa", "X", 4, Some(3), "aaaXaaa");
+        // start too large.
+        case("aaa", "XY", 123, None, "aaaXY");
+        // count too small or large.
+        case("aaa", "X", 4, Some(-123), "aaaXaaa");
+        case("aaa_", "X", 4, Some(123), "aaaX");
+        // very large start and count
+        case("aaa", "X", i32::MAX, Some(i32::MAX), "aaaX");
 
-        for (s, new_sub_str, start, count, expected) in cases {
+        #[track_caller]
+        fn case(s: &str, new_sub_str: &str, start: i32, count: Option<i32>, expected: &str) {
             let mut writer = String::new();
             match count {
                 None => overlay(s, new_sub_str, start, &mut writer),