Move async generator object to the stack

Author: HalidOdat

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

@@ -20,7 +20,7 @@ use crate::{
     string::common::StaticJsStrings,
     symbol::JsSymbol,
     value::JsValue,
-    vm::{CallFrame, CompletionRecord, GeneratorResumeKind},
+    vm::{CallFrame, CallFrameFlags, CompletionRecord, GeneratorResumeKind},
     Context, JsArgs, JsData, JsError, JsResult, JsString,
 };
 use boa_gc::{custom_trace, Finalize, Trace};
@@ -74,6 +74,10 @@ impl GeneratorContext {
 
         frame.fp = 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 subsuquent calls.
+        frame.flags |= CallFrameFlags::LOCALS_ALREADY_PUSHED;
+
         Self {
             call_frame: Some(frame),
             stack,
@@ -107,6 +111,13 @@ impl GeneratorContext {
         assert!(self.call_frame.is_some());
         result
     }
+
+    /// Returns the async generator object, if the function that this [`GeneratorContext`] is from an async generator, [`None`] otherwise.
+    pub(crate) fn async_generator_object(&self) -> Option<JsObject> {
+        self.call_frame
+            .as_ref()
+            .and_then(|frame| frame.async_generator_object(&self.stack))
+    }
 }
 
 /// The internal representation of a `Generator` object.

core/engine/src/bytecompiler/mod.rs

@@ -255,6 +255,8 @@ pub struct ByteCompiler<'ctx> {
     /// The number of arguments expected.
     pub(crate) length: u32,
 
+    pub(crate) locals_count: u32,
+
     /// \[\[ThisMode\]\]
     pub(crate) this_mode: ThisMode,
 
@@ -327,6 +329,7 @@ impl<'ctx> ByteCompiler<'ctx> {
             params: FormalParameterList::default(),
             current_open_environments_count: 0,
 
+            locals_count: 0,
             current_stack_value_count: 0,
             code_block_flags,
             handlers: ThinVec::default(),
@@ -1520,9 +1523,17 @@ impl<'ctx> ByteCompiler<'ctx> {
         }
         self.r#return(false);
 
+        if self.is_async_generator() {
+            self.locals_count += 1;
+        }
+        for handler in &mut self.handlers {
+            handler.stack_count += self.locals_count;
+        }
+
         CodeBlock {
             name: self.function_name,
             length: self.length,
+            locals_count: self.locals_count,
             this_mode: self.this_mode,
             params: self.params,
             bytecode: self.bytecode.into_boxed_slice(),

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

@@ -25,6 +25,9 @@ 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;
     }
 }
 
@@ -40,10 +43,6 @@ pub struct CallFrame {
     pub(crate) argument_count: u32,
     pub(crate) promise_capability: Option<PromiseCapability>,
 
-    // When an async generator is resumed, the generator object is needed
-    // to fulfill the steps 4.e-j in [AsyncGeneratorStart](https://tc39.es/ecma262/#sec-asyncgeneratorstart).
-    pub(crate) async_generator: Option<JsObject>,
-
     // Iterators and their `[[Done]]` flags that must be closed when an abrupt completion is thrown.
     pub(crate) iterators: ThinVec<IteratorRecord>,
 
@@ -123,6 +122,7 @@ 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_LOCAL_INDEX: u32 = 0;
 
     /// Creates a new `CallFrame` with the provided `CodeBlock`.
     pub(crate) fn new(
@@ -138,7 +138,6 @@ impl CallFrame {
             env_fp: 0,
             argument_count: 0,
             promise_capability: None,
-            async_generator: None,
             iterators: ThinVec::new(),
             binding_stack: Vec::new(),
             loop_iteration_count: 0,
@@ -201,6 +200,28 @@ impl CallFrame {
         vm.stack.truncate(fp as usize);
     }
 
+    /// Returns the async generator object, if the function that this [`GeneratorContext`] is from an async generator, [`None`] otherwise.
+    pub(crate) fn async_generator_object(&self, stack: &[JsValue]) -> Option<JsObject> {
+        if !self.code_block().is_async_generator() {
+            return None;
+        }
+
+        self.local(Self::ASYNC_GENERATOR_OBJECT_LOCAL_INDEX, stack)
+            .as_object()
+            .cloned()
+    }
+
+    /// Returns the local 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);
+        let at = self.fp + index;
+        &stack[at as usize]
+    }
+
     /// Does this have the [`CallFrameFlags::EXIT_EARLY`] flag.
     pub(crate) fn exit_early(&self) -> bool {
         self.flags.contains(CallFrameFlags::EXIT_EARLY)
@@ -213,6 +234,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)
+    }
 }
 
 /// ---- `CallFrame` stack methods ----

core/engine/src/vm/code_block.rs

@@ -177,6 +177,8 @@ pub struct CodeBlock {
     /// The number of arguments expected.
     pub(crate) length: u32,
 
+    pub(crate) locals_count: u32,
+
     /// \[\[ThisMode\]\]
     pub(crate) this_mode: ThisMode,
 
@@ -216,6 +218,7 @@ impl CodeBlock {
             name,
             flags: Cell::new(flags),
             length,
+            locals_count: 0,
             this_mode: ThisMode::Global,
             params: FormalParameterList::default(),
             handlers: ThinVec::default(),
@@ -286,6 +289,13 @@ impl CodeBlock {
         self.flags.get().contains(CodeBlockFlags::IS_GENERATOR)
     }
 
+    /// Returns true if this function a async generator function.
+    pub(crate) fn is_async_generator(&self) -> bool {
+        self.flags
+            .get()
+            .contains(CodeBlockFlags::IS_ASYNC | CodeBlockFlags::IS_GENERATOR)
+    }
+
     /// Returns true if this function an async function.
     pub(crate) fn is_ordinary(&self) -> bool {
         !self.is_async() && !self.is_generator()

core/engine/src/vm/mod.rs

@@ -162,6 +162,17 @@ 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,
+        //       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;
+            self.stack.resize_with(
+                current_stack_length + locals_count as usize,
+                JsValue::undefined,
+            );
+        }
+
         self.frames.push(frame);
     }
 

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

@@ -49,17 +49,16 @@ impl Operation for Await {
                     // d. Resume the suspended evaluation of asyncContext using NormalCompletion(value) as the result of the operation that suspended it.
                     let mut gen = captures.borrow_mut().take().expect("should only run once");
 
+                    // NOTE: We need to get the object before resuming, since it could clear the stack.
+                    let async_generator = gen.async_generator_object();
+
                     gen.resume(
                         Some(args.get_or_undefined(0).clone()),
                         GeneratorResumeKind::Normal,
                         context,
                     );
 
-                    if let Some(async_generator) = gen
-                        .call_frame
-                        .as_ref()
-                        .and_then(|f| f.async_generator.clone())
-                    {
+                    if let Some(async_generator) = async_generator {
                         async_generator
                             .downcast_mut::<AsyncGenerator>()
                             .expect("must be async generator")
@@ -92,17 +91,16 @@ impl Operation for Await {
 
                     let mut gen = captures.borrow_mut().take().expect("should only run once");
 
+                    // NOTE: We need to get the object before resuming, since it could clear the stack.
+                    let async_generator = gen.async_generator_object();
+
                     gen.resume(
                         Some(args.get_or_undefined(0).clone()),
                         GeneratorResumeKind::Throw,
                         context,
                     );
 
-                    if let Some(async_generator) = gen
-                        .call_frame
-                        .as_ref()
-                        .and_then(|f| f.async_generator.clone())
-                    {
+                    if let Some(async_generator) = async_generator {
                         async_generator
                             .downcast_mut::<AsyncGenerator>()
                             .expect("must be async generator")

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

@@ -41,7 +41,7 @@ impl Operation for Generator {
         let this_function_object =
             active_function.expect("active function should be set to the generator");
 
-        let frame = GeneratorContext::from_current(context);
+        let mut frame = GeneratorContext::from_current(context);
 
         let proto = this_function_object
             .get(PROTOTYPE, context)
@@ -64,7 +64,7 @@ impl Operation for Generator {
                 proto,
                 AsyncGenerator {
                     state: AsyncGeneratorState::SuspendedStart,
-                    context: Some(frame),
+                    context: None,
                     queue: VecDeque::new(),
                 },
             )
@@ -73,23 +73,29 @@ impl Operation for Generator {
                 context.root_shape(),
                 proto,
                 crate::builtins::generator::Generator {
-                    state: GeneratorState::SuspendedStart { context: frame },
+                    state: GeneratorState::Completed,
                 },
             )
         };
 
         if r#async {
-            let gen_clone = generator.clone();
+            let fp = frame
+                .call_frame
+                .as_ref()
+                .map_or(0, |frame| frame.fp as usize);
+            frame.stack[fp] = generator.clone().into();
+
             let mut gen = generator
                 .downcast_mut::<AsyncGenerator>()
                 .expect("must be object here");
-            let gen_context = gen.context.as_mut().expect("must exist");
-            // TODO: try to move this to the context itself.
-            gen_context
-                .call_frame
-                .as_mut()
-                .expect("should have a call frame initialized")
-                .async_generator = Some(gen_clone);
+
+            gen.context = Some(frame);
+        } else {
+            let mut gen = generator
+                .downcast_mut::<crate::builtins::generator::Generator>()
+                .expect("must be object here");
+
+            gen.state = GeneratorState::SuspendedStart { context: frame };
         }
 
         context.vm.set_return_value(generator.into());
@@ -111,14 +117,13 @@ impl Operation for AsyncGeneratorClose {
 
     fn execute(context: &mut Context) -> JsResult<CompletionType> {
         // Step 3.e-g in [AsyncGeneratorStart](https://tc39.es/ecma262/#sec-asyncgeneratorstart)
-        let generator_object = context
+        let async_generator_object = context
             .vm
             .frame()
-            .async_generator
-            .clone()
+            .async_generator_object(&context.vm.stack)
             .expect("There should be a object");
 
-        let mut gen = generator_object
+        let mut gen = async_generator_object
             .downcast_mut::<AsyncGenerator>()
             .expect("must be async generator");
 
@@ -136,7 +141,7 @@ impl Operation for AsyncGeneratorClose {
         } else {
             AsyncGenerator::complete_step(&next, Ok(return_value), true, None, context);
         }
-        AsyncGenerator::drain_queue(&generator_object, context);
+        AsyncGenerator::drain_queue(&async_generator_object, context);
 
         Ok(CompletionType::Normal)
     }

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

@@ -38,14 +38,13 @@ impl Operation for AsyncGeneratorYield {
     fn execute(context: &mut Context) -> JsResult<CompletionType> {
         let value = context.vm.pop();
 
-        let async_gen = context
+        let async_generator_object = context
             .vm
             .frame()
-            .async_generator
-            .clone()
+            .async_generator_object(&context.vm.stack)
             .expect("`AsyncGeneratorYield` must only be called inside async generators");
         let completion = Ok(value);
-        let next = async_gen
+        let next = async_generator_object
             .downcast_mut::<AsyncGenerator>()
             .expect("must be async generator object")
             .queue
@@ -55,7 +54,7 @@ impl Operation for AsyncGeneratorYield {
         // TODO: 7. Let previousContext be the second to top element of the execution context stack.
         AsyncGenerator::complete_step(&next, completion, false, None, context);
 
-        let mut generator_object_mut = async_gen.borrow_mut();
+        let mut generator_object_mut = async_generator_object.borrow_mut();
         let gen = generator_object_mut
             .downcast_mut::<AsyncGenerator>()
             .expect("must be async generator object");