Move PromiseCapability to stack

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

@@ -1523,8 +1523,13 @@ 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.locals_count += 3;
+            if self.is_generator() {
+                // NOTE: +1 for the async generator function
+                self.locals_count += 1;
+            }
         }
         for handler in &mut self.handlers {
             handler.stack_count += self.locals_count;

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,
@@ -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,
@@ -221,6 +225,68 @@ impl CallFrame {
             .cloned()
     }
 
+    pub(crate) fn promise_capability(&self, stack: &[JsValue]) -> Option<PromiseCapability> {
+        if !self.code_block().is_async() {
+            return None;
+        }
+
+        let promise = self
+            .local(Self::PROMISE_CAPABILITY_PROMISE_REGISTER_INDEX, stack)
+            .as_object()
+            .cloned()?;
+        let resolve = self
+            .local(Self::PROMISE_CAPABILITY_RESOLVE_REGISTER_INDEX, stack)
+            .as_object()
+            .cloned()
+            .and_then(JsFunction::from_object)?;
+        let reject = self
+            .local(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_local(
+            Self::PROMISE_CAPABILITY_PROMISE_REGISTER_INDEX,
+            promise_capability
+                .map(PromiseCapability::promise)
+                .cloned()
+                .map_or_else(JsValue::undefined, Into::into),
+            stack,
+        );
+        self.set_local(
+            Self::PROMISE_CAPABILITY_RESOLVE_REGISTER_INDEX,
+            promise_capability
+                .map(PromiseCapability::resolve)
+                .cloned()
+                .map_or_else(JsValue::undefined, Into::into),
+            stack,
+        );
+        self.set_local(
+            Self::PROMISE_CAPABILITY_REJECT_REGISTER_INDEX,
+            promise_capability
+                .map(PromiseCapability::reject)
+                .cloned()
+                .map_or_else(JsValue::undefined, Into::into),
+            stack,
+        );
+    }
+
     /// Returns the local at the given index.
     ///
     /// # Panics
@@ -232,6 +298,17 @@ impl CallFrame {
         &stack[at as usize]
     }
 
+    /// Returns the local at the given index.
+    ///
+    /// # Panics
+    ///
+    /// If the index is out of bounds.
+    pub(crate) fn set_local(&self, index: u32, value: JsValue, stack: &mut [JsValue]) {
+        debug_assert!(index < self.code_block().locals_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)

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 {

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

@@ -13,7 +13,7 @@ use crate::{
     vm::{
         call_frame::GeneratorResumeKind,
         opcode::{Operation, ReThrow},
-        CompletionType,
+        CallFrame, CompletionType,
     },
     Context, JsError, JsObject, JsResult, JsValue,
 };
@@ -79,11 +79,12 @@ impl Operation for Generator {
         };
 
         if r#async {
-            let fp = frame
+            let rp = frame
                 .call_frame
                 .as_ref()
                 .map_or(0, |frame| frame.rp as usize);
-            frame.stack[fp] = generator.clone().into();
+            frame.stack[rp + CallFrame::ASYNC_GENERATOR_OBJECT_REGISTER_INDEX as usize] =
+                generator.clone().into();
 
             let mut gen = generator
                 .downcast_mut::<AsyncGenerator>()