JavaScript: kept code and function source text count toward the heap

A collection counts each piece of code a closure or frame keeps: its instructions, tables,
names, nested functions and its source text, each shared buffer once. Code that eval and the
Function constructor compile at run time is charged when it is made, after a heap reservation.
Function.prototype.toString reserves and charges the text it copies.

crates/js/src/builtins.rs

@@ -7,7 +7,7 @@ use std::rc::Rc;
 
 use crate::bytecode::{CodeKind, Source};
 use crate::compiler::{CompileFail, compile_constructed_function};
-use crate::heap::{Class, Func, HIDDEN, NativeFn, NativeKind, ObjId, Object, Prop, key};
+use crate::heap::{Class, Func, HIDDEN, NativeFn, NativeKind, ObjId, Object, Prop, key, string_bytes};
 use crate::interp::{Abrupt, ErrorType, Interp, JsResult};
 use crate::value::{Str, Value};
 
@@ -255,6 +255,7 @@ fn function_ctor(interp: &mut Interp, _: Value, args: &[Value], _: bool) -> JsRe
         Err(CompileFail::Syntax(e)) => return Err(interp.throw(ErrorType::Syntax, &e.message)),
         Err(CompileFail::Internal(m)) => return Err(Abrupt::Internal(m)),
     };
+    interp.charge_code(&code, 0)?;
     let env = interp.global_env;
     interp.make_closure(code, env).map(Value::Object)
 }
@@ -269,7 +270,11 @@ fn function_to_string(interp: &mut Interp, this: Value, _: &[Value], _: bool) ->
     let text = match func {
         Some(Func::Closure { code, .. }) => {
             let (start, end) = (code.span.0 as usize, code.span.1 as usize);
-            String::from_utf16_lossy(code.source.text.get(start..end).unwrap_or_default())
+            let units = code.source.text.get(start..end).unwrap_or_default();
+            let bytes = string_bytes(units.len());
+            interp.reserve_heap(bytes)?;
+            interp.heap.bytes = interp.heap.bytes.saturating_add(bytes);
+            return Ok(Value::String(Str::from_units(units.to_vec())));
         }
         Some(Func::Native { name, .. }) => format!("function {name}() {{ [native code] }}"),
         Some(Func::Host { name, .. }) => format!("function {name}() {{ [native code] }}"),

crates/js/src/heap.rs

@@ -10,7 +10,7 @@
 use std::collections::{HashMap, HashSet};
 use std::rc::Rc;
 
-use crate::bytecode::Code;
+use crate::bytecode::{Code, Handler, Op};
 use crate::value::{Str, Value};
 
 /// A property name: its UTF-16 code units.
@@ -289,6 +289,8 @@ pub(crate) const OBJECT_BYTES: usize = 96;
 pub(crate) const ENV_BYTES: usize = 64;
 /// The estimated bytes of a rope node.
 pub(crate) const ROPE_BYTES: usize = 48;
+/// The estimated bytes of a piece of code without its instructions, tables and strings.
+const CODE_BYTES: usize = 384;
 
 pub(crate) fn string_bytes(len: usize) -> usize {
     32 + 2 * len
@@ -455,9 +457,56 @@ fn mark_value(v: &Value, work: &mut Vec<u32>, strings: &mut StringSizes) {
     }
 }
 
+/// The estimated bytes of `code` with the functions nested in it and its source text, for
+/// charging code made at run time to the heap before anything refers to it.
+pub(crate) fn code_bytes(code: &Code) -> usize {
+    let mut sizes = StringSizes::default();
+    mark_code(code, &mut sizes);
+    sizes.bytes
+}
+
+/// Adds the estimated size of `code` and the functions nested in it, counting each piece of code,
+/// each source text and each string once however many closures share them.
 fn mark_code(code: &Code, strings: &mut StringSizes) {
-    for s in &code.strings {
-        strings.add(s);
+    let mut stack = vec![code];
+    while let Some(code) = stack.pop() {
+        if !strings.seen.insert(code as *const Code as usize) {
+            continue;
+        }
+        if strings.seen.insert(Rc::as_ptr(&code.source) as usize) {
+            strings.bytes += code.source.lines.heap_bytes();
+            // The text can share its buffer with the string eval was given.
+            if strings.seen.insert(Rc::as_ptr(&code.source.text) as *const () as usize) {
+                strings.bytes += string_bytes(code.source.text.len());
+            }
+        }
+        let handles = code.strings.len()
+            + 2 * code.regexps.len()
+            + code.names.len()
+            + code.functions.len()
+            + code.params.len()
+            + code.vars.len()
+            + 2 * code.fun_decls.len();
+        let tables = code.finally_tables.iter().map(|t| 24 + 4 * t.len()).sum::<usize>();
+        strings.bytes += CODE_BYTES
+            + size_of::<Op>() * code.ops.len()
+            + 4 * code.positions.len()
+            + 8 * code.numbers.len()
+            + size_of::<Handler>() * code.handlers.len()
+            + 16 * handles
+            + tables;
+        for s in &code.strings {
+            strings.add(s);
+        }
+        for (pattern, flags) in &code.regexps {
+            strings.add(pattern);
+            strings.add(flags);
+        }
+        let keys = code.names.iter().chain(&code.params).chain(&code.vars).chain(code.fun_decls.iter().map(|(k, _)| k));
+        for k in keys.chain(&code.self_name) {
+            strings.add_key(k);
+        }
+        stack.extend(code.functions.iter().map(|f| &**f));
     }
 }
 

crates/js/src/interp.rs

@@ -24,7 +24,7 @@ use crate::bytecode::{Code, CodeKind, Op, Source};
 use crate::compiler::{CompileFail, compile_program};
 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, key, string_bytes,
+    Object, PROP_BYTES, Prop, PropMap, ROPE_BYTES, Roots, WRITABLE, code_bytes, key, string_bytes,
 };
 use crate::value::{Str, Value, array_index, index_key, number_to_str, string_to_number, to_int32, to_uint32};
 use crate::{LimitKind, Limits, Options};
@@ -400,7 +400,7 @@ impl Interp {
 
     /// Ends the run with the heap limit when `bytes` more would not fit in it after a
     /// collection. Everything the caller still needs must be reachable from the roots.
-    fn reserve_heap(&mut self, bytes: usize) -> JsResult<()> {
+    pub(crate) fn reserve_heap(&mut self, bytes: usize) -> JsResult<()> {
         if self.heap.bytes.saturating_add(bytes) <= self.limits.heap_bytes {
             return Ok(());
         }
@@ -412,6 +412,16 @@ impl Interp {
         Ok(())
     }
 
+    /// Charges code compiled at run time to the heap, reserving its estimated size and `extra`
+    /// more first. The code is not reachable from the roots yet, so a collection here leaves it
+    /// out of the live total and the charge adds it.
+    pub(crate) fn charge_code(&mut self, code: &Code, extra: usize) -> JsResult<()> {
+        let bytes = code_bytes(code);
+        self.reserve_heap(bytes.saturating_add(extra))?;
+        self.heap.bytes = self.heap.bytes.saturating_add(bytes);
+        Ok(())
+    }
+
     pub(crate) fn collect_garbage(&mut self) {
         let mut roots = Roots::default();
         roots.values.extend(self.stack.iter());
@@ -1769,7 +1779,7 @@ impl Interp {
         self.check_source_units(s.len()).map_err(|e| self.compile_error_to_abrupt(e))?;
         let units = s.flat();
         let code = self.compile_source(&units, strict, CodeKind::Eval).map_err(|e| self.compile_error_to_abrupt(e))?;
-        self.reserve_heap(declaration_bytes(&code).saturating_add(ENV_BYTES))?;
+        self.charge_code(&code, declaration_bytes(&code).saturating_add(ENV_BYTES))?;
         let frame = self.frame()?;
         let (caller_env, caller_var_env, this) = (frame.env, frame.var_env, frame.this.clone());
         let (env, var_env) = if code.strict {
@@ -1789,7 +1799,7 @@ impl Interp {
         if self.reentry >= self.limits.reentry {
             return Err(Abrupt::Limit(LimitKind::Reentry));
         }
-        self.reserve_heap(declaration_bytes(&code).saturating_add(ENV_BYTES))?;
+        self.charge_code(&code, declaration_bytes(&code).saturating_add(ENV_BYTES))?;
         let (env, global) = (self.global_env, self.global);
         let (env, var_env) = if code.strict && is_eval {
             let e = self.heap.alloc_env(Env::Decl { bindings: PropMap::default(), outer: Some(env) });

crates/js/src/interp_tests.rs

@@ -285,8 +285,9 @@ fn names_from_source_count_toward_the_heap_when_they_are_bound() {
     // Without a collection in between, the estimate grows by at least each new name's code units:
     // a global variable (an object property), a binding made by direct eval, and a parameter of a
     // function made by the Function constructor.
-    let mut realm = small(Limits::default());
-    let name = "n".repeat(100_000);
+    // Each case runs in a new realm, so the charges for the code itself do not add up to a
+    // collection.
+    let name = "n".repeat(50_000);
     let cases = [
         format!("var {name} = 1;"),
         format!("(function () {{ eval('var {name}2 = 1; {name}2'); }})();"),
@@ -294,12 +295,13 @@ fn names_from_source_count_toward_the_heap_when_they_are_bound() {
     ];
     let mut short = Vec::new();
     for (i, src) in cases.iter().enumerate() {
+        let mut realm = small(Limits::default());
         let collections = realm.interp().heap.collections;
         let before = realm.interp().heap.bytes;
         realm.run(src).unwrap();
         assert_eq!(realm.interp().heap.collections, collections, "a collection ran in case {i}");
         let grown = realm.interp().heap.bytes - before;
-        if grown < crate::heap::string_bytes(100_000) {
+        if grown < crate::heap::string_bytes(50_000) {
             short.push((i, grown));
         }
     }
@@ -316,6 +318,40 @@ fn strings_from_host_functions_count_toward_the_heap() {
     assert_eq!(limit(&mut with_big(limits), "big(1000000).length"), LimitKind::Heap);
 }
 
+#[test]
+fn functions_kept_across_runs_count_their_code_toward_the_heap() {
+    // Each run keeps a function whose source holds a 64 Ki-unit comment, and drops the string the
+    // source was made from. The kept code fills the heap after a few runs.
+    let limits = Limits { heap_bytes: 1 << 20, ..Limits::default() };
+    let mut realm = small(limits);
+    assert_eq!(realm.run("var keep = []; keep.length"), Ok(JsValue::Number(0.0)));
+    let step = "var s = 'x'; for (var i = 0; i < 16; i++) s += s; keep[keep.length] = Function('/*' + s + '*/'); s = null; keep.length";
+    let mut ended = None;
+    for run in 1..=20 {
+        match realm.run(step) {
+            Ok(v) => assert_eq!(v, JsValue::Number(f64::from(run))),
+            Err(RunError { kind: RunErrorKind::Limit(k), .. }) => {
+                ended = Some((run, k));
+                break;
+            }
+            Err(e) => panic!("run {run}: {e:?}"),
+        }
+    }
+    let Some((run, kind)) = ended else { panic!("20 kept functions fit in a 1 MiB heap") };
+    assert_eq!(kind, LimitKind::Heap);
+    assert!(run > 2, "the heap filled on run {run}");
+}
+
+#[test]
+fn function_source_text_counts_toward_the_heap() {
+    let limits = Limits { heap_bytes: 1 << 20, ..Limits::default() };
+    let comment = "x".repeat(300_000);
+    let src = format!("var f = function () {{ /*{comment}*/ }}; f.toString().length");
+    assert_eq!(limit(&mut small(limits.clone()), &src), LimitKind::Heap);
+    let mut realm = small(limits);
+    assert_eq!(realm.run("(function f(a) { return a; }).toString()"), Ok(JsValue::String("function f(a) { return a; }".into())));
+}
+
 #[test]
 fn declarations_check_the_heap_before_they_are_bound() {
     // Each script declares more than fits in the heap. The run ends with the heap limit before

crates/js/src/lib.rs

@@ -155,6 +155,11 @@ impl LineTable {
         (u32::try_from(line).unwrap_or(u32::MAX), offset.saturating_sub(start) + 1)
     }
 
+    /// The estimated bytes the table holds.
+    pub(crate) fn heap_bytes(&self) -> usize {
+        32 + 4 * self.starts.len()
+    }
+
     pub(crate) fn line_count(&self) -> u32 {
         if self.empty {
             return 0;