From c8c5a587ac637aa1521c17c631fe0070aa1dc994 Mon Sep 17 00:00:00 2001
From: Scott McMurray <scottmcm@users.noreply.github.com>
Date: Fri, 5 May 2023 02:29:40 -0700
Subject: [PATCH] Tune the `is_ascii` implementation used for short slices

---
 library/core/src/slice/ascii.rs  | 29 +++++++++++++++++---------
 library/core/src/slice/mod.rs    |  4 ++++
 tests/assembly/slice-is_ascii.rs | 35 ++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 10 deletions(-)
 create mode 100644 tests/assembly/slice-is_ascii.rs

diff --git a/library/core/src/slice/ascii.rs b/library/core/src/slice/ascii.rs
index 6a6c0c9ba8ba1..f3311f76a7f06 100644
--- a/library/core/src/slice/ascii.rs
+++ b/library/core/src/slice/ascii.rs
@@ -268,6 +268,24 @@ const fn contains_nonascii(v: usize) -> bool {
     (NONASCII_MASK & v) != 0
 }
 
+/// ASCII test *without* the chunk-at-a-time optimizations.
+///
+/// This is carefully structured to produce nice small code -- it's smaller in
+/// `-O` than what the "obvious" ways produces under `-C opt-level=s`.  If you
+/// touch it, be sure to run (and update if needed) the assembly test.
+#[unstable(feature = "str_internals", issue = "none")]
+#[doc(hidden)]
+#[inline]
+pub const fn is_ascii_simple(mut bytes: &[u8]) -> bool {
+    while let [rest @ .., last] = bytes {
+        if !last.is_ascii() {
+            break;
+        }
+        bytes = rest;
+    }
+    bytes.is_empty()
+}
+
 /// Optimized ASCII test that will use usize-at-a-time operations instead of
 /// byte-at-a-time operations (when possible).
 ///
@@ -293,16 +311,7 @@ const fn is_ascii(s: &[u8]) -> bool {
     // We also do this for architectures where `size_of::<usize>()` isn't
     // sufficient alignment for `usize`, because it's a weird edge case.
     if len < USIZE_SIZE || len < align_offset || USIZE_SIZE < mem::align_of::<usize>() {
-        // FIXME: once iterators and closures can be used in `const fn`,
-        // return s.iter().all(|b| b.is_ascii());
-        let mut i = 0;
-        while i < len {
-            if !s[i].is_ascii() {
-                return false;
-            }
-            i += 1;
-        }
-        return true;
+        return is_ascii_simple(s);
     }
 
     // We always read the first word unaligned, which means `align_offset` is
diff --git a/library/core/src/slice/mod.rs b/library/core/src/slice/mod.rs
index d4981af90d19e..4c891ba550f3c 100644
--- a/library/core/src/slice/mod.rs
+++ b/library/core/src/slice/mod.rs
@@ -44,6 +44,10 @@ mod raw;
 mod rotate;
 mod specialize;
 
+#[unstable(feature = "str_internals", issue = "none")]
+#[doc(hidden)]
+pub use ascii::is_ascii_simple;
+
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use iter::{Chunks, ChunksMut, Windows};
 #[stable(feature = "rust1", since = "1.0.0")]
diff --git a/tests/assembly/slice-is_ascii.rs b/tests/assembly/slice-is_ascii.rs
new file mode 100644
index 0000000000000..b3e1fee15a715
--- /dev/null
+++ b/tests/assembly/slice-is_ascii.rs
@@ -0,0 +1,35 @@
+// revisions: WIN LIN
+// [WIN] only-windows
+// [LIN] only-linux
+// assembly-output: emit-asm
+// compile-flags: --crate-type=lib -O -C llvm-args=-x86-asm-syntax=intel
+// min-llvm-version: 14
+// only-x86_64
+// ignore-sgx
+// ignore-debug
+
+#![feature(str_internals)]
+
+// CHECK-LABEL: is_ascii_simple_demo:
+#[no_mangle]
+pub fn is_ascii_simple_demo(bytes: &[u8]) -> bool {
+    // Linux (System V): pointer is rdi; length is rsi
+    // Windows: pointer is rcx; length is rdx.
+
+    // CHECK-NOT: mov
+    // CHECK-NOT: test
+    // CHECK-NOT: cmp
+
+    // CHECK: .[[LOOPHEAD:.+]]:
+    // CHECK-NEXT: mov [[TEMP:.+]], [[LEN:rsi|rdx]]
+    // CHECK-NEXT: sub [[LEN]], 1
+    // CHECK-NEXT: jb .[[LOOPEXIT:.+]]
+    // CHECK-NEXT: cmp byte ptr [{{rdi|rcx}} + [[TEMP]] - 1], 0
+    // CHECK-NEXT: jns .[[LOOPHEAD]]
+
+    // CHECK-NEXT: .[[LOOPEXIT]]:
+    // CHECK-NEXT: test [[TEMP]], [[TEMP]]
+    // CHECK-NEXT: sete al
+    // CHECK-NEXT: ret
+    core::slice::is_ascii_simple(bytes)
+}