Skip to content

Commit 7e25617

Browse files
youknowoneyouknowone
authored
str: resolve subscripts and slices through the code point index (#8526)
`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
1 parent d609108 commit 7e25617

2 files changed

Lines changed: 128 additions & 99 deletions

File tree

‎crates/common/src/str.rs‎

Lines changed: 75 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,16 @@ pub enum PyKindStr<'a> {
113113
Wtf8(&'a Wtf8),
114114
}
115115

116+
/// How far from an end an index is resolved by walking rather than by building
117+
/// the code point index.
118+
///
119+
/// PyPy spells this `MAX_UNROLL_NEXT_CODEPOINT_POS`, in a guard that also asks
120+
/// the JIT whether the index is a constant, so that the walk unrolls. There is
121+
/// no JIT here to ask, and the walk is short rather than free -- but four steps
122+
/// still beat a pass over the whole buffer, and skipping the build is what
123+
/// keeps `s[0]` and `s[1:-1]` on a long string from paying for a table.
124+
const MAX_WALK_TO_INDEX: usize = 4;
125+
116126
#[derive(Debug, Clone)]
117127
pub struct StrData {
118128
data: Box<Wtf8>,
@@ -405,11 +415,74 @@ impl StrData {
405415
.byte_offset(&self.data, index)
406416
}
407417

418+
/// The byte offset of code point `index`, for a caller that resolves one
419+
/// index and stops.
420+
///
421+
/// Building the table costs a pass over the whole buffer, so it is worth it
422+
/// only for a caller that comes back; an index within
423+
/// [`MAX_WALK_TO_INDEX`] steps of either end is cheaper to walk to, and
424+
/// walking keeps `s[0]` on a long string from paying for a table it will
425+
/// never use again. Anything further in builds, on the reasoning that a
426+
/// string indexed once in the middle tends to be indexed again.
427+
fn char_index_to_byte_once(&self, index: usize) -> usize {
428+
if index <= MAX_WALK_TO_INDEX {
429+
return self
430+
.data
431+
.code_point_indices()
432+
.nth(index)
433+
.map_or(self.data.len(), |(byte, _)| byte);
434+
}
435+
let from_end = self.char_len() - index;
436+
if from_end <= MAX_WALK_TO_INDEX {
437+
return self
438+
.data
439+
.code_point_indices()
440+
.nth_back(from_end - 1)
441+
.map_or(self.data.len(), |(byte, _)| byte);
442+
}
443+
self.char_index_to_byte(index)
444+
}
445+
446+
/// The byte range spanned by the code points in `range`.
447+
///
448+
/// A range that reaches within [`MAX_WALK_TO_INDEX`] of *both* ends is
449+
/// walked to for the same reason a single index near one end is -- a slice
450+
/// like `s[1:-1]` should not build a table over the whole string.
451+
#[must_use]
452+
pub fn char_range_to_bytes(&self, range: core::ops::Range<usize>) -> core::ops::Range<usize> {
453+
if self.kind.is_ascii() {
454+
return range;
455+
}
456+
let from_end = self.char_len() - range.end;
457+
if range.start <= MAX_WALK_TO_INDEX && from_end <= MAX_WALK_TO_INDEX {
458+
// Two walks over disjoint ends, each of at most MAX_WALK_TO_INDEX
459+
// steps -- one iterator driven from both sides would have them meet
460+
// on a short string.
461+
let start = self
462+
.data
463+
.code_point_indices()
464+
.nth(range.start)
465+
.map_or(self.data.len(), |(byte, _)| byte);
466+
let end = match from_end {
467+
0 => self.data.len(),
468+
n => self
469+
.data
470+
.code_point_indices()
471+
.nth_back(n - 1)
472+
.map_or(self.data.len(), |(byte, _)| byte),
473+
};
474+
return start..end;
475+
}
476+
self.char_index_to_byte(range.start)..self.char_index_to_byte(range.end)
477+
}
478+
408479
pub fn nth_char(&self, index: usize) -> CodePoint {
409480
match self.as_str_kind() {
410481
PyKindStr::Ascii(s) => s[index].into(),
411-
PyKindStr::Utf8(s) => s.chars().nth(index).unwrap().into(),
412-
PyKindStr::Wtf8(w) => w.code_points().nth(index).unwrap(),
482+
_ => self.data[self.char_index_to_byte_once(index)..]
483+
.code_points()
484+
.next()
485+
.unwrap(),
413486
}
414487
}
415488
}

‎crates/vm/src/builtins/str.rs‎

Lines changed: 53 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -1807,6 +1807,31 @@ pub(crate) fn init(ctx: &'static Context) {
18071807
PyStrIterator::extend_class(ctx, ctx.types.str_iterator_type);
18081808
}
18091809

1810+
impl PyStr {
1811+
/// The code points at `indices`, in that order, as a new string.
1812+
///
1813+
/// Each index is resolved through the string's own index table, so the
1814+
/// cost is one lookup per collected character rather than a walk to the
1815+
/// furthest one. The iterator's length is the result's character count,
1816+
/// which is why it has to be exact.
1817+
fn gather_chars(&self, indices: impl ExactSizeIterator<Item = usize>) -> Self {
1818+
let char_len = indices.len();
1819+
// Not ascii, so the code points are at least two bytes each.
1820+
let mut out = Wtf8Buf::with_capacity(2 * char_len);
1821+
let s = self.as_wtf8();
1822+
for index in indices {
1823+
out.push(
1824+
s[self.data.char_index_to_byte(index)..]
1825+
.code_points()
1826+
.next()
1827+
.expect("index is below the character count"),
1828+
);
1829+
}
1830+
// SAFETY: char_len is accurate
1831+
unsafe { Self::new_with_char_len(out, char_len) }
1832+
}
1833+
}
1834+
18101835
impl SliceableSequenceOp for PyStr {
18111836
type Item = CodePoint;
18121837
type Sliced = Self;
@@ -1816,125 +1841,56 @@ impl SliceableSequenceOp for PyStr {
18161841
}
18171842

18181843
fn do_slice(&self, range: Range<usize>) -> Self::Sliced {
1819-
match self.as_str_kind() {
1820-
PyKindStr::Ascii(s) => s[range].into(),
1821-
PyKindStr::Utf8(s) => {
1822-
let char_len = range.len();
1823-
let out = rustpython_common::str::get_chars(s, range);
1824-
// SAFETY: char_len is accurate
1825-
unsafe { Self::new_with_char_len(out, char_len) }
1826-
}
1827-
PyKindStr::Wtf8(w) => {
1828-
let char_len = range.len();
1829-
let out = rustpython_common::str::get_codepoints(w, range);
1830-
// SAFETY: char_len is accurate
1831-
unsafe { Self::new_with_char_len(out, char_len) }
1832-
}
1844+
if let PyKindStr::Ascii(s) = self.as_str_kind() {
1845+
return s[range].into();
18331846
}
1847+
// Both ends resolve through the string's own index, so the slice is a
1848+
// byte reslice rather than a walk to `range.start` and another to
1849+
// `range.end`.
1850+
let char_len = range.len();
1851+
let bytes = self.data.char_range_to_bytes(range);
1852+
let out = &self.as_wtf8()[bytes];
1853+
// SAFETY: char_len is accurate
1854+
unsafe { Self::new_with_char_len(out.to_owned(), char_len) }
18341855
}
18351856

18361857
fn do_slice_reverse(&self, range: Range<usize>) -> Self::Sliced {
1837-
match self.as_str_kind() {
1838-
PyKindStr::Ascii(s) => {
1839-
let mut out = s[range].to_owned();
1840-
out.as_mut_slice().reverse();
1841-
out.into()
1842-
}
1843-
PyKindStr::Utf8(s) => {
1844-
let char_len = range.len();
1845-
let mut out = String::with_capacity(2 * char_len);
1846-
out.extend(
1847-
s.chars()
1848-
.rev()
1849-
.skip(self.char_len() - range.end)
1850-
.take(range.len()),
1851-
);
1852-
// SAFETY: char_len is accurate
1853-
unsafe { Self::new_with_char_len(out, range.len()) }
1854-
}
1855-
PyKindStr::Wtf8(w) => {
1856-
let char_len = range.len();
1857-
let mut out = Wtf8Buf::with_capacity(2 * char_len);
1858-
out.extend(
1859-
w.code_points()
1860-
.rev()
1861-
.skip(self.char_len() - range.end)
1862-
.take(range.len()),
1863-
);
1864-
// SAFETY: char_len is accurate
1865-
unsafe { Self::new_with_char_len(out, char_len) }
1866-
}
1858+
if let PyKindStr::Ascii(s) = self.as_str_kind() {
1859+
let mut out = s[range].to_owned();
1860+
out.as_mut_slice().reverse();
1861+
return out.into();
18671862
}
1863+
let char_len = range.len();
1864+
let bytes = self.data.char_range_to_bytes(range);
1865+
let mut out = Wtf8Buf::with_capacity(bytes.len());
1866+
out.extend(self.as_wtf8()[bytes].code_points().rev());
1867+
// SAFETY: char_len is accurate
1868+
unsafe { Self::new_with_char_len(out, char_len) }
18681869
}
18691870

18701871
fn do_stepped_slice(&self, range: Range<usize>, step: usize) -> Self::Sliced {
1871-
match self.as_str_kind() {
1872-
PyKindStr::Ascii(s) => s[range]
1872+
if let PyKindStr::Ascii(s) = self.as_str_kind() {
1873+
return s[range]
18731874
.as_slice()
18741875
.iter()
18751876
.copied()
18761877
.step_by(step)
18771878
.collect::<AsciiString>()
1878-
.into(),
1879-
PyKindStr::Utf8(s) => {
1880-
let char_len = range.len().div_ceil(step);
1881-
let mut out = String::with_capacity(2 * char_len);
1882-
out.extend(s.chars().skip(range.start).take(range.len()).step_by(step));
1883-
// SAFETY: char_len is accurate
1884-
unsafe { Self::new_with_char_len(out, char_len) }
1885-
}
1886-
PyKindStr::Wtf8(w) => {
1887-
let char_len = range.len().div_ceil(step);
1888-
let mut out = Wtf8Buf::with_capacity(2 * char_len);
1889-
out.extend(
1890-
w.code_points()
1891-
.skip(range.start)
1892-
.take(range.len())
1893-
.step_by(step),
1894-
);
1895-
// SAFETY: char_len is accurate
1896-
unsafe { Self::new_with_char_len(out, char_len) }
1897-
}
1879+
.into();
18981880
}
1881+
self.gather_chars(range.step_by(step))
18991882
}
19001883

19011884
fn do_stepped_slice_reverse(&self, range: Range<usize>, step: usize) -> Self::Sliced {
1902-
match self.as_str_kind() {
1903-
PyKindStr::Ascii(s) => s[range]
1885+
if let PyKindStr::Ascii(s) = self.as_str_kind() {
1886+
return s[range]
19041887
.chars()
19051888
.rev()
19061889
.step_by(step)
19071890
.collect::<AsciiString>()
1908-
.into(),
1909-
PyKindStr::Utf8(s) => {
1910-
let char_len = range.len().div_ceil(step);
1911-
// not ascii, so the codepoints have to be at least 2 bytes each
1912-
let mut out = String::with_capacity(2 * char_len);
1913-
out.extend(
1914-
s.chars()
1915-
.rev()
1916-
.skip(self.char_len() - range.end)
1917-
.take(range.len())
1918-
.step_by(step),
1919-
);
1920-
// SAFETY: char_len is accurate
1921-
unsafe { Self::new_with_char_len(out, char_len) }
1922-
}
1923-
PyKindStr::Wtf8(w) => {
1924-
let char_len = range.len().div_ceil(step);
1925-
// not ascii, so the codepoints have to be at least 2 bytes each
1926-
let mut out = Wtf8Buf::with_capacity(2 * char_len);
1927-
out.extend(
1928-
w.code_points()
1929-
.rev()
1930-
.skip(self.char_len() - range.end)
1931-
.take(range.len())
1932-
.step_by(step),
1933-
);
1934-
// SAFETY: char_len is accurate
1935-
unsafe { Self::new_with_char_len(out, char_len) }
1936-
}
1891+
.into();
19371892
}
1893+
self.gather_chars(range.rev().step_by(step))
19381894
}
19391895

19401896
fn empty() -> Self::Sliced {

0 commit comments

Comments
 (0)