diff --git a/crates/common/src/str.rs b/crates/common/src/str.rs index 72aecb7931e..326f76c43cb 100644 --- a/crates/common/src/str.rs +++ b/crates/common/src/str.rs @@ -443,7 +443,8 @@ impl StrData { self.char_index_to_byte(index) } - /// The byte range spanned by the code points in `range`. + /// The byte range spanned by the code points in `range`, whose end must not + /// exceed the string's code point count. /// /// 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 @@ -476,6 +477,25 @@ impl StrData { self.char_index_to_byte(range.start)..self.char_index_to_byte(range.end) } + /// The character index of the character starting at byte offset `bytepos`, + /// the inverse of [`Self::char_index_to_byte`]. + /// + /// `bytepos` must be a character boundary at or before the end. + /// + /// Logarithmic rather than constant, because the index is keyed the other + /// way -- but a search whose bounds came from `char_index_to_byte` has the + /// table already, and this is what turns a byte offset back into the answer + /// a caller asked for in characters. + pub fn byte_to_char_index(&self, bytepos: usize) -> usize { + if self.kind.is_ascii() { + return bytepos; + } + let char_len = self.char_len(); + self.index + .get_or_build(&self.data, char_len) + .char_index_at_byte(&self.data, bytepos, char_len) + } + 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 index a6121f1c03d..4b2e5ebee5b 100644 --- a/crates/common/src/wtf8_index.rs +++ b/crates/common/src/wtf8_index.rs @@ -104,6 +104,63 @@ impl Wtf8Index { } } + /// The index of the code point starting at byte offset `bytepos`, the + /// inverse of [`Self::byte_offset`]. + /// + /// `data` must be the buffer the table was built for, `char_len` its code + /// point count, and `bytepos` a code point boundary at or before its end. + /// + /// Logarithmic rather than constant: the table is keyed by code point + /// index, so going the other way is a search through it. The bracketing + /// below is what keeps that search short -- a code point occupies one to + /// four bytes, which pins the answer to a narrow band around `bytepos` + /// before the first comparison. + #[must_use] + pub fn char_index_at_byte(&self, data: &Wtf8, bytepos: usize, char_len: usize) -> usize { + let bytes_remaining = data.len() - bytepos; + // At least one byte per remaining code point, and at most four, so the + // group holding the answer lies between these. + let mut group_min = + usize::max(bytepos / 4, char_len.saturating_sub(bytes_remaining + 1)) >> 6; + let mut group_max = usize::min(bytepos, char_len.saturating_sub(bytes_remaining / 4)) >> 6; + while group_min < group_max { + let middle = group_min.midpoint(group_max) + 1; + if bytepos < self.groups[middle].base { + group_max = middle - 1; + } else { + group_min = middle; + } + } + + let base = self.groups[group_min].base; + if base == bytepos { + return group_min << 6; + } + // Walk the group's entries to the last one at or before `bytepos`, + // then step the remaining code points, of which there are at most + // three -- an entry covers four. + let entries = if group_min == self.groups.len() - 1 { + ((char_len - 1) >> 2) & 0x0F + } else { + 16 + }; + let mut index = group_min << 6; + let mut pos = base; + for entry in 0..entries { + let at = base + self.groups[group_min].ofs[entry] as usize; + if at >= bytepos { + break; + } + pos = at; + index = (group_min << 6) + (entry << 2) + 1; + } + while pos < bytepos { + pos = next_pos(data, pos); + index += 1; + } + index + } + /// The table's heap footprint, in bytes. #[must_use] pub fn byte_size(&self) -> usize { @@ -153,21 +210,34 @@ mod tests { use super::*; use crate::wtf8::{CodePoint, Wtf8Buf}; - /// Every index of `s`, against the offsets its own iterator reports. + /// Every index of `s`, both ways, 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()); + let char_len = expected.len(); + let index = Wtf8Index::new(s, char_len); for (i, &want) in expected.iter().enumerate() { assert_eq!( index.byte_offset(s, i), want, - "index {i} of {s:?} ({} code points)", - expected.len() + "index {i} of {s:?} ({char_len} code points)" + ); + assert_eq!( + index.char_index_at_byte(s, want, char_len), + i, + "byte {want} of {s:?} ({char_len} code points)" ); } + // One past the last code point is a boundary too, and the searches that + // use this ask for it as an end bound. + assert_eq!( + index.char_index_at_byte(s, s.len(), char_len), + char_len, + "end of {s:?}" + ); } fn wtf8(s: &str) -> Wtf8Buf { diff --git a/crates/vm/src/anystr.rs b/crates/vm/src/anystr.rs index 69ba525267a..45a69dbfe2f 100644 --- a/crates/vm/src/anystr.rs +++ b/crates/vm/src/anystr.rs @@ -147,7 +147,11 @@ pub(crate) trait AnyStr { fn as_bytes(&self) -> &[u8]; fn elements(&self) -> impl Iterator; fn get_bytes(&self, range: Range) -> &Self; - // FIXME: get_chars is expensive for str + /// The characters in `range`, which for a `str` payload means walking to + /// both bounds -- the payload does not carry the string's character index. + /// `PyStr` therefore converts its own ranges and does not reach the search + /// helpers below through this; what remains are the byte strings, where a + /// character range is already a byte range. fn get_chars(&self, range: Range) -> &Self; fn bytes_len(&self) -> usize; // NOTE: str::chars().count() consumes the O(n) time. But pystr::char_len does cache. diff --git a/crates/vm/src/builtins/str.rs b/crates/vm/src/builtins/str.rs index 06e36738603..2e099424a23 100644 --- a/crates/vm/src/builtins/str.rs +++ b/crates/vm/src/builtins/str.rs @@ -9,7 +9,7 @@ use super::{ use crate::{ AsObject, Context, Py, PyExact, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult, TryFromBorrowedObject, VirtualMachine, - anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper, adjust_indices}, + anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper, StringRange, adjust_indices}, atomic_func, bytes_inner::{swapcase_ascii, title_ascii}, cformat::cformat_string, @@ -719,6 +719,13 @@ impl PyStr { self.data.char_index_to_byte(index) } + /// The character index of the character starting at byte offset `bytepos`, + /// which must be a character boundary at or before the end. + #[inline] + pub fn byte_to_char_index(&self, bytepos: usize) -> usize { + self.data.byte_to_char_index(bytepos) + } + #[pymethod] #[inline(always)] pub const fn isascii(&self) -> bool { @@ -924,11 +931,12 @@ impl PyStr { #[pymethod] fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult { - let (affix, substr) = - match options.prepare(self.as_wtf8(), self.len(), |s, r| s.get_chars(r)) { - Some(x) => x, - None => return Ok(false), - }; + let (affix, substr) = match options.prepare(self.as_wtf8(), self.len(), |s, r| { + &s[self.data.char_range_to_bytes(r)] + }) { + Some(x) => x, + None => return Ok(false), + }; substr.py_starts_ends_with( &affix, "endswith", @@ -944,11 +952,12 @@ impl PyStr { options: anystr::StartsEndsWithArgs, vm: &VirtualMachine, ) -> PyResult { - let (affix, substr) = - match options.prepare(self.as_wtf8(), self.len(), |s, r| s.get_chars(r)) { - Some(x) => x, - None => return Ok(false), - }; + let (affix, substr) = match options.prepare(self.as_wtf8(), self.len(), |s, r| { + &s[self.data.char_range_to_bytes(r)] + }) { + Some(x) => x, + None => return Ok(false), + }; substr.py_starts_ends_with( &affix, "startswith", @@ -1167,42 +1176,52 @@ impl PyStr { Ok(vm.ctx.new_str(joined)) } - // FIXME: two traversals of str is expensive + /// The bytes the character range `range` spans and the byte offset it + /// starts at, or `None` if the range is inverted. + /// + /// The bounds go through the string's character index, so reaching a range + /// deep in the subject costs a lookup rather than a walk to it. #[inline] - fn _to_char_idx(r: &Wtf8, byte_idx: usize) -> usize { - r[..byte_idx].code_points().count() + fn char_range_bytes(&self, range: Range) -> Option<(usize, &Wtf8)> { + if !range.is_normal() { + return None; + } + let bytes = self.data.char_range_to_bytes(range); + Some((bytes.start, &self.as_wtf8()[bytes])) } + /// Searches the character range `range` with `find`, which answers in bytes + /// relative to the range, and reports the hit as a character index. #[inline] fn _find(&self, args: FindArgs, find: F) -> Option where F: Fn(&Wtf8, &Wtf8) -> Option, { let (sub, range) = args.get_value(self.len()); - self.as_wtf8().py_find(sub.as_wtf8(), range, find) + let (start, haystack) = self.char_range_bytes(range)?; + let found = find(haystack, sub.as_wtf8())?; + Some(self.byte_to_char_index(start + found)) } #[pymethod] fn find(&self, args: FindArgs) -> isize { - self._find(args, |r, s| Some(Self::_to_char_idx(r, r.find(s)?))) - .map_or(-1, |v| v as isize) + self._find(args, Wtf8::find).map_or(-1, |v| v as isize) } #[pymethod] fn rfind(&self, args: FindArgs) -> isize { - self._find(args, |r, s| Some(Self::_to_char_idx(r, r.rfind(s)?))) - .map_or(-1, |v| v as isize) + self._find(args, Wtf8::rfind).map_or(-1, |v| v as isize) } #[pymethod] fn index(&self, args: FindArgs, vm: &VirtualMachine) -> PyResult { - self._find(args, |r, s| Some(Self::_to_char_idx(r, r.find(s)?))) + self._find(args, Wtf8::find) .ok_or_else(|| vm.new_value_error("substring not found")) } #[pymethod] fn rindex(&self, args: FindArgs, vm: &VirtualMachine) -> PyResult { - self._find(args, |r, s| Some(Self::_to_char_idx(r, r.rfind(s)?))) + self._find(args, Wtf8::rfind) .ok_or_else(|| vm.new_value_error("substring not found")) } @@ -1275,8 +1294,18 @@ impl PyStr { #[pymethod] fn count(&self, args: FindArgs) -> usize { let (needle, range) = args.get_value(self.len()); - self.as_wtf8() - .py_count(needle.as_wtf8(), range, |h, n| h.find_iter(n).count()) + let chars = range.len(); + self.char_range_bytes(range).map_or(0, |(_, haystack)| { + if needle.is_empty() { + // An empty needle sits between every pair of characters and at + // both ends, so it occurs once more than the range holds + // characters. Counting it in the bytes would answer in encoded + // positions instead. + chars + 1 + } else { + haystack.find_iter(needle.as_wtf8()).count() + } + }) } #[pymethod] diff --git a/extra_tests/snippets/builtin_str.py b/extra_tests/snippets/builtin_str.py index fde9deb8e0b..6eead5ddbfb 100644 --- a/extra_tests/snippets/builtin_str.py +++ b/extra_tests/snippets/builtin_str.py @@ -170,6 +170,15 @@ assert "aaa".count("a", 2, 2) == 0 assert "aaa".count("a", 2, 1) == 0 +# An empty needle is counted in characters, not in encoded positions. +assert "".count("") == 1 +assert "abc".count("") == 4 +assert "가나다".count("") == 4 +assert "가나다".count("", 1) == 3 +assert "가나다".count("", 1, 2) == 2 +assert "가나다".count("", 4, 4) == 0 +assert "a\U0001f600b".count("") == 4 + assert "___a__".find("a") == 3 assert "___a__".find("a", -10) == 3 assert "___a__".find("a", -3) == 3