Scripts: a flattened string stays counted while the next copy is made

- the relational operators root both strings while flattening them, so the first copy still counts when the second one starts a collection
- join roots its separator for the whole join and each element while it is written, and JSON.stringify roots a string while it writes it
- a String object flattens its string when it is made, with the copy counted, so reading one of its characters copies nothing

crates/js/src/builtins/array.rs

@@ -279,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 = 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
+    // 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> {
@@ -302,8 +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;
-        let units = flat_charged(interp, &s)?;
-        out.push_units(interp, &units)?;
+        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))

crates/js/src/builtins/json.rs

@@ -633,10 +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) => {
-            let units = flat_charged(interp, s)?;
-            write_quoted(interp, out, &units)
-        }
+        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

@@ -2626,7 +2626,11 @@ impl Interp {
             Value::Undefined | Value::Null => return Err(self.type_error("cannot convert undefined or null to an object")),
             Value::Bool(b) => (Class::Boolean(*b), self.intrinsics.boolean_proto),
             Value::Number(n) => (Class::Number(*n), self.intrinsics.number_proto),
-            Value::String(s) => (Class::String(s.clone()), self.intrinsics.string_proto),
+            // Flattened now, with its copy counted, so reading a character copies nothing.
+            Value::String(s) => {
+                flat_charged(self, s)?;
+                (Class::String(s.clone()), self.intrinsics.string_proto)
+            }
         };
         Ok(self.heap.alloc_object(Object::new(class, Some(proto))))
     }
@@ -2740,8 +2744,11 @@ impl Interp {
             (self.to_primitive(x, Hint::Number)?, py)
         };
         if let (Value::String(a), Value::String(b)) = (&px, &py) {
-            let (a, b) = (flat_charged(self, a)?, flat_charged(self, b)?);
-            return Ok(Some(*a < *b));
+            // 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)?;

crates/js/src/string_builtins.rs

@@ -214,6 +214,8 @@ fn string_ctor(interp: &mut Interp, _: Value, args: &[Value], construct: bool) -
     if !construct {
         return Ok(Value::String(s));
     }
+    // Flattened now, with its copy counted, so reading a character copies nothing.
+    flat_charged(interp, &s)?;
     let proto = interp.intrinsics.string_proto;
     Ok(Value::Object(interp.heap.alloc_object(Object::new(Class::String(s), Some(proto)))))
 }

crates/js/src/text_tests.rs

@@ -314,15 +314,43 @@ 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();
+    let not_stopped: Vec<_> = [
+        "X < 'a'",
+        "X == 1",
+        "+X",
+        "X[0]",
+        "new String(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:?}");
 }