JavaScript: sort keys and replacer names count toward the heap limit

Array.prototype.sort without a comparison function and the property list
of a JSON.stringify replacer array keep a string for each element. A key
that comes from flattening a concatenation is a new buffer: it is now
charged to the heap, and keys that are not elements are rooted so that a
collection counts them. Keys that share an existing string's buffer are
not charged again.

crates/js/src/builtins/array.rs

@@ -641,9 +641,22 @@ fn sort_elements(interp: &mut Interp, o: ObjId, len: u64, compare: &Value) -> Js
         let mut keyed = Vec::with_capacity(values.len());
         for v in values {
             interp.use_fuel()?;
-            let s = interp.to_str(&v)?.flat();
-            interp.reserve_heap(crate::heap::string_bytes(s.len()))?;
-            keyed.push((s, v));
+            let s = interp.to_str(&v)?;
+            // Flattening a concatenation makes a new buffer, so it is charged. A key that is
+            // not an element is rooted so that collections count it.
+            let rope = s.buffer_addr().is_none();
+            let bytes = crate::heap::string_bytes(s.len());
+            if rope {
+                interp.reserve_heap(bytes)?;
+            }
+            let units = s.flat();
+            if rope {
+                interp.heap.bytes = interp.heap.bytes.saturating_add(bytes);
+            }
+            if !matches!(v, Value::String(_)) {
+                interp.temp_roots.push(Value::String(s));
+            }
+            keyed.push((units, v));
         }
         merge_sort(interp, &mut keyed, |_, a, b| Ok(a.0 > b.0))?;
         values = keyed.into_iter().map(|(_, v)| v).collect();

crates/js/src/builtins/json.rs

@@ -487,17 +487,27 @@ fn property_list(interp: &mut Interp, r: ObjId) -> JsResult<Vec<Key>> {
         interp.use_fuel()?;
         let v = interp.get_from(r, &key_of(f64::from(i)), &Value::Object(r))?;
         let item = match &v {
-            Value::String(s) => Some(s.flat()),
-            Value::Number(n) => Some(number_to_str(*n).flat()),
+            Value::String(s) => Some(s.clone()),
+            Value::Number(_) => Some(interp.primitive_to_str(&v)),
             Value::Object(o) if interp.heap.object(*o).is_some_and(|x| matches!(x.class, Class::String(_) | Class::Number)) => {
-                Some(interp.to_str(&v)?.flat())
+                Some(interp.to_str(&v)?)
             }
             _ => None,
         };
-        if let Some(k) = item
-            && seen.insert(k.clone())
-        {
-            interp.reserve_heap(string_bytes(k.len()))?;
+        let Some(s) = item else { continue };
+        // Flattening a concatenation makes a new buffer, so it is charged.
+        let rope = s.buffer_addr().is_none();
+        let bytes = string_bytes(s.len());
+        if rope {
+            interp.reserve_heap(bytes)?;
+        }
+        let k = s.flat();
+        if seen.insert(k.clone()) {
+            if rope {
+                interp.heap.bytes = interp.heap.bytes.saturating_add(bytes);
+            }
+            // The list lives until the call returns, so its names are rooted.
+            interp.temp_roots.push(Value::String(s));
             list.push(k);
         }
     }

crates/js/src/builtins_tests.rs

@@ -278,6 +278,32 @@ fn built_ins_end_with_a_limit_instead_of_running_on() {
     assert_eq!(limit(&mut realm, "var s = 'x'; for (var i = 0; i < 20; i++) s += s; JSON.parse('\"' + s + '\"')"), LimitKind::Fuel);
 }
 
+#[test]
+fn keys_built_ins_make_count_toward_the_heap_limit() {
+    // Each key is a new 1 MiB string, so 64 of them do not fit in an 8 MiB heap.
+    let limits = Limits { heap_bytes: 8 << 20, ..Limits::default() };
+    let big = "var s = 'x'; for (var i = 0; i < 19; i++) s += s; ";
+    let sort = format!("{big} var a = []; for (var i = 0; i < 64; i++) a[i] = {{n: i, toString: function () {{ return s + this.n; }}}}; a.sort();");
+    assert_eq!(limit(&mut small(limits.clone()), &sort), LimitKind::Heap);
+    let list = format!(
+        "{big} var a = [], c = 0; for (var i = 0; i < 64; i++) Object.defineProperty(a, i, {{get: function () {{ return s + c++; }}}}); JSON.stringify({{}}, a);"
+    );
+    assert_eq!(limit(&mut small(limits), &list), LimitKind::Heap);
+    // Keys that share a string's storage are not charged again: with half the string's size to
+    // spare, sorting objects whose keys are that string still fits.
+    let setup = format!(
+        "{big} var keep = [s + 0, s + 1]; keep.sort(); var a = []; for (var i = 0; i < 64; i++) a[i] = s; a.sort(); \
+         var b = []; for (var i = 0; i < 64; i++) b[i] = {{toString: function () {{ return s; }}}}; 0"
+    );
+    let mut measure = small(Limits::default());
+    assert_eq!(measure.run(&setup), Ok(JsValue::Number(0.0)));
+    measure.interp().collect_garbage();
+    let live = measure.interp().heap.bytes;
+    let mut realm = small(Limits { heap_bytes: live + (512 << 10), ..Limits::default() });
+    assert_eq!(realm.run(&setup), Ok(JsValue::Number(0.0)));
+    assert_eq!(realm.run("b.sort(); JSON.stringify({}, a); JSON.stringify({}, [s, s, s]); b.length"), Ok(JsValue::Number(64.0)));
+}
+
 #[test]
 fn values_built_ins_hold_survive_collections() {
     // A small heap makes the callbacks' garbage start collections while the results are built.