mirrors/RustPython
0
0 Fork 0
mirror of https://github.com/RustPython/RustPython.git synced 2026-10-28 18:59:04 +08:00
Code Activity

str: resolve subscripts and slices through the code point index (#8526)

Browse Source 
`nth_char` and the four `SliceableSequenceOp` methods each walked the buffer
to reach a character index, so `s[i]` and `s[a:b]` on a non-ASCII string were
O(i) and O(b), and a loop over either was quadratic. They now resolve through
`char_index_to_byte`, which is what the index table was added for: a plain
slice becomes a byte reslice, and a stepped slice one lookup per collected
character.

An index within four code points of an end is still walked to, and so is a
slice that reaches within four of both. PyPy draws the same line with
`MAX_UNROLL_NEXT_CODEPOINT_POS`, in a guard that also asks the JIT whether the
index is a constant; there is no JIT here, but the reason to skip the build
survives it -- `s[0]` on a long string should not pay for a table.

The stepped slices took their character count from `(range.len() / step) + 1`,
which overshoots whenever the last step lands short: `"aéc"[::3]` reported a
length of 2 for a one-character string, and `reversed()` on it read past the
end of the buffer and panicked. The count is now the index iterator's own
length, so it cannot drift from the characters actually collected.

Assisted-by: Claude
This commit is contained in:
Jeong, YunWon authored and GitHub committed 2026-08-15 01:03:09 +09:00
1 parent d6091086e0
commit 7e25617f76
2 files changed
+128 -99
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/common/src/str.rs
+75 -2
View File
@@ -113,6 +113,16 @@ pub enum PyKindStr<'a> {
Wtf8(&'a Wtf8),
}
/// How far from an end an index is resolved by walking rather than by building
/// the code point index.
///
/// PyPy spells this `MAX_UNROLL_NEXT_CODEPOINT_POS`, in a guard that also asks
/// the JIT whether the index is a constant, so that the walk unrolls. There is
/// no JIT here to ask, and the walk is short rather than free -- but four steps
/// still beat a pass over the whole buffer, and skipping the build is what
/// keeps `s[0]` and `s[1:-1]` on a long string from paying for a table.
const MAX_WALK_TO_INDEX: usize = 4;
#[derive(Debug, Clone)]
pub struct StrData {
data: Box<Wtf8>,
@@ -405,11 +415,74 @@ impl StrData {
.byte_offset(&self.data, index)
}
/// The byte offset of code point `index`, for a caller that resolves one
/// index and stops.
///
/// Building the table costs a pass over the whole buffer, so it is worth it
/// only for a caller that comes back; an index within
/// [`MAX_WALK_TO_INDEX`] steps of either end is cheaper to walk to, and
/// walking keeps `s[0]` on a long string from paying for a table it will
/// never use again. Anything further in builds, on the reasoning that a
/// string indexed once in the middle tends to be indexed again.
fn char_index_to_byte_once(&self, index: usize) -> usize {
if index <= MAX_WALK_TO_INDEX {
return self
.data
.code_point_indices()
.nth(index)
.map_or(self.data.len(), |(byte, _)| byte);
}
let from_end = self.char_len() - index;
if from_end <= MAX_WALK_TO_INDEX {
return self
.data
.code_point_indices()
.nth_back(from_end - 1)
.map_or(self.data.len(), |(byte, _)| byte);
}
self.char_index_to_byte(index)
}
/// The byte range spanned by the code points in `range`.
///
/// A range that reaches within [`MAX_WALK_TO_INDEX`] of *both* ends is
/// walked to for the same reason a single index near one end is -- a slice
/// like `s[1:-1]` should not build a table over the whole string.
#[must_use]
pub fn char_range_to_bytes(&self, range: core::ops::Range<usize>) -> core::ops::Range<usize> {
if self.kind.is_ascii() {
return range;
}
let from_end = self.char_len() - range.end;
if range.start <= MAX_WALK_TO_INDEX && from_end <= MAX_WALK_TO_INDEX {
// Two walks over disjoint ends, each of at most MAX_WALK_TO_INDEX
// steps -- one iterator driven from both sides would have them meet
// on a short string.
let start = self
.data
.code_point_indices()
.nth(range.start)
.map_or(self.data.len(), |(byte, _)| byte);
let end = match from_end {
0 => self.data.len(),
n => self
.data
.code_point_indices()
.nth_back(n - 1)
.map_or(self.data.len(), |(byte, _)| byte),
};
return start..end;
}
self.char_index_to_byte(range.start)..self.char_index_to_byte(range.end)
}
pub fn nth_char(&self, index: usize) -> CodePoint {
match self.as_str_kind() {
PyKindStr::Ascii(s) => s[index].into(),
PyKindStr::Utf8(s) => s.chars().nth(index).unwrap().into(),
PyKindStr::Wtf8(w) => w.code_points().nth(index).unwrap(),
_ => self.data[self.char_index_to_byte_once(index)..]
.code_points()
.next()
.unwrap(),
}
}
}
crates/vm/src/builtins/str.rs
+53 -97
View File
@@ -1807,6 +1807,31 @@ pub(crate) fn init(ctx: &'static Context) {
PyStrIterator::extend_class(ctx, ctx.types.str_iterator_type);
}
impl PyStr {
/// The code points at `indices`, in that order, as a new string.
///
/// Each index is resolved through the string's own index table, so the
/// cost is one lookup per collected character rather than a walk to the
/// furthest one. The iterator's length is the result's character count,
/// which is why it has to be exact.
fn gather_chars(&self, indices: impl ExactSizeIterator<Item = usize>) -> Self {
let char_len = indices.len();
// Not ascii, so the code points are at least two bytes each.
let mut out = Wtf8Buf::with_capacity(2 * char_len);
let s = self.as_wtf8();
for index in indices {
out.push(
s[self.data.char_index_to_byte(index)..]
.code_points()
.next()
.expect("index is below the character count"),
);
}
// SAFETY: char_len is accurate
unsafe { Self::new_with_char_len(out, char_len) }
}
}
impl SliceableSequenceOp for PyStr {
type Item = CodePoint;
type Sliced = Self;
@@ -1816,125 +1841,56 @@ impl SliceableSequenceOp for PyStr {
}
fn do_slice(&self, range: Range<usize>) -> Self::Sliced {
match self.as_str_kind() {
PyKindStr::Ascii(s) => s[range].into(),
PyKindStr::Utf8(s) => {
let char_len = range.len();
let out = rustpython_common::str::get_chars(s, range);
// SAFETY: char_len is accurate
unsafe { Self::new_with_char_len(out, char_len) }
}
PyKindStr::Wtf8(w) => {
let char_len = range.len();
let out = rustpython_common::str::get_codepoints(w, range);
// SAFETY: char_len is accurate
unsafe { Self::new_with_char_len(out, char_len) }
}
if let PyKindStr::Ascii(s) = self.as_str_kind() {
return s[range].into();
}
// Both ends resolve through the string's own index, so the slice is a
// byte reslice rather than a walk to `range.start` and another to
// `range.end`.
let char_len = range.len();
let bytes = self.data.char_range_to_bytes(range);
let out = &self.as_wtf8()[bytes];
// SAFETY: char_len is accurate
unsafe { Self::new_with_char_len(out.to_owned(), char_len) }
}
fn do_slice_reverse(&self, range: Range<usize>) -> Self::Sliced {
match self.as_str_kind() {
PyKindStr::Ascii(s) => {
let mut out = s[range].to_owned();
out.as_mut_slice().reverse();
out.into()
}
PyKindStr::Utf8(s) => {
let char_len = range.len();
let mut out = String::with_capacity(2 * char_len);
out.extend(
s.chars()
.rev()
.skip(self.char_len() - range.end)
.take(range.len()),
);
// SAFETY: char_len is accurate
unsafe { Self::new_with_char_len(out, range.len()) }
}
PyKindStr::Wtf8(w) => {
let char_len = range.len();
let mut out = Wtf8Buf::with_capacity(2 * char_len);
out.extend(
w.code_points()
.rev()
.skip(self.char_len() - range.end)
.take(range.len()),
);
// SAFETY: char_len is accurate
unsafe { Self::new_with_char_len(out, char_len) }
}
if let PyKindStr::Ascii(s) = self.as_str_kind() {
let mut out = s[range].to_owned();
out.as_mut_slice().reverse();
return out.into();
}
let char_len = range.len();
let bytes = self.data.char_range_to_bytes(range);
let mut out = Wtf8Buf::with_capacity(bytes.len());
out.extend(self.as_wtf8()[bytes].code_points().rev());
// SAFETY: char_len is accurate
unsafe { Self::new_with_char_len(out, char_len) }
}
fn do_stepped_slice(&self, range: Range<usize>, step: usize) -> Self::Sliced {
match self.as_str_kind() {
PyKindStr::Ascii(s) => s[range]
if let PyKindStr::Ascii(s) = self.as_str_kind() {
return s[range]
.as_slice()
.iter()
.copied()
.step_by(step)
.collect::<AsciiString>()
.into(),
PyKindStr::Utf8(s) => {
let char_len = range.len().div_ceil(step);
let mut out = String::with_capacity(2 * char_len);
out.extend(s.chars().skip(range.start).take(range.len()).step_by(step));
// SAFETY: char_len is accurate
unsafe { Self::new_with_char_len(out, char_len) }
}
PyKindStr::Wtf8(w) => {
let char_len = range.len().div_ceil(step);
let mut out = Wtf8Buf::with_capacity(2 * char_len);
out.extend(
w.code_points()
.skip(range.start)
.take(range.len())
.step_by(step),
);
// SAFETY: char_len is accurate
unsafe { Self::new_with_char_len(out, char_len) }
}
.into();
}
self.gather_chars(range.step_by(step))
}
fn do_stepped_slice_reverse(&self, range: Range<usize>, step: usize) -> Self::Sliced {
match self.as_str_kind() {
PyKindStr::Ascii(s) => s[range]
if let PyKindStr::Ascii(s) = self.as_str_kind() {
return s[range]
.chars()
.rev()
.step_by(step)
.collect::<AsciiString>()
.into(),
PyKindStr::Utf8(s) => {
let char_len = range.len().div_ceil(step);
// not ascii, so the codepoints have to be at least 2 bytes each
let mut out = String::with_capacity(2 * char_len);
out.extend(
s.chars()
.rev()
.skip(self.char_len() - range.end)
.take(range.len())
.step_by(step),
);
// SAFETY: char_len is accurate
unsafe { Self::new_with_char_len(out, char_len) }
}
PyKindStr::Wtf8(w) => {
let char_len = range.len().div_ceil(step);
// not ascii, so the codepoints have to be at least 2 bytes each
let mut out = Wtf8Buf::with_capacity(2 * char_len);
out.extend(
w.code_points()
.rev()
.skip(self.char_len() - range.end)
.take(range.len())
.step_by(step),
);
// SAFETY: char_len is accurate
unsafe { Self::new_with_char_len(out, char_len) }
}
.into();
}
self.gather_chars(range.rev().step_by(step))
}
fn empty() -> Self::Sliced {