perf(vm): Reduce suspended Vm memory footprint

aapoalas · aapoalas · commit 7f722a737a27 · 2024-10-11T15:33:29.000+03:00
diff --git a/nova_vm/src/ecmascript/builtins/control_abstraction_objects/async_function_objects/await_reaction.rs b/nova_vm/src/ecmascript/builtins/control_abstraction_objects/async_function_objects/await_reaction.rs
@@ -18,11 +18,12 @@ use crate::{
                 promise_prototype::inner_promise_then,
             },
             promise::Promise,
+            ECMAScriptFunction,
         },
         execution::{Agent, ExecutionContext},
         types::Value,
     },
-    engine::{Executable, ExecutionResult, Vm},
+    engine::{ExecutionResult, SuspendedVm},
     heap::{CompactionLists, CreateHeapData, Heap, HeapMarkAndSweep, WorkQueues},
 };
 
@@ -66,10 +67,12 @@ impl AwaitReactionIdentifier {
         // 3. d. Resume the suspended evaluation of asyncContext using NormalCompletion(v) as the result of the operation that suspended it.
         // 5. d. Resume the suspended evaluation of asyncContext using ThrowCompletion(reason) as the result of the operation that suspended it.
         let vm = agent[self].vm.take().unwrap();
-        let executable = agent[self].executable.take().unwrap();
+        let async_function = agent[self].async_function.unwrap();
+        // SAFETY: We keep the async function alive.
+        let executable = unsafe { agent[async_function].compiled_bytecode.unwrap().as_ref() };
         let execution_result = match reaction_type {
-            PromiseReactionType::Fulfill => vm.resume(agent, &executable, value),
-            PromiseReactionType::Reject => vm.resume_throw(agent, &executable, value),
+            PromiseReactionType::Fulfill => vm.resume(agent, executable, value),
+            PromiseReactionType::Reject => vm.resume_throw(agent, executable, value),
         };
 
         match execution_result {
@@ -97,7 +100,6 @@ impl AwaitReactionIdentifier {
                 // [27.7.5.3 Await ( value )](https://tc39.es/ecma262/#await)
                 // 8. Remove asyncContext from the execution context stack and restore the execution context that is at the top of the execution context stack as the running execution context.
                 agent[self].vm = Some(vm);
-                agent[self].executable = Some(executable);
                 agent[self].execution_context = Some(agent.execution_context_stack.pop().unwrap());
 
                 // `handler` corresponds to the `fulfilledClosure` and `rejectedClosure` functions,
@@ -162,8 +164,8 @@ impl HeapMarkAndSweep for AwaitReactionIdentifier {
 
 #[derive(Debug)]
 pub(crate) struct AwaitReaction {
-    pub(crate) vm: Option<Vm>,
-    pub(crate) executable: Option<Executable>,
+    pub(crate) vm: Option<SuspendedVm>,
+    pub(crate) async_function: Option<ECMAScriptFunction>,
     pub(crate) execution_context: Option<ExecutionContext>,
     pub(crate) return_promise_capability: PromiseCapability,
 }
@@ -178,13 +180,13 @@ impl CreateHeapData<AwaitReaction, AwaitReactionIdentifier> for Heap {
 impl HeapMarkAndSweep for AwaitReaction {
     fn mark_values(&self, queues: &mut WorkQueues) {
         self.vm.mark_values(queues);
-        self.executable.mark_values(queues);
+        self.async_function.mark_values(queues);
         self.return_promise_capability.mark_values(queues);
     }
 
     fn sweep_values(&mut self, compactions: &CompactionLists) {
         self.vm.sweep_values(compactions);
-        self.executable.sweep_values(compactions);
+        self.async_function.sweep_values(compactions);
         self.return_promise_capability.sweep_values(compactions);
     }
 }
diff --git a/nova_vm/src/ecmascript/builtins/control_abstraction_objects/generator_objects.rs b/nova_vm/src/ecmascript/builtins/control_abstraction_objects/generator_objects.rs
@@ -15,7 +15,7 @@ use crate::{
             InternalMethods, InternalSlots, IntoObject, IntoValue, Object, OrdinaryObject, Value,
         },
     },
-    engine::{Executable, ExecutionResult, Vm},
+    engine::{Executable, ExecutionResult, SuspendedVm, Vm},
     heap::{
         indexes::{BaseIndex, GeneratorIndex},
         CompactionLists, CreateHeapData, Heap, HeapMarkAndSweep, WorkQueues,
@@ -337,7 +337,7 @@ pub struct GeneratorHeapData {
 
 #[derive(Debug)]
 pub(crate) enum VmOrArguments {
-    Vm(Vm),
+    Vm(SuspendedVm),
     Arguments(Box<[Value]>),
 }
 
diff --git a/nova_vm/src/ecmascript/syntax_directed_operations/function_definitions.rs b/nova_vm/src/ecmascript/syntax_directed_operations/function_definitions.rs
@@ -300,8 +300,18 @@ pub(crate) fn evaluate_async_function_body(
     //} else {
     // 4. Else,
     // a. Perform AsyncFunctionStart(promiseCapability, FunctionBody).
-    let data = CompileFunctionBodyData::new(agent, function_object);
-    let exe = Executable::compile_function_body(agent, data);
+    let exe = if let Some(exe) = agent[function_object].compiled_bytecode {
+        // SAFETY: exe is a non-null pointer pointing to an initialized
+        // Executable that cannot have a live mutable reference to it.
+        unsafe { exe.as_ref() }
+    } else {
+        let data = CompileFunctionBodyData::new(agent, function_object);
+        let exe = Box::new(Executable::compile_function_body(agent, data));
+        let exe = NonNull::from(Box::leak(exe));
+        agent[function_object].compiled_bytecode = Some(exe);
+        // SAFETY: Same as above, only more.
+        unsafe { exe.as_ref() }
+    };
 
     // AsyncFunctionStart will run the function until it returns, throws or gets suspended with
     // an await.
@@ -329,7 +339,7 @@ pub(crate) fn evaluate_async_function_body(
             // cloning it would mess up the execution context stack.
             let handler = PromiseReactionHandler::Await(agent.heap.create(AwaitReaction {
                 vm: Some(vm),
-                executable: Some(exe),
+                async_function: Some(function_object),
                 execution_context: Some(agent.running_execution_context().clone()),
                 return_promise_capability: promise_capability,
             }));
diff --git a/nova_vm/src/engine/bytecode.rs b/nova_vm/src/engine/bytecode.rs
@@ -12,4 +12,4 @@ pub(crate) use executable::{
     NamedEvaluationParameter, SendableRef,
 };
 pub(crate) use instructions::{Instruction, InstructionIter};
-pub(crate) use vm::{instanceof_operator, ExecutionResult, Vm};
+pub(crate) use vm::{instanceof_operator, ExecutionResult, SuspendedVm, Vm};
diff --git a/nova_vm/src/engine/bytecode/vm.rs b/nova_vm/src/engine/bytecode/vm.rs
@@ -64,8 +64,14 @@ unsafe impl Sync for EmptyParametersList {}
 pub(crate) enum ExecutionResult {
     Return(Value),
     Throw(JsError),
-    Await { vm: Vm, awaited_value: Value },
-    Yield { vm: Vm, yielded_value: Value },
+    Await {
+        vm: SuspendedVm,
+        awaited_value: Value,
+    },
+    Yield {
+        vm: SuspendedVm,
+        yielded_value: Value,
+    },
 }
 impl ExecutionResult {
     pub(crate) fn into_js_result(self) -> JsResult<Value> {
@@ -113,6 +119,55 @@ pub(crate) struct Vm {
     reference: Option<Reference>,
 }
 
+#[derive(Debug)]
+pub(crate) struct SuspendedVm {
+    ip: usize,
+    /// Note: Stack is non-empty only if the code awaits inside a call
+    /// expression. This is reasonably rare that we can expect the stack to
+    /// usually be empty. In this case this Box is an empty dangling pointer
+    /// and no heap data clone is required.
+    stack: Box<[Value]>,
+    /// Note: Reference stack is non-empty only if the code awaits inside a
+    /// call expression. This means that usually no heap data clone is
+    /// required.
+    reference_stack: Box<[Reference]>,
+    /// Note: Iterator stack is non-empty only if the code awaits inside a
+    /// for-in or for-of loop. This means that often no heap data clone is
+    /// required.
+    iterator_stack: Box<[VmIterator]>,
+    /// Note: Exception jump stack is non-empty only if the code awaits inside
+    /// a try block. This means that often no heap data clone is required.
+    exception_jump_target_stack: Box<[ExceptionJumpTarget]>,
+}
+
+impl SuspendedVm {
+    pub(crate) fn resume(
+        self,
+        agent: &mut Agent,
+        executable: &Executable,
+        value: Value,
+    ) -> ExecutionResult {
+        let vm = Vm::from_suspended(self);
+        vm.resume(agent, executable, value)
+    }
+
+    pub(crate) fn resume_throw(
+        self,
+        agent: &mut Agent,
+        executable: &Executable,
+        err: Value,
+    ) -> ExecutionResult {
+        // Optimisation: Avoid unsuspending the Vm if we're just going to throw
+        // out of it immediately.
+        if self.exception_jump_target_stack.is_empty() {
+            let err = JsError::new(err);
+            return ExecutionResult::Throw(err);
+        }
+        let vm = Vm::from_suspended(self);
+        vm.resume_throw(agent, executable, err)
+    }
+}
+
 impl Vm {
     fn new() -> Self {
         Self {
@@ -127,6 +182,29 @@ impl Vm {
         }
     }
 
+    fn suspend(self) -> SuspendedVm {
+        SuspendedVm {
+            ip: self.ip,
+            stack: self.stack.into_boxed_slice(),
+            reference_stack: self.reference_stack.into_boxed_slice(),
+            iterator_stack: self.iterator_stack.into_boxed_slice(),
+            exception_jump_target_stack: self.exception_jump_target_stack.into_boxed_slice(),
+        }
+    }
+
+    fn from_suspended(suspended: SuspendedVm) -> Self {
+        Self {
+            ip: suspended.ip,
+            stack: suspended.stack.into_vec(),
+            reference_stack: suspended.reference_stack.into_vec(),
+            iterator_stack: suspended.iterator_stack.into_vec(),
+            exception_jump_target_stack: suspended.exception_jump_target_stack.into_vec(),
+            result: None,
+            exception: None,
+            reference: None,
+        }
+    }
+
     fn fetch_identifier(&self, exe: &Executable, index: usize) -> String {
         String::try_from(exe.constants[index])
             .expect("Invalid identifier index: Value was not a String")
@@ -183,7 +261,7 @@ impl Vm {
         vm.inner_execute(agent, executable)
     }
 
-    pub(crate) fn resume(
+    pub fn resume(
         mut self,
         agent: &mut Agent,
         executable: &Executable,
@@ -193,7 +271,7 @@ impl Vm {
         self.inner_execute(agent, executable)
     }
 
-    pub(crate) fn resume_throw(
+    pub fn resume_throw(
         mut self,
         agent: &mut Agent,
         executable: &Executable,
@@ -217,14 +295,14 @@ impl Vm {
                 Ok(ContinuationKind::Yield) => {
                     let yielded_value = self.result.take().unwrap();
                     return ExecutionResult::Yield {
-                        vm: self,
+                        vm: self.suspend(),
                         yielded_value,
                     };
                 }
                 Ok(ContinuationKind::Await) => {
                     let awaited_value = self.result.take().unwrap();
                     return ExecutionResult::Await {
-                        vm: self,
+                        vm: self.suspend(),
                         awaited_value,
                     };
                 }
@@ -2205,30 +2283,18 @@ impl HeapMarkAndSweep for ExceptionJumpTarget {
     }
 }
 
-impl HeapMarkAndSweep for Vm {
+impl HeapMarkAndSweep for SuspendedVm {
     fn mark_values(&self, queues: &mut WorkQueues) {
-        self.stack.as_slice().mark_values(queues);
-        self.reference_stack.as_slice().mark_values(queues);
-        self.iterator_stack.as_slice().mark_values(queues);
-        self.exception_jump_target_stack
-            .as_slice()
-            .mark_values(queues);
-        self.result.mark_values(queues);
-        self.exception.mark_values(queues);
-        self.reference.mark_values(queues);
+        self.stack.mark_values(queues);
+        self.reference_stack.mark_values(queues);
+        self.iterator_stack.mark_values(queues);
+        self.exception_jump_target_stack.mark_values(queues);
     }
 
     fn sweep_values(&mut self, compactions: &CompactionLists) {
-        self.stack.as_mut_slice().sweep_values(compactions);
-        self.reference_stack
-            .as_mut_slice()
-            .sweep_values(compactions);
-        self.iterator_stack.as_mut_slice().sweep_values(compactions);
-        self.exception_jump_target_stack
-            .as_mut_slice()
-            .sweep_values(compactions);
-        self.result.sweep_values(compactions);
-        self.exception.sweep_values(compactions);
-        self.reference.sweep_values(compactions);
+        self.stack.sweep_values(compactions);
+        self.reference_stack.sweep_values(compactions);
+        self.iterator_stack.sweep_values(compactions);
+        self.exception_jump_target_stack.sweep_values(compactions);
     }
 }
