Merge pull request #6783 from dotCooCoo/fix/js-flat-charges-heap

Scripts: operators, JSON.stringify and join count a flattened string against the heap limit

crates/js/src/builtins/array.rs

@@ -13,6 +13,7 @@ use super::{
 };
 use crate::heap::{Class, NativeKind, ObjId};
 use crate::interp::{Abrupt, Interp, JsResult, strict_eq};
+use crate::string_builtins::flat_charged;
 use crate::value::{Str, Value, to_uint32};
 
 pub(super) fn install(interp: &mut Interp) -> JsResult<()> {
@@ -278,17 +279,20 @@ fn join_with(interp: &mut Interp, o: ObjId, separator: &Value, item: fn(&mut Int
     if len == 0 {
         return Ok(Value::str(""));
     }
-    let flat = sep.flat();
-    let sep = &flat[..];
-    // A result that cannot fit is refused before any element is read.
-    let least = (sep.len() as u64).saturating_mul(u64::from(len - 1));
-    if least > (interp.limits.heap_bytes / 4) as u64 {
-        return Err(Abrupt::Limit(crate::LimitKind::Heap));
-    }
-    interp.join_stack.push(o);
-    let r = join_elements(interp, o, len, sep, item);
-    interp.join_stack.pop();
-    r
+    // The separator stays rooted for the whole join, so its flattened copy stays counted.
+    interp.rooted(&[Value::String(sep.clone())], |interp| {
+        let flat = flat_charged(interp, &sep)?;
+        let sep = &flat[..];
+        // A result that cannot fit is refused before any element is read.
+        let least = (sep.len() as u64).saturating_mul(u64::from(len - 1));
+        if least > (interp.limits.heap_bytes / 4) as u64 {
+            return Err(Abrupt::Limit(crate::LimitKind::Heap));
+        }
+        interp.join_stack.push(o);
+        let r = join_elements(interp, o, len, sep, item);
+        interp.join_stack.pop();
+        r
+    })
 }
 
 fn join_elements(interp: &mut Interp, o: ObjId, len: u32, sep: &[u16], item: fn(&mut Interp, Value) -> JsResult<Option<Str>>) -> JsResult<Value> {
@@ -301,7 +305,10 @@ fn join_elements(interp: &mut Interp, o: ObjId, len: u32, sep: &[u16], item: fn(
         let Some(s) = item(interp, v)? else { continue };
         out.push_repeated(interp, sep, k.saturating_sub(seps))?;
         seps = k;
-        out.push_units(interp, &s.flat())?;
+        interp.rooted(&[Value::String(s.clone())], |interp| {
+            let units = flat_charged(interp, &s)?;
+            out.push_units(interp, &units)
+        })?;
     }
     out.push_repeated(interp, sep, u64::from(len - 1).saturating_sub(seps))?;
     Ok(out.finish(interp))
@@ -472,7 +479,7 @@ fn unshift(interp: &mut Interp, this: Value, args: &[Value], _: bool) -> JsResul
     with_this(interp, &this, |interp, o| {
         let len = u64::from(length_of(interp, o)?);
         let count = args.len() as u64;
-        // With no arguments each element is still read and written back in place (steps 6 and 7).
+        // With no arguments each element is still read and written back in place (step 6).
         let maps = vec![Map::Shift(-1), Map::Shift(count as i64 - 1)];
         let mut it = Indices::new(interp, o, 1, len + 1, false, maps)?;
         while let Some(k) = it.next(interp)? {

crates/js/src/builtins/json.rs

@@ -13,6 +13,7 @@ use std::collections::HashSet;
 use super::{Builder, ELEMENT_BYTES, arg, data, key_of, length_of, method, new_array, new_string, to_integer};
 use crate::heap::{Class, Key, OBJECT_BYTES, OPEN, ObjId, Object, Prop, key, string_bytes};
 use crate::interp::{Desc, ErrorType, Hint, Interp, JsResult};
+use crate::string_builtins::flat_charged;
 use crate::value::{Str, Value, number_to_str};
 
 pub(super) fn install(interp: &mut Interp) -> JsResult<()> {
@@ -473,7 +474,7 @@ fn settings(interp: &mut Interp, replacer: &Value, space: &Value) -> JsResult<Se
             let n = to_integer(interp, &space)?.clamp(0.0, 10.0) as usize;
             s.gap = vec![u16::from(b' '); n];
         }
-        Value::String(t) => s.gap = t.flat().iter().take(10).copied().collect(),
+        Value::String(t) => s.gap = flat_charged(interp, &t)?.iter().take(10).copied().collect(),
         _ => {}
     }
     Ok(s)
@@ -632,7 +633,10 @@ fn write_primitive(interp: &mut Interp, out: &mut Builder, v: &Value) -> JsResul
         Value::Bool(b) => out.push_str(interp, if *b { "true" } else { "false" }),
         Value::Number(n) if n.is_finite() => out.push_units(interp, &number_to_str(*n).flat()),
         Value::Number(_) => out.push_str(interp, "null"),
-        Value::String(s) => write_quoted(interp, out, &s.flat()),
+        Value::String(s) => interp.rooted(std::slice::from_ref(v), |i| {
+            let units = flat_charged(i, s)?;
+            write_quoted(i, out, &units)
+        }),
         Value::Undefined | Value::Object(_) => Ok(()),
     }
 }

crates/js/src/interp.rs

@@ -27,6 +27,7 @@ use crate::heap::{
     ArgMap, Binding, CONFIGURABLE, Class, ENUMERABLE, ENV_BYTES, Env, EnvId, ForIn, Func, HIDDEN, Heap, Key, NativeKind, OBJECT_BYTES, OPEN, ObjId,
     Object, PROP_BYTES, Prop, PropMap, ROPE_BYTES, Roots, WRITABLE, code_bytes, key, string_bytes,
 };
+use crate::string_builtins::flat_charged;
 use crate::value::{Str, Value, array_index, index_key, number_to_str, string_to_number, to_int32, to_uint32};
 use crate::{LimitKind, Limits, Options};
 
@@ -205,6 +206,12 @@ pub(crate) struct Interp {
     pub(crate) date_host: crate::date_builtins::DateHost,
     /// The bytes native code holds outside the heap's cells through [`HeldBytes`].
     pub(crate) held: Rc<HeldCount>,
+    /// Bytes of the copies [`Self::get_own_property`] made flattening the rope of a String
+    /// object, which the heap's total adds before the next limit check.
+    flattened: std::cell::Cell<usize>,
+    /// Fuel for the units those copies took and the rope nodes it read in place, which the next
+    /// step takes.
+    pub(crate) flattened_fuel: std::cell::Cell<u64>,
     /// The global code's default XML namespace URI (ECMA-357 section 12.1).
     pub(crate) xml_default_uri: Key,
     /// Set by the host to end the current slice of a resumable run at the next safe point.
@@ -450,6 +457,8 @@ impl Interp {
             files: None,
             date_host: crate::date_builtins::DateHost::default(),
             held: Rc::default(),
+            flattened: std::cell::Cell::new(0),
+            flattened_fuel: std::cell::Cell::new(0),
             xml_default_uri: Rc::from(Vec::new()),
             pause: Rc::default(),
             in_slice: false,
@@ -569,6 +578,7 @@ impl Interp {
 
     fn tick(&mut self) -> JsResult<()> {
         self.use_fuel()?;
+        crate::regexp_vm::use_fuel_n(self, self.flattened_fuel.replace(0))?;
         self.settle_held();
         if self.heap.bytes > self.gc_threshold {
             self.collect_garbage();
@@ -595,10 +605,12 @@ impl Interp {
         Ok(())
     }
 
-    /// Takes the bytes that dropped [`HeldBytes`] released off the heap's total.
+    /// Takes the bytes that dropped [`HeldBytes`] released off the heap's total, and adds the
+    /// bytes of the ropes [`Self::get_own_property`] flattened.
     fn settle_held(&mut self) {
         let released = self.held.released.replace(0);
-        self.heap.bytes = self.heap.bytes.saturating_sub(released);
+        let flattened = self.flattened.replace(0);
+        self.heap.bytes = self.heap.bytes.saturating_sub(released).saturating_add(flattened);
     }
 
     /// Collects garbage now and sets when the next collection runs.
@@ -680,6 +692,7 @@ impl Interp {
         self.heap.collect(&roots);
         // The collection counts the live total again, so bytes already released are not in it.
         self.held.released.set(0);
+        self.flattened.set(0);
         self.heap.bytes = self.heap.bytes.saturating_add(self.held_bytes).saturating_add(self.held.now.get());
     }
 
@@ -1235,7 +1248,33 @@ impl Interp {
                 if let Some(i) = array_index(k)
                     && (i as usize) < s.len()
                 {
-                    let unit = s.flat().get(i as usize).copied().unwrap_or(0);
+                    let unit = if s.is_rope() {
+                        let owed = self.flattened_fuel.get();
+                        let flattened = self.flattened.get().saturating_add(string_bytes(s.len()));
+                        if self.heap.bytes.saturating_add(flattened) > self.limits.heap_bytes {
+                            // The copy does not fit without a collection, which needs `&mut self`:
+                            // the unit is read along the rope instead, a unit of fuel per node.
+                            let Some((unit, visited)) = s.unit_at(i as usize, self.fuel.saturating_sub(owed)) else {
+                                return Err(Abrupt::Limit(LimitKind::Fuel));
+                            };
+                            self.flattened_fuel.set(owed.saturating_add(visited));
+                            unit
+                        } else {
+                            // Flattening copies the units. The run ends here when the fuel for
+                            // the copy is not left; otherwise the next step takes it, and the
+                            // next limit check adds the copy's bytes.
+                            let fuel = owed.saturating_add((s.len() / crate::string_builtins::UNITS_PER_FUEL) as u64 + 1);
+                            if fuel > self.fuel {
+                                return Err(Abrupt::Limit(LimitKind::Fuel));
+                            }
+                            self.flattened_fuel.set(fuel);
+                            self.flattened.set(flattened);
+                            s.flat().get(i as usize).copied()
+                        }
+                    } else {
+                        s.flat().get(i as usize).copied()
+                    }
+                    .unwrap_or(0);
                     return Ok(Some(Prop::Data { value: Value::String(Str::from_units(vec![unit])), flags: ENUMERABLE }));
                 }
             }
@@ -1312,7 +1351,7 @@ impl Interp {
             if let Some(i) = array_index(k)
                 && (i as usize) < s.len()
             {
-                let unit = s.flat().get(i as usize).copied().unwrap_or(0);
+                let unit = flat_charged(self, s)?.get(i as usize).copied().unwrap_or(0);
                 return Ok(Value::String(Str::from_units(vec![unit])));
             }
         }
@@ -2420,7 +2459,7 @@ impl Interp {
         let strict = self.strict();
         self.check_source_units(s.len()).map_err(|e| self.compile_error_to_abrupt(e))?;
         self.prepare_source(s.len())?;
-        let units = s.flat();
+        let units = self.flat_reserved(s)?;
         let code = self.compile_source(&units, strict, CodeKind::Eval).map_err(|e| self.compile_error_to_abrupt(e))?;
         self.charge_code(&code, declaration_bytes(&code).saturating_add(ENV_BYTES))?;
         let frame = self.frame()?;
@@ -2585,7 +2624,7 @@ impl Interp {
             Value::Null => 0.0,
             Value::Bool(b) => f64::from(u8::from(*b)),
             Value::Number(n) => *n,
-            Value::String(s) => string_to_number(&s.flat()),
+            Value::String(s) => string_to_number(&flat_charged(self, s)?),
             Value::Object(_) => {
                 let p = self.to_primitive(v, Hint::Number)?;
                 return self.to_number(&p);
@@ -2619,19 +2658,28 @@ impl Interp {
     #[allow(clippy::wrong_self_convention)]
     pub(crate) fn to_key(&mut self, v: &Value) -> JsResult<Key> {
         match v {
-            Value::String(s) => {
-                let rope = s.is_rope();
-                let flat = s.flat();
-                if rope {
-                    self.heap.bytes += string_bytes(flat.len());
-                }
-                Ok(flat)
-            }
+            Value::String(s) => self.flat_reserved(s),
             Value::Number(n) if *n >= 0.0 && *n < 4_294_967_295.0 && n.fract() == 0.0 => Ok(index_key(*n as u32)),
-            _ => Ok(self.to_str(v)?.flat()),
+            _ => {
+                let s = self.to_str(v)?;
+                self.flat_reserved(&s)
+            }
         }
     }
 
+    /// The code units of `s`. Flattening a rope makes a copy of its units, so for a rope that
+    /// copy's bytes are reserved and added to the heap first, without taking fuel. The rope stays
+    /// rooted while the copy is reserved; everything else the caller still needs must be
+    /// reachable from the roots.
+    fn flat_reserved(&mut self, s: &Str) -> JsResult<Key> {
+        if s.is_rope() {
+            let bytes = string_bytes(s.len());
+            self.rooted(&[Value::String(s.clone())], |i| i.reserve_heap(bytes))?;
+            self.heap.bytes = self.heap.bytes.saturating_add(bytes);
+        }
+        Ok(s.flat())
+    }
+
     /// ToObject (section 9.9).
     #[allow(clippy::wrong_self_convention)]
     pub(crate) fn to_object(&mut self, v: &Value) -> JsResult<ObjId> {
@@ -2754,7 +2802,11 @@ impl Interp {
             (self.to_primitive(x, Hint::Number)?, py)
         };
         if let (Value::String(a), Value::String(b)) = (&px, &py) {
-            return Ok(Some(*a.flat() < *b.flat()));
+            // Both stay rooted, so the first copy is still counted while the second is made.
+            return self.rooted(&[px.clone(), py.clone()], |i| {
+                let (a, b) = (flat_charged(i, a)?, flat_charged(i, b)?);
+                Ok(Some(*a < *b))
+            });
         }
         let nx = self.to_number(&px)?;
         let ny = self.to_number(&py)?;
@@ -2779,7 +2831,9 @@ impl Interp {
                 | (Value::String(_), Value::String(_))
                 | (Value::Bool(_), Value::Bool(_))
                 | (Value::Object(_), Value::Object(_)) => return Ok(strict_eq(&x, &y)),
-                (Value::Number(n), Value::String(s)) | (Value::String(s), Value::Number(n)) => return Ok(*n == string_to_number(&s.flat())),
+                (Value::Number(n), Value::String(s)) | (Value::String(s), Value::Number(n)) => {
+                    return Ok(*n == string_to_number(&flat_charged(self, s)?));
+                }
                 (Value::Bool(v), _) => x = Value::Number(f64::from(u8::from(*v))),
                 (_, Value::Bool(v)) => y = Value::Number(f64::from(u8::from(*v))),
                 (Value::Number(_) | Value::String(_), Value::Object(_)) => y = self.to_primitive(&y, Hint::Default)?,

crates/js/src/realm.rs

@@ -231,6 +231,7 @@ impl Realm {
             return Err(RunError { kind: RunErrorKind::Internal, message: message.clone(), line: 0, column: 0 });
         }
         interp.fuel = interp.limits.fuel;
+        interp.flattened_fuel.set(0);
         interp.source_units = 0;
         interp.throw_site = None;
         interp.last_thrown = Value::Undefined;

crates/js/src/string_builtins.rs

@@ -109,10 +109,11 @@ fn push_held<T>(interp: &mut Interp, held: &mut HeldBytes, list: &mut Vec<T>, it
 }
 
 /// The code units of `s`. Flattening a rope makes a copy of its units, so for a rope that copy's
-/// bytes are reserved and added to the heap, and its units charged as fuel, first.
+/// bytes are reserved and added to the heap, and its units charged as fuel, first. The rope stays
+/// rooted while the copy is reserved, so what it holds is still counted.
 pub(crate) fn flat_charged(interp: &mut Interp, s: &Str) -> JsResult<Rc<[u16]>> {
     if s.is_rope() {
-        reserve_units(interp, s.len())?;
+        interp.rooted(&[Value::String(s.clone())], |interp| reserve_units(interp, s.len()))?;
     }
     Ok(s.flat())
 }

crates/js/src/text_tests.rs

@@ -309,6 +309,131 @@ fn flattening_a_rope_counts_against_the_heap() {
     }
 }
 
+#[test]
+fn operators_and_builtins_count_a_flattened_rope_against_the_heap() {
+    // Three ropes of about 524,288 units, each 1 MiB once flat, with a 3 MiB heap.
+    let build = "var s = 'x'; for (var i = 0; i < 19; i++) s = s + s; var t = s + 'y', u = s + 'z';";
+    let limits = Limits { heap_bytes: 3 << 20, ..Limits::default() };
+    let not_stopped: Vec<_> =
+        ["X < 'a'", "X == 1", "+X", "X[0]", "({})[{ toString: function () { return X; } }]", "JSON.stringify(0, null, X)", "[1].join(X)"]
+            .into_iter()
+            .filter(|read| {
+                let src = format!("{build} {}; {}; {}; 0", read.replace('X', "s"), read.replace('X', "t"), read.replace('X', "u"));
+                let result = Realm::with_options(Options::default(), limits.clone()).run(&src);
+                !matches!(result, Err(ref e) if e.kind == RunErrorKind::Limit(LimitKind::Heap))
+            })
+            .collect();
+    assert!(not_stopped.is_empty(), "not stopped by the heap limit: {not_stopped:?}");
+}
+
+#[test]
+fn a_flattened_temporary_stays_counted_while_the_next_copy_is_made() {
+    // Ropes of 262,145 units, each about half the heap once flat: one fits, two don't. Each is
+    // made by a function, so nothing but the native code holds it while it is flattened.
+    let build = "var s = 'x'; for (var i = 0; i < 18; i++) s = s + s; function f(c) { return s + c; }";
+    let limits = Limits { heap_bytes: 4 * 262_146, ..Limits::default() };
+    let run = |read: &str| Realm::with_options(Options::default(), limits.clone()).run(&format!("{build} {read}; 0"));
+    assert!(run("f('a')[0]").is_ok(), "one flattened copy fits");
+    let not_stopped: Vec<_> = [
+        "({ valueOf: function () { return f('a'); } }) < ({ valueOf: function () { return f('b'); } })",
+        "JSON.stringify({ toJSON: function () { return f('a'); } })",
+        "[1, 2].join({ toString: function () { return f('a'); } })",
+        "[{ toString: function () { return f('a'); } }].join('')",
+    ]
+    .into_iter()
+    .filter(|read| !matches!(run(read), Err(ref e) if e.kind == RunErrorKind::Limit(LimitKind::Heap)))
+    .collect();
+    assert!(not_stopped.is_empty(), "not stopped by the heap limit: {not_stopped:?}");
+}
+
+#[test]
+fn a_converted_rope_counts_what_it_holds_while_its_copy_is_made() {
+    // `mk` returns a rope whose left half is a flat string of about half the heap, held by
+    // nothing else. Its flattened copy is the other half, so the two don't fit together.
+    let build = "var s = 'x'; for (var i = 0; i < 18; i++) s = s + s; function f(c) { return s + c; } \
+                 function mk() { var a = f('a'); a[0]; return a + 'b'; }";
+    let limits = Limits { heap_bytes: 4 * 262_146, ..Limits::default() };
+    let not_stopped: Vec<_> = [
+        "+({ valueOf: mk })",
+        "({})[{ toString: mk }]",
+        "({ valueOf: mk }) == 1",
+        "++({ get x() { return mk(); } }).x",
+        "JSON.stringify(0, null, (function () { var o = new String(''); o.toString = mk; return o; })())",
+    ]
+    .into_iter()
+    .filter(|read| {
+        let result = Realm::with_options(Options::default(), limits.clone()).run(&format!("{build} {read}; 0"));
+        !matches!(result, Err(ref e) if e.kind == RunErrorKind::Limit(LimitKind::Heap))
+    })
+    .collect();
+    assert!(not_stopped.is_empty(), "not stopped by the heap limit: {not_stopped:?}");
+}
+
+#[test]
+fn a_string_object_counts_its_flattened_rope_against_the_heap() {
+    // The same three ropes, each wrapped in a String object and read by index.
+    let build = "var s = 'x'; for (var i = 0; i < 19; i++) s = s + s; var t = s + 'y', u = s + 'z'; \
+                 String.prototype.first = function () { return this[0]; };";
+    let limits = Limits { heap_bytes: 3 << 20, ..Limits::default() };
+    // Wrapping alone reads no units, so it makes no copy.
+    let lengths =
+        Realm::with_options(Options::default(), limits.clone()).run(&format!("{build} Object(s).length + new String(t).length + Object(u).length"));
+    assert_eq!(lengths, Ok(JsValue::Number(1_572_866.0)));
+    // Two copies fit. The third read finds no room for its copy and reads the unit along the
+    // rope, so the live heap stays within the limit.
+    for (read, expected) in [
+        ("Object(X)[0]", "xxx"),
+        ("new String(X)[0]", "xxx"),
+        ("0 in Object(X)", "truetruetrue"),
+        ("Object(X).hasOwnProperty(0)", "truetruetrue"),
+        ("X.first()", "xxx"),
+        ("Object(X)[524288]", "undefinedyz"),
+    ] {
+        let src = format!("{build} '' + ({}) + ({}) + ({})", read.replace('X', "s"), read.replace('X', "t"), read.replace('X', "u"));
+        let mut realm = Realm::with_options(Options::default(), limits.clone());
+        assert_eq!(realm.run(&src), Ok(JsValue::String(expected.into())), "{read}");
+        let interp = realm.interp();
+        interp.collect_garbage();
+        assert!(interp.heap.bytes <= limits.heap_bytes, "{read}: {} bytes live", interp.heap.bytes);
+    }
+    let wrapped =
+        Realm::with_options(Options::default(), Limits::default()).run(&format!("{build} Object(t)[524288] + new String(u)[524288] + t.first()"));
+    assert_eq!(wrapped, Ok(JsValue::String("yzx".into())));
+}
+
+#[test]
+fn a_string_object_takes_fuel_for_the_rope_it_flattens() {
+    let limits = Limits { fuel: 40_000, ..Limits::default() };
+    let build = "var s = ' '; for (var i = 0; i < 20; i++) s = s + s; Object(s).length";
+    assert!(Realm::with_options(Options::default(), limits.clone()).run(build).is_ok());
+    for read in ["Object(s)[0]", "0 in new String(s)"] {
+        let e = fail_with(Realm::with_options(Options::default(), limits.clone()), &format!("{build}; {read}"));
+        assert_eq!(e.kind, RunErrorKind::Limit(LimitKind::Fuel), "{read}: {e:?}");
+    }
+    // A run that ends with an error right after the read leaves the next run its whole fuel.
+    let mut realm = Realm::with_options(Options::default(), Limits { fuel: 100_000, ..Limits::default() });
+    let defined = realm.run("var s = ' '; for (var i = 0; i < 20; i++) s = s + s; Object.defineProperty(new String(s), 0, { value: 'a' })");
+    assert!(matches!(defined, Err(ref e) if e.kind == RunErrorKind::Thrown), "{defined:?}");
+    assert!(realm.run("for (var j = 0; j < 5000; j++) {} j").is_ok());
+}
+
+#[test]
+fn a_call_keeps_its_operands_while_a_string_receiver_is_boxed() {
+    // `g` leaves a garbage string, so the heap is close to its limit when the rope receiver of a
+    // non-strict function is boxed. The callee and the argument object are off the stack by
+    // then, and must survive anything boxing the receiver does.
+    let build = "var s = 'x'; for (var i = 0; i < 17; i++) s = s + s; function f(c) { return s + c; } \
+                 function mk() { var a = f('a'); a[0]; return a + 'b'; } 0";
+    let mut realm = Realm::with_options(Options::default(), Limits::default());
+    assert_eq!(realm.run(build), Ok(JsValue::Number(0.0)));
+    let interp = realm.interp();
+    interp.collect_garbage();
+    interp.limits.heap_bytes = interp.heap.bytes + 2 * crate::heap::string_bytes(131_073) + (64 << 10);
+    interp.reset_gc_threshold();
+    let src = "(function () { var g = f('g'); g[0]; })(); (function (o) { return o.x; }).call(mk(), { x: 1 })";
+    assert_eq!(realm.run(src), Ok(JsValue::Number(1.0)));
+}
+
 #[test]
 fn compiling_a_pattern_counts_against_the_heap() {
     let src = "var s = '[\\\\s\\\\S]a*'; for (var i = 0; i < 14; i++) s = s + s; new RegExp(s).source.length";

crates/js/src/value.rs

@@ -163,6 +163,34 @@ impl Str {
         flat
     }
 
+    /// The code unit at `i` read along the rope without flattening it, visiting at most
+    /// `max_nodes` nodes: `None` when that is not enough, otherwise the unit (when `i` is in
+    /// range) and the number of nodes visited.
+    pub(crate) fn unit_at(&self, mut i: usize, max_nodes: u64) -> Option<(Option<u16>, u64)> {
+        let mut node = self.clone();
+        let mut visited = 0u64;
+        loop {
+            visited += 1;
+            if visited > max_nodes {
+                return None;
+            }
+            let next = match &*node.0.borrow() {
+                Node::Flat(f) => return Some((f.get(i).copied(), visited)),
+                Node::Concat { left, right, .. } => {
+                    let left_len = left.len();
+                    if i < left_len {
+                        left.clone()
+                    } else {
+                        i -= left_len;
+                        right.clone()
+                    }
+                }
+                Node::Taken => return Some((None, visited)),
+            };
+            node = next;
+        }
+    }
+
     pub(crate) fn to_rust_string(&self) -> String {
         String::from_utf16_lossy(&self.flat())
     }