From b2714d36efea8d06fa3d416b05bff1b58c0c1554 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Tue, 11 Aug 2020 14:43:05 +0200
Subject: [PATCH] fix: String: take/drop characters, not bytes

---
 src/Init/Data/String/Basic.lean          | 24 ++++++++++++------------
 tests/lean/string_imp2.lean              |  2 ++
 tests/lean/string_imp2.lean.expected.out |  2 ++
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/src/Init/Data/String/Basic.lean b/src/Init/Data/String/Basic.lean
index 94301f02ce..54ccd1c708 100644
--- a/src/Init/Data/String/Basic.lean
+++ b/src/Init/Data/String/Basic.lean
@@ -368,6 +368,14 @@ namespace Substring
 | ⟨s, b, _⟩, p =>
   if p = b then p else s.prev (b+p) - b
 
+def nextn : Substring → Nat → String.Pos → String.Pos
+| ss, 0,   p => p
+| ss, i+1, p => ss.nextn i (ss.next p)
+
+def prevn : Substring → String.Pos → Nat → String.Pos
+| ss, 0,   p => p
+| ss, i+1, p => ss.prevn i (ss.prev p)
+
 @[inline] def front (s : Substring) : Char :=
 s.get 0
 
@@ -376,24 +384,16 @@ match s with
 | ⟨s, b, e⟩ => (String.posOfAux s c e b) - b
 
 @[inline] def drop : Substring → Nat → Substring
-| ⟨s, b, e⟩, n =>
-  if b + n ≥ e then "".toSubstring
-  else ⟨s, b+n, e⟩
+| ss@⟨s, b, e⟩, n => ⟨s, ss.nextn n b, e⟩
 
 @[inline] def dropRight : Substring → Nat → Substring
-| ⟨s, b, e⟩, n =>
-  if e - n ≤ b then "".toSubstring
-  else ⟨s, b, e - n⟩
+| ss@⟨s, b, e⟩, n => ⟨s, b, ss.prevn n e⟩
 
 @[inline] def take : Substring → Nat → Substring
-| ⟨s, b, e⟩, n =>
-  let e := if b + n ≥ e then e else b + n;
-  ⟨s, b, e⟩
+| ss@⟨s, b, e⟩, n => ⟨s, b, ss.nextn n b⟩
 
 @[inline] def takeRight : Substring → Nat → Substring
-| ⟨s, b, e⟩, n =>
-  let b := if e - n ≤ b then b else e - n;
-  ⟨s, b, e⟩
+| ss@⟨s, b, e⟩, n => ⟨s, ss.prevn n e, e⟩
 
 @[inline] def atEnd : Substring → String.Pos → Bool
 | ⟨s, b, e⟩, p => b + p == e
diff --git a/tests/lean/string_imp2.lean b/tests/lean/string_imp2.lean
index 722041021c..7bfd1d5ebe 100644
--- a/tests/lean/string_imp2.lean
+++ b/tests/lean/string_imp2.lean
@@ -46,3 +46,5 @@ it₁.remainingToString ++ "-" ++ it₂.remainingToString
 #eval "αβ".mkIterator.next.prev.hasPrev
 #eval "abc" == "abc"
 #eval "abc" == "abd"
+#eval "αβγ".drop 1
+#eval "αβγ".takeRight 1
diff --git a/tests/lean/string_imp2.lean.expected.out b/tests/lean/string_imp2.lean.expected.out
index 0b0ec8edef..c2c60493ed 100644
--- a/tests/lean/string_imp2.lean.expected.out
+++ b/tests/lean/string_imp2.lean.expected.out
@@ -35,3 +35,5 @@ true
 false
 true
 false
+"βγ"
+"γ"