Skip to content

Commit 3511d1a

Browse files
youknowoneyouknowone
committed
str: take search bounds through the character index
find, rfind, index, rindex, count, startswith and endswith resolved their start/stop bounds with AnyStr::get_chars, which walks the payload to both bounds, and find reported its hit by counting the code points in front of it -- so a sweep over a subject's own indices walked it once per call. Slice the bounds with PyStr::char_index_to_byte instead, and map the hit back with Wtf8Index::char_index_at_byte, added here as the table's inverse: a bracketed search over the groups, then at most an entry's worth of steps. The payload alone does not say whether it is ASCII, so get_chars walked there too; that path is now off it as well. n=8000, sweeping the bound over the subject, best of 3, interleaved: '가나다라'*n/4 'abcd'*n/4 find 17.72ms -> 0.16ms 3.48ms -> 0.11ms startswith 17.60ms -> 0.09ms count 22.54ms -> 3.72ms count stays linear in the range it is given, as it is in CPython. Assisted-by: Claude
1 parent e9b6510 commit 3511d1a

4 files changed

Lines changed: 146 additions & 33 deletions

File tree

‎crates/common/src/str.rs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -405,6 +405,25 @@ impl StrData {
405405
.byte_offset(&self.data, index)
406406
}
407407

408+
/// The character index of the character starting at byte offset `bytepos`,
409+
/// the inverse of [`Self::char_index_to_byte`].
410+
///
411+
/// `bytepos` must be a character boundary at or before the end.
412+
///
413+
/// Logarithmic rather than constant, because the index is keyed the other
414+
/// way -- but a search whose bounds came from `char_index_to_byte` has the
415+
/// table already, and this is what turns a byte offset back into the answer
416+
/// a caller asked for in characters.
417+
pub fn byte_to_char_index(&self, bytepos: usize) -> usize {
418+
if self.kind.is_ascii() {
419+
return bytepos;
420+
}
421+
let char_len = self.char_len();
422+
self.index
423+
.get_or_build(&self.data, char_len)
424+
.char_index_at_byte(&self.data, bytepos, char_len)
425+
}
426+
408427
pub fn nth_char(&self, index: usize) -> CodePoint {
409428
match self.as_str_kind() {
410429
PyKindStr::Ascii(s) => s[index].into(),

‎crates/common/src/wtf8_index.rs‎

Lines changed: 74 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,63 @@ impl Wtf8Index {
104104
}
105105
}
106106

107+
/// The index of the code point starting at byte offset `bytepos`, the
108+
/// inverse of [`Self::byte_offset`].
109+
///
110+
/// `data` must be the buffer the table was built for, `char_len` its code
111+
/// point count, and `bytepos` a code point boundary at or before its end.
112+
///
113+
/// Logarithmic rather than constant: the table is keyed by code point
114+
/// index, so going the other way is a search through it. The bracketing
115+
/// below is what keeps that search short -- a code point occupies one to
116+
/// four bytes, which pins the answer to a narrow band around `bytepos`
117+
/// before the first comparison.
118+
#[must_use]
119+
pub fn char_index_at_byte(&self, data: &Wtf8, bytepos: usize, char_len: usize) -> usize {
120+
let bytes_remaining = data.len() - bytepos;
121+
// At least one byte per remaining code point, and at most four, so the
122+
// group holding the answer lies between these.
123+
let mut group_min =
124+
usize::max(bytepos / 4, char_len.saturating_sub(bytes_remaining + 1)) >> 6;
125+
let mut group_max = usize::min(bytepos, char_len.saturating_sub(bytes_remaining / 4)) >> 6;
126+
while group_min < group_max {
127+
let middle = group_min.midpoint(group_max) + 1;
128+
if bytepos < self.groups[middle].base {
129+
group_max = middle - 1;
130+
} else {
131+
group_min = middle;
132+
}
133+
}
134+
135+
let base = self.groups[group_min].base;
136+
if base == bytepos {
137+
return group_min << 6;
138+
}
139+
// Walk the group's entries to the last one at or before `bytepos`,
140+
// then step the remaining code points, of which there are at most
141+
// three -- an entry covers four.
142+
let entries = if group_min == self.groups.len() - 1 {
143+
((char_len - 1) >> 2) & 0x0F
144+
} else {
145+
16
146+
};
147+
let mut index = group_min << 6;
148+
let mut pos = base;
149+
for entry in 0..entries {
150+
let at = base + self.groups[group_min].ofs[entry] as usize;
151+
if at >= bytepos {
152+
break;
153+
}
154+
pos = at;
155+
index = (group_min << 6) + (entry << 2) + 1;
156+
}
157+
while pos < bytepos {
158+
pos = next_pos(data, pos);
159+
index += 1;
160+
}
161+
index
162+
}
163+
107164
/// The table's heap footprint, in bytes.
108165
#[must_use]
109166
pub fn byte_size(&self) -> usize {
@@ -153,21 +210,34 @@ mod tests {
153210
use super::*;
154211
use crate::wtf8::{CodePoint, Wtf8Buf};
155212

156-
/// Every index of `s`, against the offsets its own iterator reports.
213+
/// Every index of `s`, both ways, against the offsets its own iterator
214+
/// reports.
157215
fn check(s: &Wtf8) {
158216
let expected: Vec<usize> = s
159217
.code_point_indices()
160218
.map(|(byte_offset, _)| byte_offset)
161219
.collect();
162-
let index = Wtf8Index::new(s, expected.len());
220+
let char_len = expected.len();
221+
let index = Wtf8Index::new(s, char_len);
163222
for (i, &want) in expected.iter().enumerate() {
164223
assert_eq!(
165224
index.byte_offset(s, i),
166225
want,
167-
"index {i} of {s:?} ({} code points)",
168-
expected.len()
226+
"index {i} of {s:?} ({char_len} code points)"
227+
);
228+
assert_eq!(
229+
index.char_index_at_byte(s, want, char_len),
230+
i,
231+
"byte {want} of {s:?} ({char_len} code points)"
169232
);
170233
}
234+
// One past the last code point is a boundary too, and the searches that
235+
// use this ask for it as an end bound.
236+
assert_eq!(
237+
index.char_index_at_byte(s, s.len(), char_len),
238+
char_len,
239+
"end of {s:?}"
240+
);
171241
}
172242

173243
fn wtf8(s: &str) -> Wtf8Buf {

‎crates/vm/src/anystr.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,11 @@ pub(crate) trait AnyStr {
147147
fn as_bytes(&self) -> &[u8];
148148
fn elements(&self) -> impl Iterator<Item = Self::Char>;
149149
fn get_bytes(&self, range: Range<usize>) -> &Self;
150-
// FIXME: get_chars is expensive for str
150+
/// The characters in `range`, which for a `str` payload means walking to
151+
/// both bounds -- the payload does not carry the string's character index.
152+
/// `PyStr` therefore converts its own ranges and does not reach the search
153+
/// helpers below through this; what remains are the byte strings, where a
154+
/// character range is already a byte range.
151155
fn get_chars(&self, range: Range<usize>) -> &Self;
152156
fn bytes_len(&self) -> usize;
153157
// NOTE: str::chars().count() consumes the O(n) time. But pystr::char_len does cache.

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

Lines changed: 48 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use super::{
99
use crate::{
1010
AsObject, Context, Py, PyExact, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult,
1111
TryFromBorrowedObject, VirtualMachine,
12-
anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper, adjust_indices},
12+
anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper, StringRange, adjust_indices},
1313
atomic_func,
1414
bytes_inner::{swapcase_ascii, title_ascii},
1515
cformat::cformat_string,
@@ -719,6 +719,13 @@ impl PyStr {
719719
self.data.char_index_to_byte(index)
720720
}
721721

722+
/// The character index of the character starting at byte offset `bytepos`,
723+
/// which must be a character boundary at or before the end.
724+
#[inline]
725+
pub fn byte_to_char_index(&self, bytepos: usize) -> usize {
726+
self.data.byte_to_char_index(bytepos)
727+
}
728+
722729
#[pymethod]
723730
#[inline(always)]
724731
pub const fn isascii(&self) -> bool {
@@ -924,11 +931,12 @@ impl PyStr {
924931

925932
#[pymethod]
926933
fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
927-
let (affix, substr) =
928-
match options.prepare(self.as_wtf8(), self.len(), |s, r| s.get_chars(r)) {
929-
Some(x) => x,
930-
None => return Ok(false),
931-
};
934+
let (affix, substr) = match options.prepare(self.as_wtf8(), self.len(), |s, r| {
935+
&s[self.char_index_to_byte(r.start)..self.char_index_to_byte(r.end)]
936+
}) {
937+
Some(x) => x,
938+
None => return Ok(false),
939+
};
932940
substr.py_starts_ends_with(
933941
&affix,
934942
"endswith",
@@ -944,11 +952,12 @@ impl PyStr {
944952
options: anystr::StartsEndsWithArgs,
945953
vm: &VirtualMachine,
946954
) -> PyResult<bool> {
947-
let (affix, substr) =
948-
match options.prepare(self.as_wtf8(), self.len(), |s, r| s.get_chars(r)) {
949-
Some(x) => x,
950-
None => return Ok(false),
951-
};
955+
let (affix, substr) = match options.prepare(self.as_wtf8(), self.len(), |s, r| {
956+
&s[self.char_index_to_byte(r.start)..self.char_index_to_byte(r.end)]
957+
}) {
958+
Some(x) => x,
959+
None => return Ok(false),
960+
};
952961
substr.py_starts_ends_with(
953962
&affix,
954963
"startswith",
@@ -1167,42 +1176,52 @@ impl PyStr {
11671176
Ok(vm.ctx.new_str(joined))
11681177
}
11691178

1170-
// FIXME: two traversals of str is expensive
1179+
/// The bytes of the character range `range`, or `None` if it is empty.
1180+
///
1181+
/// Both bounds go through the string's character index, so reaching a
1182+
/// range deep in the subject costs a lookup rather than a walk to it.
11711183
#[inline]
1172-
fn _to_char_idx(r: &Wtf8, byte_idx: usize) -> usize {
1173-
r[..byte_idx].code_points().count()
1184+
fn char_range_bytes(&self, range: Range<usize>) -> Option<(usize, &Wtf8)> {
1185+
if !range.is_normal() {
1186+
return None;
1187+
}
1188+
let start = self.char_index_to_byte(range.start);
1189+
let end = self.char_index_to_byte(range.end);
1190+
Some((start, &self.as_wtf8()[start..end]))
11741191
}
11751192

1193+
/// Searches the character range `range` with `find`, which answers in bytes
1194+
/// relative to the range, and reports the hit as a character index.
11761195
#[inline]
11771196
fn _find<F>(&self, args: FindArgs, find: F) -> Option<usize>
11781197
where
11791198
F: Fn(&Wtf8, &Wtf8) -> Option<usize>,
11801199
{
11811200
let (sub, range) = args.get_value(self.len());
1182-
self.as_wtf8().py_find(sub.as_wtf8(), range, find)
1201+
let (start, haystack) = self.char_range_bytes(range)?;
1202+
let found = find(haystack, sub.as_wtf8())?;
1203+
Some(self.byte_to_char_index(start + found))
11831204
}
11841205

11851206
#[pymethod]
11861207
fn find(&self, args: FindArgs) -> isize {
1187-
self._find(args, |r, s| Some(Self::_to_char_idx(r, r.find(s)?)))
1188-
.map_or(-1, |v| v as isize)
1208+
self._find(args, Wtf8::find).map_or(-1, |v| v as isize)
11891209
}
11901210

11911211
#[pymethod]
11921212
fn rfind(&self, args: FindArgs) -> isize {
1193-
self._find(args, |r, s| Some(Self::_to_char_idx(r, r.rfind(s)?)))
1194-
.map_or(-1, |v| v as isize)
1213+
self._find(args, Wtf8::rfind).map_or(-1, |v| v as isize)
11951214
}
11961215

11971216
#[pymethod]
11981217
fn index(&self, args: FindArgs, vm: &VirtualMachine) -> PyResult<usize> {
1199-
self._find(args, |r, s| Some(Self::_to_char_idx(r, r.find(s)?)))
1218+
self._find(args, Wtf8::find)
12001219
.ok_or_else(|| vm.new_value_error("substring not found"))
12011220
}
12021221

12031222
#[pymethod]
12041223
fn rindex(&self, args: FindArgs, vm: &VirtualMachine) -> PyResult<usize> {
1205-
self._find(args, |r, s| Some(Self::_to_char_idx(r, r.rfind(s)?)))
1224+
self._find(args, Wtf8::rfind)
12061225
.ok_or_else(|| vm.new_value_error("substring not found"))
12071226
}
12081227

@@ -1275,15 +1294,16 @@ impl PyStr {
12751294
#[pymethod]
12761295
fn count(&self, args: FindArgs) -> usize {
12771296
let (needle, range) = args.get_value(self.len());
1278-
self.as_wtf8().py_count(needle.as_wtf8(), range, |h, n| {
1279-
if n.is_empty() {
1297+
let chars = range.len();
1298+
self.char_range_bytes(range).map_or(0, |(_, haystack)| {
1299+
if needle.is_empty() {
12801300
// An empty needle sits between every pair of characters and at
1281-
// both ends, so it occurs once more than the haystack holds
1282-
// characters. Searching for it in the bytes would instead
1283-
// answer in encoded positions.
1284-
h.code_points().count() + 1
1301+
// both ends, so it occurs once more than the range holds
1302+
// characters. Counting it in the bytes would answer in encoded
1303+
// positions instead.
1304+
chars + 1
12851305
} else {
1286-
h.find_iter(n).count()
1306+
haystack.find_iter(needle.as_wtf8()).count()
12871307
}
12881308
})
12891309
}

0 commit comments

Comments
 (0)