Optimize string operations

ober

fbb60d682afbbd1f073c9a074fec77db7a12f20d

diff --git a/lib/std/misc/string-more.ss b/lib/std/misc/string-more.ss
index 8709109..c95afad 100644
--- a/lib/std/misc/string-more.ss
+++ b/lib/std/misc/string-more.ss
@@ -18,31 +18,44 @@
   (import (chezscheme)
           (only (jerboa core) def))
 
+  (def (string-region=? a ai b bi n)
+    (let loop ([k 0])
+      (or (fx= k n)
+          (and (char=? (string-ref a (fx+ ai k))
+                       (string-ref b (fx+ bi k)))
+               (loop (fx+ k 1))))))
+
   ;; Test if str starts with prefix
   (def (string-prefix? prefix str)
     (let ([plen (string-length prefix)]
           [slen (string-length str)])
       (and (<= plen slen)
-           (string=? (substring str 0 plen) prefix))))
+           (string-region=? prefix 0 str 0 plen))))
 
   ;; Test if str ends with suffix
   (def (string-suffix? suffix str)
     (let ([suflen (string-length suffix)]
           [slen (string-length str)])
       (and (<= suflen slen)
-           (string=? (substring str (- slen suflen) slen) suffix))))
+           (string-region=? suffix 0 str (- slen suflen) suflen))))
 
   ;; Test if str contains substring
   (def (string-contains? needle str)
     (let ([nlen (string-length needle)]
           [slen (string-length str)])
-      (if (> nlen slen)
-          #f
-          (let loop ([i 0])
-            (cond
-              [(> (+ i nlen) slen) #f]
-              [(string=? (substring str i (+ i nlen)) needle) #t]
-              [else (loop (+ i 1))])))))
+      (cond
+        [(zero? nlen) #t]
+        [(> nlen slen) #f]
+        [else
+         (let ([first (string-ref needle 0)]
+               [limit (- slen nlen)])
+           (let loop ([i 0])
+             (cond
+               [(> i limit) #f]
+               [(and (char=? (string-ref str i) first)
+                     (string-region=? str i needle 0 nlen))
+                #t]
+               [else (loop (+ i 1))])))])))
 
   ;; Trim whitespace from both ends
   (def (string-trim-both str)
@@ -59,20 +72,48 @@
 
   ;; Join list of strings with separator
   (def (string-join strs sep)
-    (if (null? strs)
-        ""
-        (let loop ([rest (cdr strs)] [acc (car strs)])
-          (if (null? rest)
-              acc
-              (loop (cdr rest)
-                    (string-append acc sep (car rest)))))))
+    (cond
+      [(null? strs) ""]
+      [(null? (cdr strs)) (car strs)]
+      [else
+       (let* ([sep-len (string-length sep)]
+              [total
+               (let loop ([rest strs] [n 0] [first? #t])
+                 (if (null? rest)
+                     n
+                     (loop (cdr rest)
+                           (+ n
+                              (string-length (car rest))
+                              (if first? 0 sep-len))
+                           #f)))]
+              [out (make-string total)])
+         (let loop ([rest strs] [pos 0] [first? #t])
+           (if (null? rest)
+               out
+               (let* ([pos (if first?
+                               pos
+                               (begin
+                                 (string-copy! sep 0 out pos sep-len)
+                                 (+ pos sep-len)))]
+                      [s (car rest)]
+                      [slen (string-length s)])
+                 (string-copy! s 0 out pos slen)
+                 (loop (cdr rest) (+ pos slen) #f)))))]))
 
   ;; Repeat string N times
   (def (string-repeat str n)
-    (let loop ([i 0] [acc ""])
-      (if (>= i n)
-          acc
-          (loop (+ i 1) (string-append acc str)))))
+    (cond
+      [(<= n 0) ""]
+      [(= n 1) str]
+      [else
+       (let* ([slen (string-length str)]
+              [out (make-string (* slen n))])
+         (let loop ([i 0] [pos 0])
+           (if (= i n)
+               out
+               (begin
+                 (string-copy! str 0 out pos slen)
+                 (loop (+ i 1) (+ pos slen))))))]))
 
   ;; Find first index where char/pred matches (or #f)
   (def (string-index str pred/char)
@@ -166,30 +207,56 @@
            (let ([slen (string-length str)]
                  [dlen (string-length delim)])
              (if (= dlen 0) (list str)
-                 (let loop ([i 0] [start 0] [acc '()])
-                   (cond
-                     [(> (+ i dlen) slen)
-                      (reverse (cons (substring str start slen) acc))]
-                     [(string=? (substring str i (+ i dlen)) delim)
-                      (loop (+ i dlen) (+ i dlen)
-                            (cons (substring str start i) acc))]
-                     [else (loop (+ i 1) start acc)])))))]))
+                 (let ([first (string-ref delim 0)]
+                       [limit (- slen dlen)])
+                   (let loop ([i 0] [start 0] [acc '()])
+                     (cond
+                       [(> i limit)
+                        (reverse (cons (substring str start slen) acc))]
+                       [(and (char=? (string-ref str i) first)
+                             (string-region=? str i delim 0 dlen))
+                        (loop (+ i dlen) (+ i dlen)
+                              (cons (substring str start i) acc))]
+                       [else (loop (+ i 1) start acc)]))))))]))
 
   ;; Replace all occurrences of old with new in string
   (def (string-replace str old new)
     (let ([slen (string-length str)]
-          [olen (string-length old)])
-      (if (= olen 0) str
-          (let loop ([i 0] [acc '()])
-            (cond
-              [(> (+ i olen) slen)
-               (apply string-append
-                      (reverse (cons (substring str i slen) acc)))]
-              [(string=? (substring str i (+ i olen)) old)
-               (loop (+ i olen) (cons new acc))]
-              [else
-               (loop (+ i 1)
-                     (cons (string (string-ref str i)) acc))])))))
+          [olen (string-length old)]
+          [nlen (string-length new)])
+      (if (= olen 0)
+          str
+          (let ([first (string-ref old 0)]
+                [limit (- slen olen)])
+            (define (match-at? i)
+              (and (<= i limit)
+                   (char=? (string-ref str i) first)
+                   (string-region=? str i old 0 olen)))
+            (let count-loop ([i 0] [count 0])
+              (cond
+                [(> i limit)
+                 (if (= count 0)
+                     str
+                     (let ([out (make-string (+ slen (* count (- nlen olen))))])
+                       (let copy-loop ([i 0] [start 0] [pos 0])
+                         (cond
+                           [(> i limit)
+                            (let ([tail (- slen start)])
+                              (string-copy! str start out pos tail)
+                              out)]
+                           [(match-at? i)
+                            (let ([chunk (- i start)])
+                              (string-copy! str start out pos chunk)
+                              (string-copy! new 0 out (+ pos chunk) nlen)
+                              (copy-loop (+ i olen)
+                                         (+ i olen)
+                                         (+ pos chunk nlen)))]
+                           [else
+                            (copy-loop (+ i 1) start pos)]))))]
+                [(match-at? i)
+                 (count-loop (+ i olen) (+ count 1))]
+                [else
+                 (count-loop (+ i 1) count)]))))))
 
   ;; Filter characters by predicate
   (def (string-filter pred str)
diff --git a/lib/std/misc/string.sls b/lib/std/misc/string.sls
index 697c1c8..be285f2 100644
--- a/lib/std/misc/string.sls
+++ b/lib/std/misc/string.sls
@@ -20,6 +20,13 @@
   (import (chezscheme)
           (std pregexp))
 
+  (define (string-region=? a ai b bi n)
+    (let loop ([k 0])
+      (or (fx= k n)
+          (and (char=? (string-ref a (fx+ ai k))
+                       (string-ref b (fx+ bi k)))
+               (loop (fx+ k 1))))))
+
   (define string-split
     (case-lambda
       ((str) (string-split str #\space))
@@ -44,23 +51,49 @@
           [len (string-length str)])
       (if (zero? slen)
         (map string (string->list str))
-        (let loop ([i 0] [start 0] [acc '()])
-          (cond
-            [(> (+ i slen) len)
-             (reverse (cons (substring str start len) acc))]
-            [(string=? (substring str i (+ i slen)) sep)
-             (loop (+ i slen) (+ i slen)
-                   (cons (substring str start i) acc))]
-            [else (loop (+ i 1) start acc)])))))
+        (let ([first (string-ref sep 0)]
+              [limit (- len slen)])
+          (let loop ([i 0] [start 0] [acc '()])
+            (cond
+              [(> i limit)
+               (reverse (cons (substring str start len) acc))]
+              [(and (char=? (string-ref str i) first)
+                    (string-region=? str i sep 0 slen))
+               (loop (+ i slen) (+ i slen)
+                     (cons (substring str start i) acc))]
+              [else (loop (+ i 1) start acc)]))))))
 
   (define string-join
     (case-lambda
       ((lst) (string-join lst " "))
       ((lst sep)
-       (if (null? lst) ""
-         (let loop ([rest (cdr lst)] [acc (car lst)])
-           (if (null? rest) acc
-             (loop (cdr rest) (string-append acc sep (car rest)))))))))
+       (cond
+         [(null? lst) ""]
+         [(null? (cdr lst)) (car lst)]
+         [else
+          (let* ([sep-len (string-length sep)]
+                 [total
+                  (let loop ([rest lst] [n 0] [first? #t])
+                    (if (null? rest)
+                        n
+                        (loop (cdr rest)
+                              (+ n
+                                 (string-length (car rest))
+                                 (if first? 0 sep-len))
+                              #f)))]
+                 [out (make-string total)])
+            (let loop ([rest lst] [pos 0] [first? #t])
+              (if (null? rest)
+                  out
+                  (let* ([pos (if first?
+                                  pos
+                                  (begin
+                                    (string-copy! sep 0 out pos sep-len)
+                                    (+ pos sep-len)))]
+                         [s (car rest)]
+                         [slen (string-length s)])
+                    (string-copy! s 0 out pos slen)
+                    (loop (cdr rest) (+ pos slen) #f)))))]))))
 
   (define (string-trim str)
     (let* ([len (string-length str)]
@@ -73,23 +106,33 @@
       (substring str start end)))
 
   (define (string-prefix? prefix str)
-    (and (<= (string-length prefix) (string-length str))
-         (string=? prefix (substring str 0 (string-length prefix)))))
+    (let ([plen (string-length prefix)]
+          [len (string-length str)])
+      (and (<= plen len)
+           (string-region=? prefix 0 str 0 plen))))
 
   (define (string-suffix? suffix str)
     (let ([slen (string-length suffix)]
           [len (string-length str)])
       (and (<= slen len)
-           (string=? suffix (substring str (- len slen) len)))))
+           (string-region=? suffix 0 str (- len slen) slen))))
 
   (define (string-contains str sub)
     (let ([slen (string-length sub)]
           [len (string-length str)])
-      (let loop ([i 0])
-        (cond
-          [(> (+ i slen) len) #f]
-          [(string=? (substring str i (+ i slen)) sub) i]
-          [else (loop (+ i 1))]))))
+      (cond
+        [(zero? slen) 0]
+        [(> slen len) #f]
+        [else
+         (let ([first (string-ref sub 0)]
+               [limit (- len slen)])
+           (let loop ([i 0])
+             (cond
+               [(> i limit) #f]
+               [(and (char=? (string-ref str i) first)
+                     (string-region=? str i sub 0 slen))
+                i]
+               [else (loop (+ i 1))])))])))
 
   (define string-index
     (case-lambda
diff --git a/lib/std/srfi/srfi-13.ss b/lib/std/srfi/srfi-13.ss
index af48e42..ad289b1 100644
--- a/lib/std/srfi/srfi-13.ss
+++ b/lib/std/srfi/srfi-13.ss
@@ -36,6 +36,13 @@
   (import (chezscheme)
           (only (jerboa core) def))
 
+  (def (string-region=? a ai b bi n)
+    (let lp ((k 0))
+      (or (fx= k n)
+          (and (char=? (string-ref a (fx+ ai k))
+                       (string-ref b (fx+ bi k)))
+               (lp (fx+ k 1))))))
+
   (def (string-index str pred . rest)
     (let* ((start (if (pair? rest) (car rest) 0))
            (end (if (and (pair? rest) (pair? (cdr rest))) (cadr rest) (string-length str)))
@@ -68,25 +75,29 @@
            (len2 (string-length s2)))
       (if (= len2 0)
         start1
-        (let lp ((i start1))
-          (cond
-            ((> (+ i len2) end1) #f)
-            ((string=? (substring s1 i (+ i len2)) s2) i)
-            (else (lp (+ i 1))))))))
+        (let ((first (string-ref s2 0))
+              (limit (- end1 len2)))
+          (let lp ((i start1))
+            (cond
+              ((> i limit) #f)
+              ((and (char=? (string-ref s1 i) first)
+                    (string-region=? s1 i s2 0 len2))
+               i)
+              (else (lp (+ i 1)))))))))
 
   (def (string-prefix? prefix str . rest)
     (let* ((start (if (pair? rest) (car rest) 0))
            (plen (string-length prefix))
            (slen (string-length str)))
       (and (<= (+ start plen) slen)
-           (string=? prefix (substring str start (+ start plen))))))
+           (string-region=? prefix 0 str start plen))))
 
   (def (string-suffix? suffix str . rest)
     (let* ((slen (string-length str))
            (end (if (pair? rest) (car rest) slen))
            (suflen (string-length suffix)))
       (and (<= suflen end)
-           (string=? suffix (substring str (- end suflen) end)))))
+           (string-region=? suffix 0 str (- end suflen) suflen))))
 
   (def (string-trim str . rest)
     (let ((pred (if (pair? rest) (car rest) char-whitespace?)))
@@ -131,10 +142,29 @@
         ((null? lst) "")
         ((null? (cdr lst)) (car lst))
         (else
-         (let lp ((rest (cdr lst)) (acc (car lst)))
-           (if (null? rest)
-             acc
-             (lp (cdr rest) (string-append acc sep (car rest)))))))))
+         (let* ((sep-len (string-length sep))
+                (total
+                 (let lp ((rest lst) (n 0) (first? #t))
+                   (if (null? rest)
+                     n
+                     (lp (cdr rest)
+                         (+ n
+                            (string-length (car rest))
+                            (if first? 0 sep-len))
+                         #f))))
+                (out (make-string total)))
+           (let lp ((rest lst) (pos 0) (first? #t))
+             (if (null? rest)
+               out
+               (let* ((pos (if first?
+                             pos
+                             (begin
+                               (string-copy! sep 0 out pos sep-len)
+                               (+ pos sep-len))))
+                      (s (car rest))
+                      (slen (string-length s)))
+                 (string-copy! s 0 out pos slen)
+                 (lp (cdr rest) (+ pos slen) #f)))))))))
 
   (def (string-concatenate lst)
     (apply string-append lst))
diff --git a/tests/test-better.ss b/tests/test-better.ss
index 3aa8549..789fe46 100644
--- a/tests/test-better.ss
+++ b/tests/test-better.ss
@@ -512,12 +512,15 @@
 (check-false (string-suffix? "xyz" "hello"))
 (check-true (string-contains? "ell" "hello"))
 (check-false (string-contains? "xyz" "hello"))
+(check-true (string-contains? "" "hello"))
 (check (string-trim-both "  hello  ") => "hello")
 (check (string-trim-both "hello") => "hello")
 (check (string-join '("a" "b" "c") ", ") => "a, b, c")
+(check (string-join '("a" "b" "c") "") => "abc")
 (check (string-join '() ", ") => "")
 (check (string-repeat "ab" 3) => "ababab")
 (check (string-repeat "x" 0) => "")
+(check (string-repeat "" 5) => "")
 (check (string-index "hello" #\l) => 2)
 (check (string-index "hello" #\z) => #f)
 (check (string-index-right "hello" #\l) => 3)
@@ -527,6 +530,9 @@
 (check (string-count "hello" #\l) => 2)
 (check (string-take-while "aaabbb" (lambda (c) (char=? c #\a))) => "aaa")
 (check (string-drop-while "aaabbb" (lambda (c) (char=? c #\a))) => "bbb")
+(check (string-split "abc" "") => '("abc"))
+(check (string-replace "abcabc" "b" "") => "acac")
+(check (string-replace "aaaa" "aa" "b") => "bb")
 
 ;; ========== #30: list-more ==========
 (printf "  #30 list-more...~n")
diff --git a/tests/test-stdlib.ss b/tests/test-stdlib.ss
index 0c14824..288b23d 100644
--- a/tests/test-stdlib.ss
+++ b/tests/test-stdlib.ss
@@ -104,13 +104,16 @@
 
 (check (string-split "a,b,c" #\,) => '("a" "b" "c"))
 (check (string-split "a::b::c" "::") => '("a" "b" "c"))
+(check (string-split "abc" "") => '("a" "b" "c"))
 (check (string-split "hello") => '("hello"))
 (check (string-join '("a" "b" "c") ",") => "a,b,c")
+(check (string-join '("a" "b" "c") "") => "abc")
 (check (string-join '("hello") " ") => "hello")
 (check (string-trim "  hello  ") => "hello")
 (check (string-prefix? "he" "hello") => #t)
 (check (string-prefix? "xx" "hello") => #f)
 (check (string-suffix? "lo" "hello") => #t)
+(check (string-contains "hello" "") => 0)
 (check (string-contains "hello world" "world") => 6)
 (check (string-contains "hello" "xyz") => #f)
 (check (string-index "hello" #\l) => 2)