diff --git a/crates/common/src/str.rs b/crates/common/src/str.rs index 8edb48fa4fa..72aecb7931e 100644 --- a/crates/common/src/str.rs +++ b/crates/common/src/str.rs @@ -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, @@ -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) -> core::ops::Range { + 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(), } } } diff --git a/crates/vm/src/builtins/str.rs b/crates/vm/src/builtins/str.rs index c25940cd3c6..06e36738603 100644 --- a/crates/vm/src/builtins/str.rs +++ b/crates/vm/src/builtins/str.rs @@ -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) -> 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) -> 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) -> 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, 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::() - .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, 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::() - .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 {