Move PromiseCapability to stack (#3528)

* Move `PromiseCapability` to stack

* Apply review

* Rename locals to register for consistency

core/engine/src/builtins/generator/mod.rs

@@ -75,8 +75,8 @@ impl GeneratorContext {
         frame.rp = CallFrame::FUNCTION_PROLOGUE + frame.argument_count;
 
         // NOTE: Since we get a pre-built call frame with stack, and we reuse them.
-        //       So we don't need to push the locals in subsequent calls.
-        frame.flags |= CallFrameFlags::LOCALS_ALREADY_PUSHED;
+        //       So we don't need to push the registers in subsequent calls.
+        frame.flags |= CallFrameFlags::REGISTERS_ALREADY_PUSHED;
 
         Self {
             call_frame: Some(frame),

core/engine/src/builtins/promise/mod.rs

@@ -168,10 +168,10 @@ pub(crate) use if_abrupt_reject_promise;
 #[derive(Debug, Clone, Finalize)]
 pub(crate) struct PromiseCapability {
     /// The `[[Promise]]` field.
-    promise: JsObject,
+    pub(crate) promise: JsObject,
 
     /// The resolving functions,
-    functions: ResolvingFunctions,
+    pub(crate) functions: ResolvingFunctions,
 }
 
 // SAFETY: manually implementing `Trace` to allow destructuring.

core/engine/src/bytecompiler/mod.rs

@@ -255,7 +255,7 @@ pub struct ByteCompiler<'ctx> {
     /// The number of arguments expected.
     pub(crate) length: u32,
 
-    pub(crate) locals_count: u32,
+    pub(crate) register_count: u32,
 
     /// \[\[ThisMode\]\]
     pub(crate) this_mode: ThisMode,
@@ -329,7 +329,7 @@ impl<'ctx> ByteCompiler<'ctx> {
             params: FormalParameterList::default(),
             current_open_environments_count: 0,
 
-            locals_count: 0,
+            register_count: 0,
             current_stack_value_count: 0,
             code_block_flags,
             handlers: ThinVec::default(),
@@ -1523,17 +1523,25 @@ impl<'ctx> ByteCompiler<'ctx> {
         }
         self.r#return(false);
 
-        if self.is_async_generator() {
-            self.locals_count += 1;
+        if self.is_async() {
+            // NOTE: +3 for the promise capability
+            self.register_count += 3;
+            if self.is_generator() {
+                // NOTE: +1 for the async generator function
+                self.register_count += 1;
+            }
         }
+
+        // NOTE: Offset the handlers stack count so we don't pop the registers
+        //       when a exception is thrown.
         for handler in &mut self.handlers {
-            handler.stack_count += self.locals_count;
+            handler.stack_count += self.register_count;
         }
 
         CodeBlock {
             name: self.function_name,
             length: self.length,
-            locals_count: self.locals_count,
+            register_count: self.register_count,
             this_mode: self.this_mode,
             params: self.params,
             bytecode: self.bytecode.into_boxed_slice(),

core/engine/src/module/source.rs

@@ -1327,7 +1327,7 @@ impl SourceTextModule {
 
         // 9. Perform ! module.ExecuteModule(capability).
         // 10. Return unused.
-        self.execute(Some(capability), context)
+        self.execute(Some(&capability), context)
             .expect("async modules cannot directly throw");
     }
 
@@ -1741,7 +1741,7 @@ impl SourceTextModule {
     /// [spec]: https://tc39.es/ecma262/#sec-source-text-module-record-execute-module
     fn execute(
         &self,
-        capability: Option<PromiseCapability>,
+        capability: Option<&PromiseCapability>,
         context: &mut Context,
     ) -> JsResult<()> {
         // 1. Let moduleContext be a new ECMAScript code execution context.
@@ -1763,21 +1763,27 @@ impl SourceTextModule {
         // 6. Set the VariableEnvironment of moduleContext to module.[[Environment]].
         // 7. Set the LexicalEnvironment of moduleContext to module.[[Environment]].
         let env_fp = environments.len() as u32;
-        let mut callframe = CallFrame::new(
+        let callframe = CallFrame::new(
             codeblock,
             Some(ActiveRunnable::Module(self.parent())),
             environments,
             realm,
         )
         .with_env_fp(env_fp)
         .with_flags(CallFrameFlags::EXIT_EARLY);
-        callframe.promise_capability = capability;
 
         // 8. Suspend the running execution context.
         context
             .vm
             .push_frame_with_stack(callframe, JsValue::undefined(), JsValue::null());
 
+        context
+            .vm
+            .frames
+            .last()
+            .expect("there should be a frame")
+            .set_promise_capability(&mut context.vm.stack, capability);
+
         // 9. If module.[[HasTLA]] is false, then
         //    a. Assert: capability is not present.
         //    b. Push moduleContext onto the execution context stack; moduleContext is now the running execution context.

core/engine/src/vm/call_frame/mod.rs

@@ -3,9 +3,12 @@
 //! This module will provides everything needed to implement the `CallFrame`
 
 use crate::{
-    builtins::{iterable::IteratorRecord, promise::PromiseCapability},
+    builtins::{
+        iterable::IteratorRecord,
+        promise::{PromiseCapability, ResolvingFunctions},
+    },
     environments::{BindingLocator, EnvironmentStack},
-    object::JsObject,
+    object::{JsFunction, JsObject},
     realm::Realm,
     vm::CodeBlock,
     JsValue,
@@ -26,8 +29,8 @@ bitflags::bitflags! {
         /// Was this [`CallFrame`] created from the `__construct__()` internal object method?
         const CONSTRUCT = 0b0000_0010;
 
-        /// Does this [`CallFrame`] need to push local variables on [`Vm::push_frame()`].
-        const LOCALS_ALREADY_PUSHED = 0b0000_0100;
+        /// Does this [`CallFrame`] need to push registers on [`Vm::push_frame()`].
+        const REGISTERS_ALREADY_PUSHED = 0b0000_0100;
     }
 }
 
@@ -44,7 +47,6 @@ pub struct CallFrame {
     pub(crate) rp: u32,
     pub(crate) argument_count: u32,
     pub(crate) env_fp: u32,
-    pub(crate) promise_capability: Option<PromiseCapability>,
 
     // Iterators and their `[[Done]]` flags that must be closed when an abrupt completion is thrown.
     pub(crate) iterators: ThinVec<IteratorRecord>,
@@ -132,7 +134,10 @@ impl CallFrame {
     pub(crate) const FUNCTION_PROLOGUE: u32 = 2;
     pub(crate) const THIS_POSITION: u32 = 2;
     pub(crate) const FUNCTION_POSITION: u32 = 1;
-    pub(crate) const ASYNC_GENERATOR_OBJECT_REGISTER_INDEX: u32 = 0;
+    pub(crate) const PROMISE_CAPABILITY_PROMISE_REGISTER_INDEX: u32 = 0;
+    pub(crate) const PROMISE_CAPABILITY_RESOLVE_REGISTER_INDEX: u32 = 1;
+    pub(crate) const PROMISE_CAPABILITY_REJECT_REGISTER_INDEX: u32 = 2;
+    pub(crate) const ASYNC_GENERATOR_OBJECT_REGISTER_INDEX: u32 = 3;
 
     /// Creates a new `CallFrame` with the provided `CodeBlock`.
     pub(crate) fn new(
@@ -147,7 +152,6 @@ impl CallFrame {
             rp: 0,
             env_fp: 0,
             argument_count: 0,
-            promise_capability: None,
             iterators: ThinVec::new(),
             binding_stack: Vec::new(),
             loop_iteration_count: 0,
@@ -216,22 +220,95 @@ impl CallFrame {
             return None;
         }
 
-        self.local(Self::ASYNC_GENERATOR_OBJECT_REGISTER_INDEX, stack)
+        self.register(Self::ASYNC_GENERATOR_OBJECT_REGISTER_INDEX, stack)
             .as_object()
             .cloned()
     }
 
-    /// Returns the local at the given index.
+    pub(crate) fn promise_capability(&self, stack: &[JsValue]) -> Option<PromiseCapability> {
+        if !self.code_block().is_async() {
+            return None;
+        }
+
+        let promise = self
+            .register(Self::PROMISE_CAPABILITY_PROMISE_REGISTER_INDEX, stack)
+            .as_object()
+            .cloned()?;
+        let resolve = self
+            .register(Self::PROMISE_CAPABILITY_RESOLVE_REGISTER_INDEX, stack)
+            .as_object()
+            .cloned()
+            .and_then(JsFunction::from_object)?;
+        let reject = self
+            .register(Self::PROMISE_CAPABILITY_REJECT_REGISTER_INDEX, stack)
+            .as_object()
+            .cloned()
+            .and_then(JsFunction::from_object)?;
+
+        Some(PromiseCapability {
+            promise,
+            functions: ResolvingFunctions { resolve, reject },
+        })
+    }
+
+    pub(crate) fn set_promise_capability(
+        &self,
+        stack: &mut [JsValue],
+        promise_capability: Option<&PromiseCapability>,
+    ) {
+        debug_assert!(
+            self.code_block().is_async(),
+            "Only async functions have a promise capability"
+        );
+
+        self.set_register(
+            Self::PROMISE_CAPABILITY_PROMISE_REGISTER_INDEX,
+            promise_capability
+                .map(PromiseCapability::promise)
+                .cloned()
+                .map_or_else(JsValue::undefined, Into::into),
+            stack,
+        );
+        self.set_register(
+            Self::PROMISE_CAPABILITY_RESOLVE_REGISTER_INDEX,
+            promise_capability
+                .map(PromiseCapability::resolve)
+                .cloned()
+                .map_or_else(JsValue::undefined, Into::into),
+            stack,
+        );
+        self.set_register(
+            Self::PROMISE_CAPABILITY_REJECT_REGISTER_INDEX,
+            promise_capability
+                .map(PromiseCapability::reject)
+                .cloned()
+                .map_or_else(JsValue::undefined, Into::into),
+            stack,
+        );
+    }
+
+    /// Returns the register at the given index.
     ///
     /// # Panics
     ///
     /// If the index is out of bounds.
-    pub(crate) fn local<'stack>(&self, index: u32, stack: &'stack [JsValue]) -> &'stack JsValue {
-        debug_assert!(index < self.code_block().locals_count);
+    pub(crate) fn register<'stack>(&self, index: u32, stack: &'stack [JsValue]) -> &'stack JsValue {
+        debug_assert!(index < self.code_block().register_count);
         let at = self.rp + index;
         &stack[at as usize]
     }
 
+    /// Sets the register at the given index.
+    ///
+    /// # Panics
+    ///
+    /// If the index is out of bounds.
+    pub(crate) fn set_register(&self, index: u32, value: JsValue, stack: &mut [JsValue]) {
+        debug_assert!(index < self.code_block().register_count);
+        let at = self.rp + index;
+        stack[at as usize] = value;
+    }
+
     /// Does this have the [`CallFrameFlags::EXIT_EARLY`] flag.
     pub(crate) fn exit_early(&self) -> bool {
         self.flags.contains(CallFrameFlags::EXIT_EARLY)
@@ -244,9 +321,10 @@ impl CallFrame {
     pub(crate) fn construct(&self) -> bool {
         self.flags.contains(CallFrameFlags::CONSTRUCT)
     }
-    /// Does this [`CallFrame`] need to push local variables on [`Vm::push_frame()`].
-    pub(crate) fn locals_already_pushed(&self) -> bool {
-        self.flags.contains(CallFrameFlags::LOCALS_ALREADY_PUSHED)
+    /// Does this [`CallFrame`] need to push registers on [`Vm::push_frame()`].
+    pub(crate) fn registers_already_pushed(&self) -> bool {
+        self.flags
+            .contains(CallFrameFlags::REGISTERS_ALREADY_PUSHED)
     }
 }
 

core/engine/src/vm/code_block.rs

@@ -177,7 +177,7 @@ pub struct CodeBlock {
     /// The number of arguments expected.
     pub(crate) length: u32,
 
-    pub(crate) locals_count: u32,
+    pub(crate) register_count: u32,
 
     /// \[\[ThisMode\]\]
     pub(crate) this_mode: ThisMode,
@@ -218,7 +218,7 @@ impl CodeBlock {
             name,
             flags: Cell::new(flags),
             length,
-            locals_count: 0,
+            register_count: 0,
             this_mode: ThisMode::Global,
             params: FormalParameterList::default(),
             handlers: ThinVec::default(),

core/engine/src/vm/mod.rs

@@ -162,13 +162,13 @@ impl Vm {
         std::mem::swap(&mut self.environments, &mut frame.environments);
         std::mem::swap(&mut self.realm, &mut frame.realm);
 
-        // NOTE: We need to check if we already pushed the locals,
+        // NOTE: We need to check if we already pushed the registers,
         //       since generator-like functions push the same call
         //       frame with pre-built stack.
-        if !frame.locals_already_pushed() {
-            let locals_count = frame.code_block().locals_count;
+        if !frame.registers_already_pushed() {
+            let register_count = frame.code_block().register_count;
             self.stack.resize_with(
-                current_stack_length + locals_count as usize,
+                current_stack_length + register_count as usize,
                 JsValue::undefined,
             );
         }

core/engine/src/vm/opcode/await/mod.rs

@@ -33,6 +33,16 @@ impl Operation for Await {
             context,
         )?;
 
+        let return_value = context
+            .vm
+            .frame()
+            .promise_capability(&context.vm.stack)
+            .as_ref()
+            .map(PromiseCapability::promise)
+            .cloned()
+            .map(JsValue::from)
+            .unwrap_or_default();
+
         let gen = GeneratorContext::from_current(context);
 
         let captures = Gc::new(GcRefCell::new(Some(gen)));
@@ -125,16 +135,6 @@ impl Operation for Await {
             context,
         );
 
-        let return_value = context
-            .vm
-            .frame()
-            .promise_capability
-            .as_ref()
-            .map(PromiseCapability::promise)
-            .cloned()
-            .map(JsValue::from)
-            .unwrap_or_default();
-
         context.vm.set_return_value(return_value);
         Ok(CompletionType::Yield)
     }
@@ -153,7 +153,12 @@ impl Operation for CreatePromiseCapability {
     const COST: u8 = 8;
 
     fn execute(context: &mut Context) -> JsResult<CompletionType> {
-        if context.vm.frame().promise_capability.is_some() {
+        if context
+            .vm
+            .frame()
+            .promise_capability(&context.vm.stack)
+            .is_some()
+        {
             return Ok(CompletionType::Normal);
         }
 
@@ -163,7 +168,12 @@ impl Operation for CreatePromiseCapability {
         )
         .expect("cannot fail per spec");
 
-        context.vm.frame_mut().promise_capability = Some(promise_capability);
+        context
+            .vm
+            .frames
+            .last()
+            .expect("there should be a frame")
+            .set_promise_capability(&mut context.vm.stack, Some(&promise_capability));
         Ok(CompletionType::Normal)
     }
 }
@@ -183,7 +193,8 @@ impl Operation for CompletePromiseCapability {
     fn execute(context: &mut Context) -> JsResult<CompletionType> {
         // If the current executing function is an async function we have to resolve/reject it's promise at the end.
         // The relevant spec section is 3. in [AsyncBlockStart](https://tc39.es/ecma262/#sec-asyncblockstart).
-        let Some(promise_capability) = context.vm.frame_mut().promise_capability.take() else {
+        let Some(promise_capability) = context.vm.frame().promise_capability(&context.vm.stack)
+        else {
             return if context.vm.pending_exception.is_some() {
                 Ok(CompletionType::Throw)
             } else {