perf: optimize String.charactersIn (#14395)

This PR attempts to make `String.charactersIn` faster by doing less work
(and in particular only `USize` and `UInt8` arithmetic, not `Nat`) in
the hot loop.

Author: TwoFX

src/Init/Data/UInt/Basic.lean

@@ -190,6 +190,15 @@ def UInt8.toAsciiLower (b : UInt8) : UInt8 :=
   -- to generate slighly better assembly.
   b + ((decide (b - 65 < 26)).toUInt8 <<< 5)
 
+/--
+If `b` is the ASCII value of a lowercase character return the corresponding
+uppercase value, otherwise leave it untouched.
+-/
+@[inline]
+def UInt8.toAsciiUpper (b : UInt8) : UInt8 :=
+  -- See comment on `toAsciiLower` above.
+  b - ((decide (b - 97 < 26)).toUInt8 <<< 5)
+
 /-- Converts a `Fin UInt16.size` into the corresponding `UInt16`. -/
 @[inline] def UInt16.ofFin (a : Fin UInt16.size) : UInt16 := ⟨⟨a⟩⟩
 

src/Init/Data/UInt/Lemmas.lean

@@ -3596,3 +3596,9 @@ theorem UInt64.lt_add_one {c : UInt64} (h : c ≠ -1) : c < c + 1 :=
 theorem USize.lt_add_one {c : USize} (h : c ≠ -1) : c < c + 1 :=
   USize.lt_iff_toBitVec_lt.2 (BitVec.lt_add_one
     (by simpa [← USize.toBitVec_inj, BitVec.neg_one_eq_allOnes] using h))
+
+theorem UInt8.toNat_ofNat_le {n : Nat} : (UInt8.ofNat n).toNat ≤ n := toNat_ofNat ▸ Nat.mod_le ..
+theorem UInt16.toNat_ofNat_le {n : Nat} : (UInt16.ofNat n).toNat ≤ n := toNat_ofNat ▸ Nat.mod_le ..
+theorem UInt32.toNat_ofNat_le {n : Nat} : (UInt32.ofNat n).toNat ≤ n := toNat_ofNat ▸ Nat.mod_le ..
+theorem UInt64.toNat_ofNat_le {n : Nat} : (UInt64.ofNat n).toNat ≤ n := toNat_ofNat ▸ Nat.mod_le ..
+theorem USize.toNat_ofNat_le {n : Nat} : (USize.ofNat n).toNat ≤ n := toNat_ofNat ▸ Nat.mod_le ..

src/Lean/Server/Completion/CompletionUtils.lean

@@ -7,33 +7,55 @@ module
 
 prelude
 public import Lean.Meta.WHNF
+import Init.Data.String.Bootstrap
+import Init.Data.UInt.Lemmas
 
 public section
 
+/--
+Returns `true` if the characters of `a` occur in `b` in order (not necessarily contiguously),
+compared case-insensitively via `Char.toLower`.
+
+This function assumes that `a` and `b` have size at most `USize.size`, which should be a safe
+assumption.
+-/
 partial def Lean.String.charactersIn (a b : String) : Bool :=
-  goFastScalar ⟨0⟩ ⟨0⟩
+  goFastScalar a.utf8ByteSize.toUSize b.utf8ByteSize.toUSize 0 0 rfl rfl
 where
   /-
-  This function is the ASCII fast path for `go`
+  This function is the ASCII fast path for `go`. It caches the two case variants of the current
+  pattern byte in `lower`/`upper` (via `scan`) so that the hot loop scanning `b` only performs a
+  bounds check, a byte load, and two comparisons per byte.
   -/
-  goFastScalar (aPos bPos : String.Pos.Raw) : Bool :=
-    if ha : ¬aPos < a.rawEndPos then
-      true
-    else if hb : ¬bPos < b.rawEndPos then
-      false
+  goFastScalar (aSize bSize aPos bPos : USize)
+      (ha : aSize = a.utf8ByteSize.toUSize) (hb : bSize = b.utf8ByteSize.toUSize) : Bool :=
+    if haPos : aPos < aSize then
+      let aByte := String.Internal.ugetUTF8Byte a aPos
+        (Std.lt_of_lt_of_le (USize.lt_iff_toNat_lt.1 haPos) (ha ▸ USize.toNat_ofNat_le))
+      -- If `a` contains a non-ASCII byte we give up on the fast path
+      if aByte &&& 0x80 != 0 then
+        go ⟨aPos.toNat⟩ ⟨bPos.toNat⟩
+      else
+        scan aByte.toAsciiLower aByte.toAsciiUpper aSize bSize aPos bPos ha hb
     else
-      let aByte := a.getUTF8Byte aPos (by simpa using ha)
-      let bByte := b.getUTF8Byte bPos (by simpa using hb)
-      -- If a or b are not UTF-8 bytes we give up on the fast path
-      if (aByte &&& 0x80 != 0) || (bByte &&& 0x80 != 0) then
-        go aPos bPos
+      true
+
+  scan (lower upper : UInt8) (aSize bSize aPos bPos : USize)
+      (ha : aSize = a.utf8ByteSize.toUSize) (hb : bSize = b.utf8ByteSize.toUSize) : Bool :=
+    if hbPos : bPos < bSize then
+      let bByte := String.Internal.ugetUTF8Byte b bPos
+        (Std.lt_of_lt_of_le (USize.lt_iff_toNat_lt.1 hbPos) (hb ▸ USize.toNat_ofNat_le))
+      -- `lower` and `upper` are ASCII, so a non-ASCII byte can never match and it is fine to
+      -- check for a match first
+      if bByte == lower || bByte == upper then
+        goFastScalar aSize bSize (aPos + 1) (bPos + 1) ha hb
+      -- If `b` contains a non-ASCII byte we give up on the fast path
+      else if bByte &&& 0x80 != 0 then
+        go ⟨aPos.toNat⟩ ⟨bPos.toNat⟩
       else
-        let bPos := ⟨bPos.byteIdx + 1⟩
-        if aByte.toAsciiLower == bByte.toAsciiLower then
-          let aPos := ⟨aPos.byteIdx + 1⟩
-          goFastScalar aPos bPos
-        else
-          goFastScalar aPos bPos
+        scan lower upper aSize bSize aPos (bPos + 1) ha hb
+    else
+      false
 
   go (aPos bPos : String.Pos.Raw) : Bool :=
     if ha : aPos.atEnd a then