_sre: drive an all-ASCII str subject over its bytes (#8520)

* sre_engine: document that StrDrive positions are character indices

StrDrive carried no documentation, so nothing recorded that a cursor's
position is a character index rather than a byte offset. Every implementation
satisfies it -- skip(n) advances position by exactly n in all three -- and the
engine depends on it: _count bounds a repeat with position + max_count and
reports the repeat length as a difference of positions, ASSERT compares
position against a lookbehind width, and search_info recovers a match start as
position - (len - 1). A drive over a variable-width encoding that stored byte
offsets would leave those type-correct and wrong, and would index a lookbehind
out of bounds.

Write the invariant down on StringCursor and on each trait method.

Also walk the group's own cursor in GROUPREF instead of counting to the
group's width, so the loop bound is the thing being stepped. Generated
machine code is unchanged.

Assisted-by: Claude

* _sre: drive an all-ASCII str subject over its bytes

with_sre_str handed every str subject to the &Wtf8 drive, which answers
count() by counting every code point and create_cursor(n) by stepping over the
first n. Both run once per Request, so a scan that restarts at successive
positions walked the subject again on every call: finditer over an ASCII
subject of n tokens was quadratic in n.

PyStr already records whether it is ASCII -- StrKind is decided when the
string is built -- and for ASCII a character index is a byte index, so the
&[u8] drive's cursor arithmetic already applies: count() is the byte length
and create_cursor() is a pointer offset. Add AsciiStr, which delegates every
StrDrive method to that impl and differs only in slice(), which reslices the
span and returns str rather than bytes.

Matching is unaffected: StrDrive carries no unicode semantics, because the
engine keys every unicode decision on the compiled pattern's opcode rather
than on the subject type.

Also bind the subject once in with_sre_str, so callers passing a temporary
(`&x.clone()`) build it once rather than per arm.

finditer over an ASCII subject, n tokens, this machine:

    n       before      after
     5000    41.60ms     1.37ms
    10000   217.30ms     2.68ms
    20000  1206.28ms     5.17ms
    40000  5016.65ms    10.31ms

Per-doubling x5.22/x5.55/x4.16 becomes x1.95/x1.93/x1.99. Collecting
m.group(0) for every match goes 5061.02ms -> 17.80ms at n=40000.

test.test_re is unchanged at 166 tests, OK (skipped=14, expected failures=6),
and a 2731-line differential over the is_ascii() boundary -- findall,
finditer spans and groups, sub, split, match, search, fullmatch across
subjects that are empty, ASCII, non-ASCII, or mixed -- is byte-identical to
CPython 3.14.6. Introducing an off-by-one in AsciiStr::slice moves 1224 of
those lines, so the comparison reaches the new code.

Assisted-by: Claude

crates/sre_engine/src/engine.rs

@@ -581,7 +581,11 @@ fn _match<S: StrDrive>(req: &Request<'_, S>, state: &mut State, mut ctx: MatchCo
                                 ..ctx
                             };
 
-                            for _ in group_start..group_end {
+                            // Walk the group itself rather than counting to its
+                            // width: `g_ctx` is already stepping over exactly
+                            // the characters being compared, so its own cursor
+                            // is the loop bound.
+                            while g_ctx.cursor.position < group_end {
                                 #[allow(clippy::redundant_closure_call)]
                                 if ctx.at_end(req)
                                     || $f(ctx.peek_char::<S>()) != $f(g_ctx.peek_char::<S>())

crates/sre_engine/src/string.rs

@@ -1,5 +1,12 @@
 use rustpython_wtf8::Wtf8;
 
+/// A position in the subject, paired with the byte pointer it resolves to.
+///
+/// `position` is a **character index**, never a byte offset. The engine does
+/// arithmetic on it directly — it subtracts two positions to get a character
+/// count, adds a repeat count to get a bound, and compares one against a
+/// lookbehind width — so the unit is part of the [`StrDrive`] contract rather
+/// than a detail each implementation may pick.
 #[derive(Debug, Clone, Copy)]
 pub struct StringCursor {
     pub(crate) ptr: *const u8,
@@ -15,15 +22,43 @@ impl Default for StringCursor {
     }
 }
 
+/// Random access over the subject being matched.
+///
+/// An implementation chooses how a character is spelled in memory — one byte
+/// for `&[u8]`, one code point for `&str` and `&Wtf8` — but **not** how
+/// positions are counted. Every position this trait produces or consumes is a
+/// character index: `count` is the subject's length in characters, and
+/// `skip(n)` advances a cursor's `position` by exactly `n`.
+///
+/// That is load-bearing, not incidental. The engine reads position arithmetic
+/// as character arithmetic in several places — `_count` bounds a repeat with
+/// `position + max_count` and reports the repeat's length as a difference of
+/// positions, `ASSERT` tests `position < back` against a lookbehind width, and
+/// `search_info` recovers a match start as `position - (len - 1)`. A drive
+/// that stored byte offsets here would leave all of those type-correct and
+/// silently wrong, and would index a lookbehind out of bounds.
+///
+/// So a drive over a variable-width encoding pays for the mapping: `count`
+/// and `create_cursor` have to resolve character indices, and cannot simply
+/// hand back byte lengths and byte offsets.
 pub trait StrDrive: Copy {
+    /// The subject's length, in characters.
     fn count(&self) -> usize;
+    /// A cursor at character index `n`.
     fn create_cursor(&self, n: usize) -> StringCursor;
+    /// Move `cursor` to character index `n`, from wherever it is now.
     fn adjust_cursor(&self, cursor: &mut StringCursor, n: usize);
+    /// Consume one character, returning it; `position` grows by one.
     fn advance(cursor: &mut StringCursor) -> u32;
+    /// The character at `cursor`, without moving it.
     fn peek(cursor: &StringCursor) -> u32;
+    /// Skip `n` characters, so `position` grows by exactly `n`.
     fn skip(cursor: &mut StringCursor, n: usize);
+    /// Step back over one character, returning it; `position` shrinks by one.
     fn back_advance(cursor: &mut StringCursor) -> u32;
+    /// The character before `cursor`, without moving it.
     fn back_peek(cursor: &StringCursor) -> u32;
+    /// Step back `n` characters, so `position` shrinks by exactly `n`.
     fn back_skip(cursor: &mut StringCursor, n: usize);
 }
 

crates/vm/src/stdlib/_sre.rs

@@ -22,7 +22,7 @@ mod _sre {
     use itertools::Itertools;
     use num_traits::ToPrimitive;
     use rustpython_sre_engine::{
-        Request, SearchIter, SreFlag, State, StrDrive,
+        Request, SearchIter, SreFlag, State, StrDrive, StringCursor,
         string::{lower_ascii, lower_unicode},
     };
 
@@ -83,6 +83,72 @@ mod _sre {
         }
     }
 
+    /// An all-ASCII `str` subject, driven over its bytes.
+    ///
+    /// For ASCII a character index *is* a byte index, so `&[u8]`'s cursor
+    /// arithmetic is already the right arithmetic: `count` is the byte length
+    /// and `create_cursor` is a pointer offset. The `&Wtf8` drive has to count
+    /// code points from the start of the subject to answer either, once per
+    /// `Request`, which makes a scan that restarts at successive positions --
+    /// `finditer`, or `re` module functions called in a loop -- walk the
+    /// subject again on every call.
+    ///
+    /// Matching is unaffected: `StrDrive` carries no unicode semantics of its
+    /// own, because the engine keys every unicode decision on the compiled
+    /// pattern's opcode rather than on the subject type. Only `slice` differs
+    /// from the `&[u8]` impl, to hand back `str` instead of `bytes`.
+    #[derive(Clone, Copy)]
+    struct AsciiStr<'a>(&'a [u8]);
+
+    impl StrDrive for AsciiStr<'_> {
+        fn count(&self) -> usize {
+            <&[u8] as StrDrive>::count(&self.0)
+        }
+
+        fn create_cursor(&self, n: usize) -> StringCursor {
+            <&[u8] as StrDrive>::create_cursor(&self.0, n)
+        }
+
+        fn adjust_cursor(&self, cursor: &mut StringCursor, n: usize) {
+            <&[u8] as StrDrive>::adjust_cursor(&self.0, cursor, n)
+        }
+
+        fn advance(cursor: &mut StringCursor) -> u32 {
+            <&[u8] as StrDrive>::advance(cursor)
+        }
+
+        fn peek(cursor: &StringCursor) -> u32 {
+            <&[u8] as StrDrive>::peek(cursor)
+        }
+
+        fn skip(cursor: &mut StringCursor, n: usize) {
+            <&[u8] as StrDrive>::skip(cursor, n)
+        }
+
+        fn back_advance(cursor: &mut StringCursor) -> u32 {
+            <&[u8] as StrDrive>::back_advance(cursor)
+        }
+
+        fn back_peek(cursor: &StringCursor) -> u32 {
+            <&[u8] as StrDrive>::back_peek(cursor)
+        }
+
+        fn back_skip(cursor: &mut StringCursor, n: usize) {
+            <&[u8] as StrDrive>::back_skip(cursor, n)
+        }
+    }
+
+    impl SreStr for AsciiStr<'_> {
+        fn slice(&self, start: usize, end: usize, vm: &VirtualMachine) -> PyObjectRef {
+            let end = end.min(self.0.len());
+            let start = start.min(end);
+            // The subject is ASCII, so any span of it is valid UTF-8 and the
+            // span is a reslice rather than a walk from the subject's start.
+            let s = str::from_utf8(&self.0[start..end]).expect("ascii subject");
+            vm.ctx.new_str(s).into()
+        }
+    }
+
     #[pyfunction]
     fn compile(
         pattern: PyObjectRef,
@@ -200,13 +266,18 @@ mod _sre {
     }
 
     macro_rules! with_sre_str {
-        ($pattern:expr, $string:expr, $vm:expr, $f:expr) => {
+        ($pattern:expr, $string:expr, $vm:expr, $f:expr) => {{
+            // Bind once: the branches only borrow the subject, and callers pass
+            // a temporary (`&x.clone()`) that would otherwise be rebuilt per arm.
+            let subject = $string;
             if $pattern.isbytes {
-                Pattern::with_bytes($string, $vm, $f)
+                Pattern::with_bytes(subject, $vm, $f)
+            } else if Pattern::is_ascii_str(subject) {
+                Pattern::with_ascii_str(subject, $vm, $f)
             } else {
-                Pattern::with_str($string, $vm, $f)
+                Pattern::with_str(subject, $vm, $f)
             }
-        };
+        }};
     }
 
     #[pyclass(with(Hashable, Comparable, Representable), flags(HAS_WEAKREF))]
@@ -221,6 +292,27 @@ mod _sre {
             f(string.as_wtf8())
         }
 
+        /// Whether a `str` subject can take the [`AsciiStr`] drive.
+        ///
+        /// `PyStr` already knows: `StrKind` is decided when the string is
+        /// built, so this is a field load rather than a scan. A non-`str`
+        /// argument answers `false` and is reported by [`Self::with_str`].
+        fn is_ascii_str(string: &PyObject) -> bool {
+            string
+                .downcast_ref::<PyStr>()
+                .is_some_and(|s| s.kind().is_ascii())
+        }
+
+        fn with_ascii_str<F, R>(string: &PyObject, vm: &VirtualMachine, f: F) -> PyResult<R>
+        where
+            F: FnOnce(AsciiStr<'_>) -> PyResult<R>,
+        {
+            let string = string.downcast_ref::<PyStr>().ok_or_else(|| {
+                vm.new_type_error(format!("expected string got '{}'", string.class()))
+            })?;
+            f(AsciiStr(string.as_wtf8().as_bytes()))
+        }
+
         fn with_bytes<F, R>(string: &PyObject, vm: &VirtualMachine, f: F) -> PyResult<R>
         where
             F: FnOnce(&[u8]) -> PyResult<R>,