perf(vm): cache string character indexes with char_index 🧵

Cargo.toml

@@ -14,6 +14,7 @@ license = "MIT"
 
 [workspace.dependencies]
 ahash = { version = "0.8.12" }
+char_index = "0.1.5"
 anyhow = "1.0.103"
 clap = { version = "4.6.1", features = ["derive"] }
 criterion = { version = "0.8.2", features = ["html_reports"] }

manual/src/reference/types/string.md

@@ -5,7 +5,9 @@ iterate over a string you get strings of length 1. Just like in Rust strings are
 UTF-8. This means that you can't store arbitrary binary data in a String.
 
 Indexing into a String is done by UTF-8 codepoint (equivalent to Rust's `char`) rather than by byte offset.
-This means that indexing into a string is `O(n)` instead of `O(1)`.
+The first indexed read builds a character index in `O(n)` time. Subsequent reads use
+that cache: `O(1)` for ASCII strings and `O(log n)` in the worst case for Unicode strings.
+Mutating the string discards the cache, so the next indexed read rebuilds it.
 
 ```ndc
 let string = "I ❤ Andy C++";

ndc_stdlib/src/index.rs

@@ -180,7 +180,7 @@ fn register_set(env: &mut FunctionRegistry<Rc<NativeFunction>>) {
 fn vm_sequence_length(v: &Value) -> Option<usize> {
     match v {
         Value::Object(obj) => match obj.as_ref() {
-            Object::String(s) => Some(s.borrow().chars().count()),
+            Object::String(s) => Some(s.char_count()),
             Object::List(l) => Some(l.borrow().len()),
             Object::Tuple(t) => Some(t.len()),
             Object::Map { entries, .. } => Some(entries.borrow().len()),
@@ -297,15 +297,16 @@ fn vm_get_at_index(container: &Value, index_value: &Value, vm: &mut Vm) -> Resul
                 }
             }
             Object::String(s) => {
-                let s = s.borrow();
+                let s = s.indexed();
                 match extract_vm_offset(index_value, size)? {
                     VmOffset::Element(idx) => {
-                        let ch = s.chars().nth(idx).expect("bounds already checked");
+                        let ch = s.get_char(idx).expect("bounds already checked");
                         Ok(Value::string(ch.to_string()))
                     }
                     VmOffset::Range(from, to) => {
-                        let result: String = s.chars().skip(from).take(to - from).collect();
-                        Ok(Value::string(result))
+                        let from = s.get_index(from).unwrap_or(s.len());
+                        let to = s.get_index(to).unwrap_or(s.len());
+                        Ok(Value::string(&s[from..to]))
                     }
                 }
             }
@@ -407,15 +408,16 @@ fn vm_get_at_index_simple(container: &Value, index_value: &Value) -> Result<Valu
                 }
             }
             Object::String(s) => {
-                let s = s.borrow();
+                let s = s.indexed();
                 match extract_vm_offset(index_value, size)? {
                     VmOffset::Element(idx) => {
-                        let ch = s.chars().nth(idx).expect("bounds already checked");
+                        let ch = s.get_char(idx).expect("bounds already checked");
                         Ok(Value::string(ch.to_string()))
                     }
                     VmOffset::Range(from, to) => {
-                        let result: String = s.chars().skip(from).take(to - from).collect();
-                        Ok(Value::string(result))
+                        let from = s.get_index(from).unwrap_or(s.len());
+                        let to = s.get_index(to).unwrap_or(s.len());
+                        Ok(Value::string(&s[from..to]))
                     }
                 }
             }

ndc_stdlib/src/sequence.rs

@@ -320,7 +320,7 @@ mod inner {
                 Object::List(l) => Ok(l.borrow().len() as i64),
                 Object::Tuple(t) => Ok(t.len() as i64),
                 Object::Deque(d) => Ok(d.borrow().len() as i64),
-                Object::String(s) => Ok(s.borrow().chars().count() as i64),
+                Object::String(s) => Ok(s.char_count() as i64),
                 Object::Map { entries, .. } => Ok(entries.borrow().len() as i64),
                 _ => Err(anyhow!(
                     "cannot determine the length of {}",

ndc_stdlib/src/string.rs

@@ -1,10 +1,9 @@
 use ndc_macros::export_module;
-use ndc_vm::value::{SeqValue, Value};
+use ndc_vm::value::{SeqValue, Value, VmString};
 
-use std::cell::RefCell;
 use std::rc::Rc;
 
-type StringRepr = Rc<RefCell<String>>;
+type StringRepr = Rc<VmString>;
 
 use anyhow::{Context, anyhow};
 

ndc_vm/Cargo.toml

@@ -8,6 +8,7 @@ trace = []
 
 [dependencies]
 ahash.workspace = true
+char_index.workspace = true
 ndc_core.workspace = true
 ndc_lexer.workspace = true
 ndc_parser.workspace = true

ndc_vm/src/iterator.rs

@@ -1,3 +1,4 @@
+use crate::value::VmString;
 use crate::{Object, OrdValue, Value, ValueIter};
 use std::cell::RefCell;
 use std::cmp::Reverse;
@@ -639,12 +640,12 @@ impl VmIterator for EnumerateIter {
 
 /// Iterates over string characters, yielding each as a single-char string
 pub struct StringIter {
-    string: Rc<RefCell<String>>,
+    string: Rc<VmString>,
     byte_offset: usize,
 }
 
 impl StringIter {
-    pub fn new(string: Rc<RefCell<String>>) -> Self {
+    pub fn new(string: Rc<VmString>) -> Self {
         Self {
             string,
             byte_offset: 0,
@@ -672,7 +673,7 @@ impl VmIterator for StringIter {
     fn deep_copy(&self) -> Option<SharedIterator> {
         let s = self.string.borrow().clone();
         Some(Rc::new(RefCell::new(Self {
-            string: Rc::new(RefCell::new(s)),
+            string: Rc::new(VmString::new(s)),
             byte_offset: self.byte_offset,
         })))
     }

ndc_vm/src/value/mod.rs

@@ -1,10 +1,12 @@
 mod function;
 mod number;
 mod numeric;
+mod string;
 
 pub use function::*;
 pub use number::{AdvancedNumber, BinaryOperatorError, ExactFraction, NumberToFloatError};
 pub use numeric::{NumericMode, NumericRef};
+pub use string::VmString;
 
 use crate::iterator::SharedIterator;
 use ndc_core::StaticType;
@@ -64,7 +66,7 @@ pub enum Value {
 #[derive(Clone)]
 pub enum Object {
     Some(Value),
-    String(Rc<RefCell<String>>),
+    String(Rc<VmString>),
     List(RefCell<Vec<Value>>),
     Tuple(Vec<Value>),
     Map {
@@ -169,12 +171,12 @@ impl Value {
     }
 
     pub fn string<S: Into<String>>(string: S) -> Self {
-        Self::Object(Rc::new(Object::String(Rc::new(RefCell::new(
+        Self::Object(Rc::new(Object::String(Rc::new(VmString::new(
             string.into(),
         )))))
     }
 
-    pub fn from_string_rc(rc: Rc<RefCell<String>>) -> Self {
+    pub fn from_string_rc(rc: Rc<VmString>) -> Self {
         Self::Object(Rc::new(Object::String(rc)))
     }
 
@@ -217,9 +219,9 @@ impl Value {
     pub fn shallow_clone(&self) -> Self {
         match self {
             Self::Object(obj) => match obj.as_ref() {
-                Object::String(rc) => Self::Object(Rc::new(Object::String(Rc::new(RefCell::new(
-                    rc.borrow().clone(),
-                ))))),
+                Object::String(rc) => Self::Object(Rc::new(Object::String(Rc::new(
+                    VmString::new(rc.borrow().clone()),
+                )))),
                 Object::List(refcell) => Self::Object(Rc::new(Object::List(RefCell::new(
                     refcell.borrow().clone(),
                 )))),
@@ -597,7 +599,7 @@ impl Object {
     pub fn deep_copy(&self) -> Self {
         match self {
             Self::Some(v) => Self::Some(v.deep_copy()),
-            Self::String(rc) => Self::String(Rc::new(RefCell::new(rc.borrow().clone()))),
+            Self::String(rc) => Self::String(Rc::new(VmString::new(rc.borrow().clone()))),
             Self::List(refcell) => Self::List(RefCell::new(
                 refcell.borrow().iter().map(Value::deep_copy).collect(),
             )),

ndc_vm/src/value/string.rs

@@ -0,0 +1,107 @@
+use char_index::OwnedIndexedChars;
+use std::cell::{Ref, RefCell, RefMut};
+
+/// Mutable UTF-8 storage with a character index built on the first indexed read.
+/// Both states own the same string allocation; mutable borrows discard the index.
+pub struct VmString {
+    storage: RefCell<Storage>,
+}
+
+enum Storage {
+    Plain(String),
+    Indexed(OwnedIndexedChars),
+}
+
+impl VmString {
+    pub fn new(string: String) -> Self {
+        Self {
+            storage: RefCell::new(Storage::Plain(string)),
+        }
+    }
+
+    pub fn borrow(&self) -> Ref<'_, String> {
+        Ref::map(self.storage.borrow(), |storage| match storage {
+            Storage::Plain(string) => string,
+            Storage::Indexed(string) => string.as_string(),
+        })
+    }
+
+    pub fn borrow_mut(&self) -> RefMut<'_, String> {
+        let mut storage = self.storage.borrow_mut();
+        if matches!(*storage, Storage::Indexed(_)) {
+            let Storage::Indexed(string) =
+                std::mem::replace(&mut *storage, Storage::Plain(String::new()))
+            else {
+                unreachable!()
+            };
+            *storage = Storage::Plain(string.into_string());
+        }
+        RefMut::map(storage, |storage| match storage {
+            Storage::Plain(string) => string,
+            Storage::Indexed(_) => unreachable!(),
+        })
+    }
+
+    /// Uses the cached count when available without building an index for `.len`.
+    pub fn char_count(&self) -> usize {
+        match &*self.storage.borrow() {
+            Storage::Plain(string) => string.chars().count(),
+            Storage::Indexed(string) => string.char_count(),
+        }
+    }
+
+    pub fn indexed(&self) -> Ref<'_, OwnedIndexedChars> {
+        if matches!(*self.storage.borrow(), Storage::Plain(_)) {
+            let mut storage = self.storage.borrow_mut();
+            let Storage::Plain(string) =
+                std::mem::replace(&mut *storage, Storage::Plain(String::new()))
+            else {
+                unreachable!()
+            };
+            *storage = Storage::Indexed(OwnedIndexedChars::new(string));
+        }
+        Ref::map(self.storage.borrow(), |storage| match storage {
+            Storage::Indexed(string) => string,
+            Storage::Plain(_) => unreachable!(),
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::VmString;
+
+    #[test]
+    fn indexed_reads_match_utf8_boundaries_across_rollovers() {
+        for source in [String::new(), "abcd".repeat(1024), "aé中😀".repeat(1024)] {
+            let string = VmString::new(source.clone());
+            let indexed = string.indexed();
+            assert_eq!(indexed.char_count(), source.chars().count());
+            for (index, (offset, ch)) in source.char_indices().enumerate() {
+                assert_eq!(indexed.get_char(index), Some(ch));
+                assert_eq!(indexed.get_index(index), Some(offset));
+            }
+            assert_eq!(indexed.get_char(indexed.char_count()), None);
+            assert_eq!(indexed.get_index(indexed.char_count()), None);
+        }
+    }
+
+    #[test]
+    fn mutation_invalidates_index_without_copying_string() {
+        let string = VmString::new("aé中😀".into());
+        let original_ptr = string.borrow().as_ptr();
+        assert_eq!(string.indexed().get_char(3), Some('😀'));
+        assert_eq!(string.borrow().as_ptr(), original_ptr);
+        {
+            let mut plain = string.borrow_mut();
+            assert_eq!(plain.as_ptr(), original_ptr);
+            plain.replace_range(1..3, "b");
+        }
+        assert_eq!(string.char_count(), 4);
+        assert_eq!(string.indexed().get_char(1), Some('b'));
+        assert_eq!(string.indexed().get_char(3), Some('😀'));
+        string.borrow_mut().clear();
+        assert_eq!(string.char_count(), 0);
+        assert_eq!(string.indexed().get_char(0), None);
+    }
+}

tests/functional/programs/003_string/001_index_cache_mutation.ndc

@@ -0,0 +1,41 @@
+// Populate the index before each mutation, including writes through an alias.
+let s = "aé中😀";
+let alias = s;
+assert_eq(s[3], "😀");
+alias[1] = "xy";
+assert_eq(s.len, 5);
+assert_eq(s[1..4], "xy中");
+assert_eq(s[-1], "😀");
+
+s[1..4] = "Ω";
+assert_eq(s.len, 3);
+assert_eq(s[0..3], "aΩ😀");
+assert_eq(s[3..3], "");
+s.append("z");
+assert_eq(s[-1], "z");
+s.reverse();
+assert_eq(s[0], "z");
+assert_eq(s[-1], "a");
+
+s ++= s;
+assert_eq(s.len, 8);
+assert_eq(s[4..8], "z😀Ωa");
+s <>= "é";
+assert_eq(s[8], "é");
+
+// Copies have independent storage, including after an indexed read.
+let shallow = s.clone();
+let deep = s.deepcopy();
+assert_eq(shallow[0], "z");
+assert_eq(deep[0], "z");
+s[0] = "中";
+assert_eq(s[0], "中");
+assert_eq(shallow[0], "z");
+assert_eq(deep[0], "z");
+
+// Exercise the generic indexing overload as well as statically typed strings.
+fn at_end(value: Any) => value[-1];
+assert_eq(at_end(s), "é");
+s[0..s.len] = "";
+assert_eq(s.len, 0);
+assert_eq(s[0..0], "");