RustPython · youknowone · Aug 14, 2026 · Aug 14, 2026 · Aug 14, 2026
diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs
@@ -21,6 +21,7 @@ pub mod rc;
 pub mod refcount;
 pub mod static_cell;
 pub mod str;
+pub mod wtf8_index;
 
 pub use rustpython_wtf8 as wtf8;
 

diff --git a/crates/common/src/str.rs b/crates/common/src/str.rs
@@ -1,7 +1,8 @@
 // spell-checker:ignore uncomputed
-use crate::atomic::{PyAtomic, Radium};
+use crate::atomic::{OncePtr, PyAtomic, Radium};
 use crate::format::CharLen;
 use crate::wtf8::{CodePoint, Wtf8, Wtf8Buf};
+use crate::wtf8_index::Wtf8Index;
 use ascii::{AsciiChar, AsciiStr, AsciiString};
 use core::fmt;
 use core::ops::{Bound, RangeBounds};
@@ -117,6 +118,55 @@ pub struct StrData {
     data: Box<Wtf8>,
     kind: StrKind,
     len: StrLen,
+    index: Wtf8IndexSlot,
+}
+
+/// A [`Wtf8Index`] built on first use.
+///
+/// The table is a pure function of `data`, so publishing it races benignly: a
+/// thread that loses the exchange drops its own copy and reads the winner's.
+#[derive(Default)]
+struct Wtf8IndexSlot(OncePtr<Wtf8Index>);
+
+impl Wtf8IndexSlot {
+    #[inline(always)]
+    fn new() -> Self {
+        Self(OncePtr::new())
+    }
+
+    #[inline]
+    fn get_or_build(&self, data: &Wtf8, char_len: usize) -> &Wtf8Index {
+        let index = self
+            .0
+            .get_or_init(|| Box::new(Wtf8Index::new(data, char_len)));
+        // The slot owns the table, never replaces it, and outlives the borrow.
+        unsafe { index.as_ref() }
+    }
+}
+
+impl fmt::Debug for Wtf8IndexSlot {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self.0.get() {
+            Some(_) => f.write_str("<built>"),
+            None => f.write_str("<unbuilt>"),
+        }
+    }
+}
+
+impl Clone for Wtf8IndexSlot {
+    /// A fresh slot: the clone copies the buffer, so it has to index that copy,
+    /// and the table is rebuilt on demand rather than eagerly here.
+    fn clone(&self) -> Self {
+        Self::new()
+    }
+}
+
+impl Drop for Wtf8IndexSlot {
+    fn drop(&mut self) {
+        if let Some(index) = self.0.get() {
+            drop(unsafe { Box::from_raw(index.as_ptr()) });
+        }
+    }
 }
 
 struct StrLen(PyAtomic<usize>);
@@ -163,6 +213,7 @@ impl Default for StrData {
             data: <Box<Wtf8>>::default(),
             kind: StrKind::Ascii,
             len: StrLen::zero(),
+            index: Wtf8IndexSlot::new(),
         }
     }
 }
@@ -193,6 +244,7 @@ impl From<Box<AsciiStr>> for StrData {
             len: value.len().into(),
             data: value.into(),
             kind: StrKind::Ascii,
+            index: Wtf8IndexSlot::new(),
         }
     }
 }
@@ -212,6 +264,7 @@ impl From<char> for StrData {
                 data: ch.to_string().into(),
                 kind: StrKind::Utf8,
                 len: 1.into(),
+                index: Wtf8IndexSlot::new(),
             }
         }
     }
@@ -226,6 +279,7 @@ impl From<CodePoint> for StrData {
                 data: Wtf8Buf::from(ch).into(),
                 kind: StrKind::Wtf8,
                 len: 1.into(),
+                index: Wtf8IndexSlot::new(),
             }
         }
     }
@@ -241,7 +295,12 @@ impl StrData {
             StrKind::Ascii => data.len().into(),
             _ => StrLen::uncomputed(),
         };
-        Self { data, kind, len }
+        Self {
+            data,
+            kind,
+            len,
+            index: Wtf8IndexSlot::new(),
+        }
     }
 
     /// # Safety
@@ -253,6 +312,7 @@ impl StrData {
             data,
             kind,
             len: char_len.into(),
+            index: Wtf8IndexSlot::new(),
         }
     }
 
@@ -322,6 +382,29 @@ impl StrData {
         len
     }
 
+    /// The byte offset the `index`-th code point starts at.
+    ///
+    /// An `index` at or past the end answers the buffer's byte length, so a
+    /// caller walking to a bound does not have to special-case it.
+    ///
+    /// O(1), but the first call on a non-ASCII string builds an index over the
+    /// whole buffer, so a caller that resolves a single index and stops is
+    /// better served by [`Self::nth_char`].
+    pub fn char_index_to_byte(&self, index: usize) -> usize {
+        // For ASCII the two units coincide, and the table would be a Nth entry
+        // saying N.
+        if self.kind.is_ascii() {
+            return index.min(self.data.len());
+        }
+        let char_len = self.char_len();
+        if index >= char_len {
+            return self.data.len();
+        }
+        self.index
+            .get_or_build(&self.data, char_len)
+            .byte_offset(&self.data, index)
+    }
+
     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
@@ -0,0 +1,229 @@
+// spell-checker:ignore rpython rlib rutf
+//! Random access into a WTF-8 buffer.
+//!
+//! WTF-8 is variable width, so a buffer's n-th code point can only be found by
+//! decoding the n-1 before it: [`Wtf8`]'s iterators are sequential, and
+//! resolving an index through them is O(n). Code that indexes the same string
+//! repeatedly -- a regex scan restarting at successive positions, say -- then
+//! walks the whole buffer once per index, which is quadratic in its length.
+//!
+//! [`Wtf8Index`] is the side table that makes the lookup O(1): one 24-byte
+//! group per 64 code points, so 0.375 bytes per code point. It is a cache, and
+//! holds no state of its own beyond the buffer's shape -- building it twice for
+//! the same buffer yields the same table.
+//!
+//! The layout is PyPy's `UTF8_INDEX_STORAGE` (`rpython/rlib/rutf8.py`).
+
+use crate::wtf8::Wtf8;
+
+/// One group of 64 code points.
+#[derive(Clone, Copy)]
+struct Group {
+    /// The byte offset the group's first code point starts at.
+    base: usize,
+    /// `ofs[i]` is the byte offset of the group's `4 * i + 1`-th code point,
+    /// relative to `base`. One entry covers four code points, so the widest
+    /// offset an entry has to hold is that of the 61st code point of a group,
+    /// at most `61 * 4 = 244` bytes in -- inside a `u8`, which is what buys the
+    /// table its density.
+    ofs: [u8; 16],
+}
+
+/// A code-point-index to byte-offset table for one WTF-8 buffer.
+pub struct Wtf8Index {
+    groups: Box<[Group]>,
+}
+
+impl Wtf8Index {
+    /// Builds the table for `data`, whose code point count is `char_len`.
+    ///
+    /// O(`data.len()`), and touches every byte, so it pays for itself only when
+    /// the caller goes on to index the buffer more than a couple of times.
+    #[must_use]
+    pub fn new(data: &Wtf8, char_len: usize) -> Self {
+        let mut groups = vec![
+            Group {
+                base: 0,
+                ofs: [0; 16],
+            };
+            char_len / 64 + 1
+        ];
+        // Signed: the countdown overshoots the last group -- the loop stops on
+        // the first negative value rather than at a group boundary.
+        let mut remaining = char_len as isize;
+        let mut base = 0;
+        let mut current = 0;
+        loop {
+            groups[current].base = base;
+            let mut next = base;
+            let mut group_filled = true;
+            for i in 0..16 {
+                // Past the end, step as if one more single-byte code point
+                // followed, so the entry stays in range and is never read.
+                next = if remaining == 0 {
+                    next + 1
+                } else {
+                    next_pos(data, next)
+                };
+                groups[current].ofs[i] = (next - base) as u8;
+                remaining -= 4;
+                if remaining < 0 {
+                    debug_assert_eq!(current + 1, groups.len());
+                    group_filled = false;
+                    break;
+                }
+                next = next_pos(data, next_pos(data, next_pos(data, next)));
+            }
+            if !group_filled {
+                break;
+            }
+            current += 1;
+            base = next;
+        }
+        Self {
+            groups: groups.into_boxed_slice(),
+        }
+    }
+
+    /// The byte offset of `data`'s `index`-th code point.
+    ///
+    /// `data` must be the buffer the table was built for, and `index` must be
+    /// below its code point count.
+    #[inline]
+    #[must_use]
+    pub fn byte_offset(&self, data: &Wtf8, index: usize) -> usize {
+        let group = &self.groups[index >> 6];
+        // The entry sits on the 4k+1-th code point of the group, so a lookup is
+        // one table read plus at most two steps in either direction.
+        let pos = group.base + group.ofs[(index >> 2) & 0x0F] as usize;
+        match index & 0x3 {
+            0 => prev_pos(data, pos),
+            1 => pos,
+            2 => next_pos(data, pos),
+            _ => next_pos(data, next_pos(data, pos)),
+        }
+    }
+
+    /// The table's heap footprint, in bytes.
+    #[must_use]
+    pub fn byte_size(&self) -> usize {
+        core::mem::size_of_val(&*self.groups)
+    }
+}
+
+/// The byte offset of the code point after the one at `pos`.
+///
+/// `data` must be well-formed WTF-8 and `pos` a code point boundary before its
+/// end -- reading only the lead byte is what makes this branch-light.
+#[inline]
+fn next_pos(data: &Wtf8, pos: usize) -> usize {
+    match data.as_bytes()[pos] {
+        0x00..=0x7F => pos + 1,
+        0x80..=0xDF => pos + 2,
+        0xE0..=0xEF => pos + 3,
+        _ => pos + 4,
+    }
+}
+
+/// The byte offset of the code point before the one at `pos`, which must not be
+/// zero.
+///
+/// A `pos` one past the end reads as the extra code point [`Wtf8Index::new`]
+/// steps over there.
+#[inline]
+fn prev_pos(data: &Wtf8, pos: usize) -> usize {
+    let data = data.as_bytes();
+    let mut pos = pos - 1;
+    if pos >= data.len() || data[pos] <= 0x7F {
+        return pos;
+    }
+    pos -= 1;
+    if data[pos] >= 0xC0 {
+        return pos;
+    }
+    pos -= 1;
+    if data[pos] >= 0xC0 {
+        return pos;
+    }
+    pos - 1
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::wtf8::{CodePoint, Wtf8Buf};
+
+    /// Every index of `s`, against the offsets its own iterator reports.
+    fn check(s: &Wtf8) {
+        let expected: Vec<usize> = s
+            .code_point_indices()
+            .map(|(byte_offset, _)| byte_offset)
+            .collect();
+        let index = Wtf8Index::new(s, expected.len());
+        for (i, &want) in expected.iter().enumerate() {
+            assert_eq!(
+                index.byte_offset(s, i),
+                want,
+                "index {i} of {s:?} ({} code points)",
+                expected.len()
+            );
+        }
+    }
+
+    fn wtf8(s: &str) -> Wtf8Buf {
+        Wtf8Buf::from(s)
+    }
+
+    #[test]
+    fn empty() {
+        check(wtf8("").as_ref());
+    }
+
+    #[test]
+    fn widths() {
+        // One case per encoded width, and the boundaries between them.
+        check(wtf8("abc").as_ref());
+        check(wtf8("\u{80}\u{7ff}").as_ref());
+        check(wtf8("\u{800}\u{ffff}").as_ref());
+        check(wtf8("\u{10000}\u{10ffff}").as_ref());
+        check(wtf8("a\u{80}\u{800}\u{10000}").as_ref());
+    }
+
+    #[test]
+    fn group_boundaries() {
+        // A group covers 64 code points and an entry four, so the interesting
+        // lengths are the ones on and around both.
+        for len in [1, 3, 4, 5, 63, 64, 65, 127, 128, 129, 255, 256, 257] {
+            for unit in ["a", "\u{80}", "\u{800}", "\u{10000}"] {
+                check(wtf8(&unit.repeat(len)).as_ref());
+            }
+            // Mixed widths, so a group's entries do not share a stride.
+            check(wtf8(&"a\u{80}\u{800}\u{10000}".repeat(len)).as_ref());
+        }
+    }
+
+    #[test]
+    fn lone_surrogates() {
+        let mut s = wtf8("a");
+        for cp in [0xD800, 0xDBFF, 0xDC00, 0xDFFF] {
+            s.push(CodePoint::from_u32(cp).unwrap());
+            s.push_str("b");
+        }
+        check(s.as_ref());
+
+        // Surrogates only, spanning more than one group.
+        let mut s = wtf8("");
+        for i in 0..200 {
+            s.push(CodePoint::from_u32(0xD800 + (i % 0x400)).unwrap());
+        }
+        check(s.as_ref());
+    }
+
+    #[test]
+    fn byte_size_is_one_group_per_64_code_points() {
+        let s = wtf8(&"\u{10000}".repeat(200));
+        let index = Wtf8Index::new(s.as_ref(), 200);
+        assert_eq!(index.byte_size(), (200 / 64 + 1) * size_of::<Group>());
+        assert_eq!(size_of::<Group>(), 24);
+    }
+}