|
| 1 | +// spell-checker:ignore rpython rlib rutf |
| 2 | +//! Random access into a WTF-8 buffer. |
| 3 | +//! |
| 4 | +//! WTF-8 is variable width, so a buffer's n-th code point can only be found by |
| 5 | +//! decoding the n-1 before it: [`Wtf8`]'s iterators are sequential, and |
| 6 | +//! resolving an index through them is O(n). Code that indexes the same string |
| 7 | +//! repeatedly -- a regex scan restarting at successive positions, say -- then |
| 8 | +//! walks the whole buffer once per index, which is quadratic in its length. |
| 9 | +//! |
| 10 | +//! [`Wtf8Index`] is the side table that makes the lookup O(1): one 24-byte |
| 11 | +//! group per 64 code points, so 0.375 bytes per code point. It is a cache, and |
| 12 | +//! holds no state of its own beyond the buffer's shape -- building it twice for |
| 13 | +//! the same buffer yields the same table. |
| 14 | +//! |
| 15 | +//! The layout is PyPy's `UTF8_INDEX_STORAGE` (`rpython/rlib/rutf8.py`). |
| 16 | +
|
| 17 | +use crate::wtf8::Wtf8; |
| 18 | + |
| 19 | +/// One group of 64 code points. |
| 20 | +#[derive(Clone, Copy)] |
| 21 | +struct Group { |
| 22 | + /// The byte offset the group's first code point starts at. |
| 23 | + base: usize, |
| 24 | + /// `ofs[i]` is the byte offset of the group's `4 * i + 1`-th code point, |
| 25 | + /// relative to `base`. One entry covers four code points, so the widest |
| 26 | + /// offset an entry has to hold is that of the 61st code point of a group, |
| 27 | + /// at most `61 * 4 = 244` bytes in -- inside a `u8`, which is what buys the |
| 28 | + /// table its density. |
| 29 | + ofs: [u8; 16], |
| 30 | +} |
| 31 | + |
| 32 | +/// A code-point-index to byte-offset table for one WTF-8 buffer. |
| 33 | +pub struct Wtf8Index { |
| 34 | + groups: Box<[Group]>, |
| 35 | +} |
| 36 | + |
| 37 | +impl Wtf8Index { |
| 38 | + /// Builds the table for `data`, whose code point count is `char_len`. |
| 39 | + /// |
| 40 | + /// O(`data.len()`), and touches every byte, so it pays for itself only when |
| 41 | + /// the caller goes on to index the buffer more than a couple of times. |
| 42 | + #[must_use] |
| 43 | + pub fn new(data: &Wtf8, char_len: usize) -> Self { |
| 44 | + let mut groups = vec![ |
| 45 | + Group { |
| 46 | + base: 0, |
| 47 | + ofs: [0; 16], |
| 48 | + }; |
| 49 | + char_len / 64 + 1 |
| 50 | + ]; |
| 51 | + // Signed: the countdown overshoots the last group -- the loop stops on |
| 52 | + // the first negative value rather than at a group boundary. |
| 53 | + let mut remaining = char_len as isize; |
| 54 | + let mut base = 0; |
| 55 | + let mut current = 0; |
| 56 | + loop { |
| 57 | + groups[current].base = base; |
| 58 | + let mut next = base; |
| 59 | + let mut group_filled = true; |
| 60 | + for i in 0..16 { |
| 61 | + // Past the end, step as if one more single-byte code point |
| 62 | + // followed, so the entry stays in range and is never read. |
| 63 | + next = if remaining == 0 { |
| 64 | + next + 1 |
| 65 | + } else { |
| 66 | + next_pos(data, next) |
| 67 | + }; |
| 68 | + groups[current].ofs[i] = (next - base) as u8; |
| 69 | + remaining -= 4; |
| 70 | + if remaining < 0 { |
| 71 | + debug_assert_eq!(current + 1, groups.len()); |
| 72 | + group_filled = false; |
| 73 | + break; |
| 74 | + } |
| 75 | + next = next_pos(data, next_pos(data, next_pos(data, next))); |
| 76 | + } |
| 77 | + if !group_filled { |
| 78 | + break; |
| 79 | + } |
| 80 | + current += 1; |
| 81 | + base = next; |
| 82 | + } |
| 83 | + Self { |
| 84 | + groups: groups.into_boxed_slice(), |
| 85 | + } |
| 86 | + } |
| 87 | + |
| 88 | + /// The byte offset of `data`'s `index`-th code point. |
| 89 | + /// |
| 90 | + /// `data` must be the buffer the table was built for, and `index` must be |
| 91 | + /// below its code point count. |
| 92 | + #[inline] |
| 93 | + #[must_use] |
| 94 | + pub fn byte_offset(&self, data: &Wtf8, index: usize) -> usize { |
| 95 | + let group = &self.groups[index >> 6]; |
| 96 | + // The entry sits on the 4k+1-th code point of the group, so a lookup is |
| 97 | + // one table read plus at most two steps in either direction. |
| 98 | + let pos = group.base + group.ofs[(index >> 2) & 0x0F] as usize; |
| 99 | + match index & 0x3 { |
| 100 | + 0 => prev_pos(data, pos), |
| 101 | + 1 => pos, |
| 102 | + 2 => next_pos(data, pos), |
| 103 | + _ => next_pos(data, next_pos(data, pos)), |
| 104 | + } |
| 105 | + } |
| 106 | + |
| 107 | + /// The table's heap footprint, in bytes. |
| 108 | + #[must_use] |
| 109 | + pub fn byte_size(&self) -> usize { |
| 110 | + core::mem::size_of_val(&*self.groups) |
| 111 | + } |
| 112 | +} |
| 113 | + |
| 114 | +/// The byte offset of the code point after the one at `pos`. |
| 115 | +/// |
| 116 | +/// `data` must be well-formed WTF-8 and `pos` a code point boundary before its |
| 117 | +/// end -- reading only the lead byte is what makes this branch-light. |
| 118 | +#[inline] |
| 119 | +fn next_pos(data: &Wtf8, pos: usize) -> usize { |
| 120 | + match data.as_bytes()[pos] { |
| 121 | + 0x00..=0x7F => pos + 1, |
| 122 | + 0x80..=0xDF => pos + 2, |
| 123 | + 0xE0..=0xEF => pos + 3, |
| 124 | + _ => pos + 4, |
| 125 | + } |
| 126 | +} |
| 127 | + |
| 128 | +/// The byte offset of the code point before the one at `pos`, which must not be |
| 129 | +/// zero. |
| 130 | +/// |
| 131 | +/// A `pos` one past the end reads as the extra code point [`Wtf8Index::new`] |
| 132 | +/// steps over there. |
| 133 | +#[inline] |
| 134 | +fn prev_pos(data: &Wtf8, pos: usize) -> usize { |
| 135 | + let data = data.as_bytes(); |
| 136 | + let mut pos = pos - 1; |
| 137 | + if pos >= data.len() || data[pos] <= 0x7F { |
| 138 | + return pos; |
| 139 | + } |
| 140 | + pos -= 1; |
| 141 | + if data[pos] >= 0xC0 { |
| 142 | + return pos; |
| 143 | + } |
| 144 | + pos -= 1; |
| 145 | + if data[pos] >= 0xC0 { |
| 146 | + return pos; |
| 147 | + } |
| 148 | + pos - 1 |
| 149 | +} |
| 150 | + |
| 151 | +#[cfg(test)] |
| 152 | +mod tests { |
| 153 | + use super::*; |
| 154 | + use crate::wtf8::{CodePoint, Wtf8Buf}; |
| 155 | + |
| 156 | + /// Every index of `s`, against the offsets its own iterator reports. |
| 157 | + fn check(s: &Wtf8) { |
| 158 | + let expected: Vec<usize> = s |
| 159 | + .code_point_indices() |
| 160 | + .map(|(byte_offset, _)| byte_offset) |
| 161 | + .collect(); |
| 162 | + let index = Wtf8Index::new(s, expected.len()); |
| 163 | + for (i, &want) in expected.iter().enumerate() { |
| 164 | + assert_eq!( |
| 165 | + index.byte_offset(s, i), |
| 166 | + want, |
| 167 | + "index {i} of {s:?} ({} code points)", |
| 168 | + expected.len() |
| 169 | + ); |
| 170 | + } |
| 171 | + } |
| 172 | + |
| 173 | + fn wtf8(s: &str) -> Wtf8Buf { |
| 174 | + Wtf8Buf::from(s) |
| 175 | + } |
| 176 | + |
| 177 | + #[test] |
| 178 | + fn empty() { |
| 179 | + check(wtf8("").as_ref()); |
| 180 | + } |
| 181 | + |
| 182 | + #[test] |
| 183 | + fn widths() { |
| 184 | + // One case per encoded width, and the boundaries between them. |
| 185 | + check(wtf8("abc").as_ref()); |
| 186 | + check(wtf8("\u{80}\u{7ff}").as_ref()); |
| 187 | + check(wtf8("\u{800}\u{ffff}").as_ref()); |
| 188 | + check(wtf8("\u{10000}\u{10ffff}").as_ref()); |
| 189 | + check(wtf8("a\u{80}\u{800}\u{10000}").as_ref()); |
| 190 | + } |
| 191 | + |
| 192 | + #[test] |
| 193 | + fn group_boundaries() { |
| 194 | + // A group covers 64 code points and an entry four, so the interesting |
| 195 | + // lengths are the ones on and around both. |
| 196 | + for len in [1, 3, 4, 5, 63, 64, 65, 127, 128, 129, 255, 256, 257] { |
| 197 | + for unit in ["a", "\u{80}", "\u{800}", "\u{10000}"] { |
| 198 | + check(wtf8(&unit.repeat(len)).as_ref()); |
| 199 | + } |
| 200 | + // Mixed widths, so a group's entries do not share a stride. |
| 201 | + check(wtf8(&"a\u{80}\u{800}\u{10000}".repeat(len)).as_ref()); |
| 202 | + } |
| 203 | + } |
| 204 | + |
| 205 | + #[test] |
| 206 | + fn lone_surrogates() { |
| 207 | + let mut s = wtf8("a"); |
| 208 | + for cp in [0xD800, 0xDBFF, 0xDC00, 0xDFFF] { |
| 209 | + s.push(CodePoint::from_u32(cp).unwrap()); |
| 210 | + s.push_str("b"); |
| 211 | + } |
| 212 | + check(s.as_ref()); |
| 213 | + |
| 214 | + // Surrogates only, spanning more than one group. |
| 215 | + let mut s = wtf8(""); |
| 216 | + for i in 0..200 { |
| 217 | + s.push(CodePoint::from_u32(0xD800 + (i % 0x400)).unwrap()); |
| 218 | + } |
| 219 | + check(s.as_ref()); |
| 220 | + } |
| 221 | + |
| 222 | + #[test] |
| 223 | + fn byte_size_is_one_group_per_64_code_points() { |
| 224 | + let s = wtf8(&"\u{10000}".repeat(200)); |
| 225 | + let index = Wtf8Index::new(s.as_ref(), 200); |
| 226 | + assert_eq!(index.byte_size(), (200 / 64 + 1) * size_of::<Group>()); |
| 227 | + assert_eq!(size_of::<Group>(), 24); |
| 228 | + } |
| 229 | +} |
0 commit comments