fix: logic error in String.Slice.takeWhile (#10868)

This PR fixes a bug in `String.Slice.takeWhile` which caused it to get
its bookkeeping wrong and panic. The new version only uses safe
operations on `String.Slice.Pos`.

(cherry picked from commit 196d501)

Author: TwoFX

src/Init/Data/String/Basic.lean

@@ -1870,7 +1870,7 @@ theorem Pos.Raw.unoffsetBy_offsetBy {p q : Pos.Raw} : (p.offsetBy q).unoffsetBy
 /-- Given a position in `s` that is at least `p₀`, obtain the corresponding position in
 `s.replaceStart p₀`. -/
 @[inline]
-def Slice.Pos.toReplaceStart {s : Slice} (p₀ : s.Pos) (pos : s.Pos) (h : p₀.offset ≤ pos.offset) :
+def Slice.Pos.toReplaceStart {s : Slice} (p₀ : s.Pos) (pos : s.Pos) (h : p₀ ≤ pos) :
     (s.replaceStart p₀).Pos where
   offset := pos.offset.unoffsetBy p₀.offset
   isValidForSlice := Pos.Raw.isValidForSlice_replaceStart.2 (by
@@ -1904,6 +1904,27 @@ theorem Slice.Pos.get_eq_get_ofReplaceStart {s : Slice} {p₀ : s.Pos} {pos : (s
     pos.get h = (ofReplaceStart pos).get (by rwa [← ofReplaceStart_endPos, ne_eq, ofReplaceStart_inj]) := by
   simp [Slice.Pos.get, Nat.add_assoc]
 
+/-- Given a position in `s.replaceEnd p₀`, obtain the corresponding position in `s`. -/
+@[inline]
+def Slice.Pos.ofReplaceEnd {s : Slice} {p₀ : s.Pos} (pos : (s.replaceEnd p₀).Pos) : s.Pos where
+  offset := pos.offset
+  isValidForSlice := (Pos.Raw.isValidForSlice_replaceEnd.1 pos.isValidForSlice).2
+
+@[simp]
+theorem Slice.Pos.offset_ofReplaceEnd {s : Slice} {p₀ : s.Pos} {pos : (s.replaceEnd p₀).Pos} :
+    (ofReplaceEnd pos).offset = pos.offset := (rfl)
+
+/-- Given a position in `s` that is at most `p₀`, obtain the corresponding position in `s.replaceEnd p₀`. -/
+@[inline]
+def Slice.Pos.toReplaceEnd {s : Slice} (p₀ : s.Pos) (pos : s.Pos) (h : pos ≤ p₀) :
+    (s.replaceEnd p₀).Pos where
+  offset := pos.offset
+  isValidForSlice := Pos.Raw.isValidForSlice_replaceEnd.2 ⟨h, pos.isValidForSlice⟩
+
+@[simp]
+theorem Slice.Pos.offset_toReplaceEnd {s : Slice} {p₀ : s.Pos} {pos : s.Pos} {h : pos ≤ p₀} :
+    (toReplaceEnd p₀ pos h).offset = pos.offset := (rfl)
+
 theorem Slice.Pos.copy_eq_append_get {s : Slice} {pos : s.Pos} (h : pos ≠ s.endPos) :
     ∃ t₁ t₂ : String, s.copy = t₁ ++ singleton (pos.get h) ++ t₂ ∧ t₁.utf8ByteSize = pos.offset.byteIdx := by
   obtain ⟨t₂, ht₂⟩ := (s.replaceStart pos).copy.eq_singleton_append (by simpa [← ValidPos.ofCopy_inj, ← ofReplaceStart_inj])

src/Init/Data/String/Pattern/Basic.lean

@@ -70,7 +70,7 @@ class ForwardPattern (ρ : Type) where
   Checks whether the slice starts with the pattern. If it does, the slice is returned with the
   prefix removed; otherwise the result is {name}`none`.
   -/
-  dropPrefix? : Slice → ρ → Option Slice
+  dropPrefix? : (s : Slice) → ρ → Option s.Pos
 
 namespace Internal
 
@@ -153,10 +153,10 @@ def defaultStartsWith (s : Slice) (pat : ρ) : Bool :=
   | _ => false
 
 @[specialize pat]
-def defaultDropPrefix? (s : Slice) (pat : ρ) : Option Slice :=
+def defaultDropPrefix? (s : Slice) (pat : ρ) : Option s.Pos :=
   let searcher := ToForwardSearcher.toSearcher s pat
   match searcher.step with
-  | .yield _ (.matched _ endPos) _ => some (s.replaceStart endPos)
+  | .yield _ (.matched _ endPos) _ => some endPos
   | _ => none
 
 @[always_inline, inline]
@@ -196,7 +196,7 @@ class BackwardPattern (ρ : Type) where
   Checks whether the slice ends with the pattern. If it does, the slice is returned with the
   suffix removed; otherwise the result is {name}`none`.
   -/
-  dropSuffix? : Slice → ρ → Option Slice
+  dropSuffix? : (s : Slice) → ρ → Option s.Pos
 
 namespace ToBackwardSearcher
 
@@ -212,10 +212,10 @@ def defaultEndsWith (s : Slice) (pat : ρ) : Bool :=
   | _ => false
 
 @[specialize pat]
-def defaultDropSuffix? (s : Slice) (pat : ρ) : Option Slice :=
+def defaultDropSuffix? (s : Slice) (pat : ρ) : Option s.Pos :=
   let searcher := ToBackwardSearcher.toSearcher s pat
   match searcher.step with
-  | .yield _ (.matched startPos _) _ => some (s.replaceEnd startPos)
+  | .yield _ (.matched startPos _) _ => some startPos
   | _ => none
 
 @[always_inline, inline]

src/Init/Data/String/Pattern/String.lean

@@ -218,9 +218,9 @@ def startsWith (s : Slice) (pat : Slice) : Bool :=
     false
 
 @[inline]
-def dropPrefix? (s : Slice) (pat : Slice) : Option Slice :=
+def dropPrefix? (s : Slice) (pat : Slice) : Option s.Pos :=
   if startsWith s pat then
-    some <| s.replaceStart <| s.pos! <| pat.rawEndPos.offsetBy s.startPos.offset
+    some <| s.pos! <| pat.rawEndPos.offsetBy s.startPos.offset
   else
     none
 
@@ -254,9 +254,9 @@ def endsWith (s : Slice) (pat : Slice) : Bool :=
     false
 
 @[inline]
-def dropSuffix? (s : Slice) (pat : Slice) : Option Slice :=
+def dropSuffix? (s : Slice) (pat : Slice) : Option s.Pos :=
   if endsWith s pat then
-    some <| s.replaceEnd <| s.pos! <| s.endPos.offset.unoffsetBy pat.rawEndPos
+    some <| s.pos! <| s.endPos.offset.unoffsetBy pat.rawEndPos
   else
     none
 

src/Init/Data/String/Slice.lean

@@ -251,7 +251,7 @@ Examples:
 -/
 @[inline]
 def dropPrefix? [ForwardPattern ρ] (s : Slice) (pat : ρ) : Option Slice :=
-  ForwardPattern.dropPrefix? s pat
+  (ForwardPattern.dropPrefix? s pat).map s.replaceStart
 
 /--
 If {name}`pat` matches a prefix of {name}`s`, returns the remainder. Returns {name}`s` unmodified
@@ -357,18 +357,18 @@ Examples:
 -/
 @[inline]
 partial def takeWhile [ForwardPattern ρ] (s : Slice) (pat : ρ) : Slice :=
-  go s
+  go s.startPos
 where
   @[specialize pat]
-  go (curr : Slice) : Slice :=
-    if let some nextCurr := dropPrefix? curr pat then
-      if curr.startInclusive.offset < nextCurr.startInclusive.offset then
+  go (curr : s.Pos) : Slice :=
+    if let some nextCurr := ForwardPattern.dropPrefix? (s.replaceStart curr) pat then
+      if (s.replaceStart curr).startPos < nextCurr then
         -- TODO: need lawful patterns to show this terminates
-        go nextCurr
+        go (Pos.ofReplaceStart nextCurr)
       else
-        s.replaceEnd <| s.pos! <| curr.startInclusive.offset
+        s.replaceEnd curr
     else
-      s.replaceEnd <| s.pos! <| curr.startInclusive.offset
+      s.replaceEnd curr
 
 /--
 Finds the position of the first match of the pattern {name}`pat` in a slice {name}`true`. If there
@@ -524,7 +524,7 @@ Examples:
 -/
 @[inline]
 def dropSuffix? [BackwardPattern ρ] (s : Slice) (pat : ρ) : Option Slice :=
-  BackwardPattern.dropSuffix? s pat
+  (BackwardPattern.dropSuffix? s pat).map s.replaceEnd
 
 /--
 If {name}`pat` matches a suffix of {name}`s`, returns the remainder. Returns {name}`s` unmodified
@@ -629,18 +629,18 @@ Examples:
 -/
 @[inline]
 partial def takeEndWhile [BackwardPattern ρ] (s : Slice) (pat : ρ) : Slice :=
-  go s
+  go s.endPos
 where
   @[specialize pat]
-  go (curr : Slice) : Slice :=
-    if let some nextCurr := dropSuffix? curr pat then
-      if nextCurr.endExclusive.offset < curr.endExclusive.offset then
+  go (curr : s.Pos) : Slice :=
+    if let some nextCurr := BackwardPattern.dropSuffix? (s.replaceEnd curr) pat then
+      if nextCurr < (s.replaceEnd curr).endPos then
         -- TODO: need lawful patterns to show this terminates
-        go nextCurr
+        go (Pos.ofReplaceEnd nextCurr)
       else
-        s.replaceStart <| s.pos! <| curr.endExclusive.offset
+        s.replaceStart curr
     else
-      s.replaceStart <| s.pos! <| curr.endExclusive.offset
+      s.replaceStart curr
 
 /--
 Finds the position of the first match of the pattern {name}`pat` in a slice {name}`true`, starting

tests/lean/run/string_slice.lean

@@ -55,6 +55,11 @@ Tests for `String.Slice` functions
 #guard "red red green blue".toSlice.takeWhile "red " == "red red ".toSlice
 #guard "red green blue".toSlice.takeWhile (fun (_ : Char) => true) == "red green blue".toSlice
 
+#guard (" ".toSlice.dropWhile ' ' |>.takeWhile Char.isLower) == "".toSlice
+#guard (" ".toSlice.dropWhile ' ' |>.takeEndWhile Char.isLower) == "".toSlice
+#guard ("∃a∃".toSlice.drop 1 |>.takeWhile Char.isLower) == "a".toSlice
+#guard ("∃a∃".toSlice.dropEnd 1 |>.takeEndWhile Char.isLower) == "a".toSlice
+
 #guard "red green blue".toSlice.dropPrefix? "red " == some "green blue".toSlice
 #guard "red green blue".toSlice.dropPrefix? "reed " == none
 #guard "red green blue".toSlice.dropPrefix? 'r' == some "ed green blue".toSlice