diff --git a/crates/sre_engine/src/engine.rs b/crates/sre_engine/src/engine.rs index 690801e0d9d..c2a6ac81975 100644 --- a/crates/sre_engine/src/engine.rs +++ b/crates/sre_engine/src/engine.rs @@ -581,7 +581,11 @@ fn _match(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::()) != $f(g_ctx.peek_char::()) diff --git a/crates/sre_engine/src/string.rs b/crates/sre_engine/src/string.rs index ca7303a2a7f..5cc1b04b9fc 100644 --- a/crates/sre_engine/src/string.rs +++ b/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); } diff --git a/crates/vm/src/stdlib/_sre.rs b/crates/vm/src/stdlib/_sre.rs index 18b4ffde818..6fe6b434702 100644 --- a/crates/vm/src/stdlib/_sre.rs +++ b/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::() + .is_some_and(|s| s.kind().is_ascii()) + } + + fn with_ascii_str(string: &PyObject, vm: &VirtualMachine, f: F) -> PyResult + where + F: FnOnce(AsciiStr<'_>) -> PyResult, + { + let string = string.downcast_ref::().ok_or_else(|| { + vm.new_type_error(format!("expected string got '{}'", string.class())) + })?; + f(AsciiStr(string.as_wtf8().as_bytes())) + } + fn with_bytes(string: &PyObject, vm: &VirtualMachine, f: F) -> PyResult where F: FnOnce(&[u8]) -> PyResult,