From 0ebff92181cdcb71acb468530ee049e4e9df8184 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Fri, 14 Aug 2026 17:51:28 +0900 Subject: [PATCH 1/2] common: index a WTF-8 buffer's code points for random access `Wtf8`'s iterators are sequential, so resolving a code point index through them is O(n) and code that indexes the same string repeatedly walks it once per index. `Wtf8Index` is a side table -- one 24-byte group per 64 code points, 0.375 bytes per code point -- that answers the same question in constant time; the layout is PyPy's `UTF8_INDEX_STORAGE`. `StrData` builds one on the first call to the new `char_index_to_byte`, in a slot published by compare-exchange, and drops it with the string. ASCII strings answer from the index itself and never build a table. A clone gets an empty slot, since it indexes its own copy of the buffer. Assisted-by: Claude --- crates/common/src/lib.rs | 1 + crates/common/src/str.rs | 87 +++++++++++- crates/common/src/wtf8_index.rs | 229 ++++++++++++++++++++++++++++++++ crates/vm/src/builtins/str.rs | 7 + 4 files changed, 322 insertions(+), 2 deletions(-) create mode 100644 crates/common/src/wtf8_index.rs diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 53a8e0d752b..d1e04b46d57 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -21,6 +21,7 @@ pub mod rc; pub mod refcount; pub mod static_cell; pub mod str; +pub mod wtf8_index; pub use rustpython_wtf8 as wtf8; diff --git a/crates/common/src/str.rs b/crates/common/src/str.rs index c006a5f4db4..8edb48fa4fa 100644 --- a/crates/common/src/str.rs +++ b/crates/common/src/str.rs @@ -1,7 +1,8 @@ // spell-checker:ignore uncomputed -use crate::atomic::{PyAtomic, Radium}; +use crate::atomic::{OncePtr, PyAtomic, Radium}; use crate::format::CharLen; use crate::wtf8::{CodePoint, Wtf8, Wtf8Buf}; +use crate::wtf8_index::Wtf8Index; use ascii::{AsciiChar, AsciiStr, AsciiString}; use core::fmt; use core::ops::{Bound, RangeBounds}; @@ -117,6 +118,55 @@ pub struct StrData { data: Box, kind: StrKind, len: StrLen, + index: Wtf8IndexSlot, +} + +/// A [`Wtf8Index`] built on first use. +/// +/// The table is a pure function of `data`, so publishing it races benignly: a +/// thread that loses the exchange drops its own copy and reads the winner's. +#[derive(Default)] +struct Wtf8IndexSlot(OncePtr); + +impl Wtf8IndexSlot { + #[inline(always)] + fn new() -> Self { + Self(OncePtr::new()) + } + + #[inline] + fn get_or_build(&self, data: &Wtf8, char_len: usize) -> &Wtf8Index { + let index = self + .0 + .get_or_init(|| Box::new(Wtf8Index::new(data, char_len))); + // The slot owns the table, never replaces it, and outlives the borrow. + unsafe { index.as_ref() } + } +} + +impl fmt::Debug for Wtf8IndexSlot { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self.0.get() { + Some(_) => f.write_str(""), + None => f.write_str(""), + } + } +} + +impl Clone for Wtf8IndexSlot { + /// A fresh slot: the clone copies the buffer, so it has to index that copy, + /// and the table is rebuilt on demand rather than eagerly here. + fn clone(&self) -> Self { + Self::new() + } +} + +impl Drop for Wtf8IndexSlot { + fn drop(&mut self) { + if let Some(index) = self.0.get() { + drop(unsafe { Box::from_raw(index.as_ptr()) }); + } + } } struct StrLen(PyAtomic); @@ -163,6 +213,7 @@ impl Default for StrData { data: >::default(), kind: StrKind::Ascii, len: StrLen::zero(), + index: Wtf8IndexSlot::new(), } } } @@ -193,6 +244,7 @@ impl From> for StrData { len: value.len().into(), data: value.into(), kind: StrKind::Ascii, + index: Wtf8IndexSlot::new(), } } } @@ -212,6 +264,7 @@ impl From for StrData { data: ch.to_string().into(), kind: StrKind::Utf8, len: 1.into(), + index: Wtf8IndexSlot::new(), } } } @@ -226,6 +279,7 @@ impl From for StrData { data: Wtf8Buf::from(ch).into(), kind: StrKind::Wtf8, len: 1.into(), + index: Wtf8IndexSlot::new(), } } } @@ -241,7 +295,12 @@ impl StrData { StrKind::Ascii => data.len().into(), _ => StrLen::uncomputed(), }; - Self { data, kind, len } + Self { + data, + kind, + len, + index: Wtf8IndexSlot::new(), + } } /// # Safety @@ -253,6 +312,7 @@ impl StrData { data, kind, len: char_len.into(), + index: Wtf8IndexSlot::new(), } } @@ -322,6 +382,29 @@ impl StrData { len } + /// The byte offset the `index`-th code point starts at. + /// + /// An `index` at or past the end answers the buffer's byte length, so a + /// caller walking to a bound does not have to special-case it. + /// + /// O(1), but the first call on a non-ASCII string builds an index over the + /// whole buffer, so a caller that resolves a single index and stops is + /// better served by [`Self::nth_char`]. + pub fn char_index_to_byte(&self, index: usize) -> usize { + // For ASCII the two units coincide, and the table would be a Nth entry + // saying N. + if self.kind.is_ascii() { + return index.min(self.data.len()); + } + let char_len = self.char_len(); + if index >= char_len { + return self.data.len(); + } + self.index + .get_or_build(&self.data, char_len) + .byte_offset(&self.data, index) + } + pub fn nth_char(&self, index: usize) -> CodePoint { match self.as_str_kind() { PyKindStr::Ascii(s) => s[index].into(), diff --git a/crates/common/src/wtf8_index.rs b/crates/common/src/wtf8_index.rs new file mode 100644 index 00000000000..a6121f1c03d --- /dev/null +++ b/crates/common/src/wtf8_index.rs @@ -0,0 +1,229 @@ +// spell-checker:ignore rpython rlib rutf +//! Random access into a WTF-8 buffer. +//! +//! WTF-8 is variable width, so a buffer's n-th code point can only be found by +//! decoding the n-1 before it: [`Wtf8`]'s iterators are sequential, and +//! resolving an index through them is O(n). Code that indexes the same string +//! repeatedly -- a regex scan restarting at successive positions, say -- then +//! walks the whole buffer once per index, which is quadratic in its length. +//! +//! [`Wtf8Index`] is the side table that makes the lookup O(1): one 24-byte +//! group per 64 code points, so 0.375 bytes per code point. It is a cache, and +//! holds no state of its own beyond the buffer's shape -- building it twice for +//! the same buffer yields the same table. +//! +//! The layout is PyPy's `UTF8_INDEX_STORAGE` (`rpython/rlib/rutf8.py`). + +use crate::wtf8::Wtf8; + +/// One group of 64 code points. +#[derive(Clone, Copy)] +struct Group { + /// The byte offset the group's first code point starts at. + base: usize, + /// `ofs[i]` is the byte offset of the group's `4 * i + 1`-th code point, + /// relative to `base`. One entry covers four code points, so the widest + /// offset an entry has to hold is that of the 61st code point of a group, + /// at most `61 * 4 = 244` bytes in -- inside a `u8`, which is what buys the + /// table its density. + ofs: [u8; 16], +} + +/// A code-point-index to byte-offset table for one WTF-8 buffer. +pub struct Wtf8Index { + groups: Box<[Group]>, +} + +impl Wtf8Index { + /// Builds the table for `data`, whose code point count is `char_len`. + /// + /// O(`data.len()`), and touches every byte, so it pays for itself only when + /// the caller goes on to index the buffer more than a couple of times. + #[must_use] + pub fn new(data: &Wtf8, char_len: usize) -> Self { + let mut groups = vec![ + Group { + base: 0, + ofs: [0; 16], + }; + char_len / 64 + 1 + ]; + // Signed: the countdown overshoots the last group -- the loop stops on + // the first negative value rather than at a group boundary. + let mut remaining = char_len as isize; + let mut base = 0; + let mut current = 0; + loop { + groups[current].base = base; + let mut next = base; + let mut group_filled = true; + for i in 0..16 { + // Past the end, step as if one more single-byte code point + // followed, so the entry stays in range and is never read. + next = if remaining == 0 { + next + 1 + } else { + next_pos(data, next) + }; + groups[current].ofs[i] = (next - base) as u8; + remaining -= 4; + if remaining < 0 { + debug_assert_eq!(current + 1, groups.len()); + group_filled = false; + break; + } + next = next_pos(data, next_pos(data, next_pos(data, next))); + } + if !group_filled { + break; + } + current += 1; + base = next; + } + Self { + groups: groups.into_boxed_slice(), + } + } + + /// The byte offset of `data`'s `index`-th code point. + /// + /// `data` must be the buffer the table was built for, and `index` must be + /// below its code point count. + #[inline] + #[must_use] + pub fn byte_offset(&self, data: &Wtf8, index: usize) -> usize { + let group = &self.groups[index >> 6]; + // The entry sits on the 4k+1-th code point of the group, so a lookup is + // one table read plus at most two steps in either direction. + let pos = group.base + group.ofs[(index >> 2) & 0x0F] as usize; + match index & 0x3 { + 0 => prev_pos(data, pos), + 1 => pos, + 2 => next_pos(data, pos), + _ => next_pos(data, next_pos(data, pos)), + } + } + + /// The table's heap footprint, in bytes. + #[must_use] + pub fn byte_size(&self) -> usize { + core::mem::size_of_val(&*self.groups) + } +} + +/// The byte offset of the code point after the one at `pos`. +/// +/// `data` must be well-formed WTF-8 and `pos` a code point boundary before its +/// end -- reading only the lead byte is what makes this branch-light. +#[inline] +fn next_pos(data: &Wtf8, pos: usize) -> usize { + match data.as_bytes()[pos] { + 0x00..=0x7F => pos + 1, + 0x80..=0xDF => pos + 2, + 0xE0..=0xEF => pos + 3, + _ => pos + 4, + } +} + +/// The byte offset of the code point before the one at `pos`, which must not be +/// zero. +/// +/// A `pos` one past the end reads as the extra code point [`Wtf8Index::new`] +/// steps over there. +#[inline] +fn prev_pos(data: &Wtf8, pos: usize) -> usize { + let data = data.as_bytes(); + let mut pos = pos - 1; + if pos >= data.len() || data[pos] <= 0x7F { + return pos; + } + pos -= 1; + if data[pos] >= 0xC0 { + return pos; + } + pos -= 1; + if data[pos] >= 0xC0 { + return pos; + } + pos - 1 +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::wtf8::{CodePoint, Wtf8Buf}; + + /// Every index of `s`, against the offsets its own iterator reports. + fn check(s: &Wtf8) { + let expected: Vec = s + .code_point_indices() + .map(|(byte_offset, _)| byte_offset) + .collect(); + let index = Wtf8Index::new(s, expected.len()); + for (i, &want) in expected.iter().enumerate() { + assert_eq!( + index.byte_offset(s, i), + want, + "index {i} of {s:?} ({} code points)", + expected.len() + ); + } + } + + fn wtf8(s: &str) -> Wtf8Buf { + Wtf8Buf::from(s) + } + + #[test] + fn empty() { + check(wtf8("").as_ref()); + } + + #[test] + fn widths() { + // One case per encoded width, and the boundaries between them. + check(wtf8("abc").as_ref()); + check(wtf8("\u{80}\u{7ff}").as_ref()); + check(wtf8("\u{800}\u{ffff}").as_ref()); + check(wtf8("\u{10000}\u{10ffff}").as_ref()); + check(wtf8("a\u{80}\u{800}\u{10000}").as_ref()); + } + + #[test] + fn group_boundaries() { + // A group covers 64 code points and an entry four, so the interesting + // lengths are the ones on and around both. + for len in [1, 3, 4, 5, 63, 64, 65, 127, 128, 129, 255, 256, 257] { + for unit in ["a", "\u{80}", "\u{800}", "\u{10000}"] { + check(wtf8(&unit.repeat(len)).as_ref()); + } + // Mixed widths, so a group's entries do not share a stride. + check(wtf8(&"a\u{80}\u{800}\u{10000}".repeat(len)).as_ref()); + } + } + + #[test] + fn lone_surrogates() { + let mut s = wtf8("a"); + for cp in [0xD800, 0xDBFF, 0xDC00, 0xDFFF] { + s.push(CodePoint::from_u32(cp).unwrap()); + s.push_str("b"); + } + check(s.as_ref()); + + // Surrogates only, spanning more than one group. + let mut s = wtf8(""); + for i in 0..200 { + s.push(CodePoint::from_u32(0xD800 + (i % 0x400)).unwrap()); + } + check(s.as_ref()); + } + + #[test] + fn byte_size_is_one_group_per_64_code_points() { + let s = wtf8(&"\u{10000}".repeat(200)); + let index = Wtf8Index::new(s.as_ref(), 200); + assert_eq!(index.byte_size(), (200 / 64 + 1) * size_of::()); + assert_eq!(size_of::(), 24); + } +} diff --git a/crates/vm/src/builtins/str.rs b/crates/vm/src/builtins/str.rs index 07325159a39..5502bcce138 100644 --- a/crates/vm/src/builtins/str.rs +++ b/crates/vm/src/builtins/str.rs @@ -712,6 +712,13 @@ impl PyStr { self.data.char_len() } + /// The byte offset the `index`-th character starts at, or the string's byte + /// length if `index` is at or past its end. + #[inline] + pub fn char_index_to_byte(&self, index: usize) -> usize { + self.data.char_index_to_byte(index) + } + #[pymethod] #[inline(always)] pub const fn isascii(&self) -> bool { From 30b2aaa9fbf68a967beb4017f25f7c7145556c8e Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Fri, 14 Aug 2026 17:51:36 +0900 Subject: [PATCH 2/2] _sre: drive a non-ASCII str subject through the string's character index The `&Wtf8` drive answers `count` and `create_cursor` by decoding from the start of the subject, so a scan that restarts at successive positions walks the subject once per position, and `slice` walks it again per extracted group. `Utf8Str` holds the `PyStr` and asks it instead: `count` is the cached character length, and `create_cursor` and `slice` resolve their positions through `char_index_to_byte`. Stepping is the `&Wtf8` drive's, unchanged. The table lives on the string, so a `Match` that outlives the scan shares it -- `group` has no cursor of its own to move relative to. `SreStr for &Wtf8` has no callers left; the `StrDrive` impl stays, since `Utf8Str` steps through it. The three subject helpers now share one downcast. Assisted-by: Claude --- crates/vm/src/stdlib/_sre.rs | 101 +++++++++++++++++++++++++++++------ 1 file changed, 85 insertions(+), 16 deletions(-) diff --git a/crates/vm/src/stdlib/_sre.rs b/crates/vm/src/stdlib/_sre.rs index 6fe6b434702..168363103e4 100644 --- a/crates/vm/src/stdlib/_sre.rs +++ b/crates/vm/src/stdlib/_sre.rs @@ -70,15 +70,76 @@ mod _sre { } } - impl SreStr for &Wtf8 { + /// A `str` subject with non-ASCII characters, driven through the string's + /// own character-index table. + /// + /// The `&Wtf8` drive answers `count` and `create_cursor` by decoding from + /// the start of the subject, so both are O(n) and a scan that restarts at + /// successive positions walks the subject once per position. `PyStr` + /// already caches its character length and can resolve a character index to + /// a byte offset in constant time, so this drive asks the string instead of + /// re-deriving: the table it builds on the first lookup is shared by every + /// later one, including by `Match` objects that outlive the scan and have + /// no cursor of their own to move relative to. + /// + /// Stepping is the `&Wtf8` drive's, unchanged -- the subject is the same + /// buffer, decoded the same way. Only the two operations that resolve a + /// position from scratch differ. + #[derive(Clone, Copy)] + struct Utf8Str<'a>(&'a Py); + + impl StrDrive for Utf8Str<'_> { + fn count(&self) -> usize { + self.0.char_len() + } + + fn create_cursor(&self, n: usize) -> StringCursor { + // `StringCursor`'s pointer is private to the engine, so the cursor + // is taken from the `&Wtf8` drive at the start of the suffix that + // begins at `n` -- an O(1) reslice -- rather than built here. + let suffix = &self.0.as_wtf8()[self.0.char_index_to_byte(n)..]; + let mut cursor = <&Wtf8 as StrDrive>::create_cursor(&suffix, 0); + cursor.position = n; + cursor + } + + fn adjust_cursor(&self, cursor: &mut StringCursor, n: usize) { + // Rebuilding is O(1), so it is never the slower branch and the + // `&Wtf8` drive's walk-or-restart choice does not apply. + *cursor = self.create_cursor(n); + } + + fn advance(cursor: &mut StringCursor) -> u32 { + <&Wtf8 as StrDrive>::advance(cursor) + } + + fn peek(cursor: &StringCursor) -> u32 { + <&Wtf8 as StrDrive>::peek(cursor) + } + + fn skip(cursor: &mut StringCursor, n: usize) { + <&Wtf8 as StrDrive>::skip(cursor, n) + } + + fn back_advance(cursor: &mut StringCursor) -> u32 { + <&Wtf8 as StrDrive>::back_advance(cursor) + } + + fn back_peek(cursor: &StringCursor) -> u32 { + <&Wtf8 as StrDrive>::back_peek(cursor) + } + + fn back_skip(cursor: &mut StringCursor, n: usize) { + <&Wtf8 as StrDrive>::back_skip(cursor, n) + } + } + + impl SreStr for Utf8Str<'_> { fn slice(&self, start: usize, end: usize, vm: &VirtualMachine) -> PyObjectRef { + let end = self.0.char_index_to_byte(end); + let start = self.0.char_index_to_byte(start).min(end); vm.ctx - .new_str( - self.code_points() - .take(end) - .skip(start) - .collect::(), - ) + .new_str(self.0.as_wtf8()[start..end].to_owned()) .into() } } @@ -275,28 +336,31 @@ mod _sre { } else if Pattern::is_ascii_str(subject) { Pattern::with_ascii_str(subject, $vm, $f) } else { - Pattern::with_str(subject, $vm, $f) + Pattern::with_utf8_str(subject, $vm, $f) } }}; } #[pyclass(with(Hashable, Comparable, Representable), flags(HAS_WEAKREF))] impl Pattern { + fn downcast_str<'a>(string: &'a PyObject, vm: &VirtualMachine) -> PyResult<&'a Py> { + string.downcast_ref::().ok_or_else(|| { + vm.new_type_error(format!("expected string got '{}'", string.class())) + }) + } + fn with_str(string: &PyObject, vm: &VirtualMachine, f: F) -> PyResult where F: FnOnce(&Wtf8) -> PyResult, { - let string = string.downcast_ref::().ok_or_else(|| { - vm.new_type_error(format!("expected string got '{}'", string.class())) - })?; - f(string.as_wtf8()) + f(Self::downcast_str(string, vm)?.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`]. + /// argument answers `false` and is reported by [`Self::with_utf8_str`]. fn is_ascii_str(string: &PyObject) -> bool { string .downcast_ref::() @@ -307,12 +371,17 @@ mod _sre { where F: FnOnce(AsciiStr<'_>) -> PyResult, { - let string = string.downcast_ref::().ok_or_else(|| { - vm.new_type_error(format!("expected string got '{}'", string.class())) - })?; + let string = Self::downcast_str(string, vm)?; f(AsciiStr(string.as_wtf8().as_bytes())) } + fn with_utf8_str(string: &PyObject, vm: &VirtualMachine, f: F) -> PyResult + where + F: FnOnce(Utf8Str<'_>) -> PyResult, + { + f(Utf8Str(Self::downcast_str(string, vm)?)) + } + fn with_bytes(string: &PyObject, vm: &VirtualMachine, f: F) -> PyResult where F: FnOnce(&[u8]) -> PyResult,