Refactor vm calling convention to allow locals (#3496)

* Refactor vm calling convention to allow locals

* Move async generator object to the stack

* Apply review and some clean-up

* Add todo for storing fp in CallFrame

* Rename ASYNC_GENERATOR_OBJECT_LOCAL_INDEX to ASYNC_GENERATOR_OBJECT_REGISTER_INDEX

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

@@ -938,10 +938,7 @@ pub(crate) fn function_call(
 
     context.vm.push_frame(frame);
 
-    let fp = context.vm.stack.len() - argument_count - CallFrame::FUNCTION_PROLOGUE;
-    context.vm.frame_mut().fp = fp as u32;
-
-    let this = context.vm.stack[fp + CallFrame::THIS_POSITION].clone();
+    let this = context.vm.frame().this(&context.vm);
 
     let lexical_this_mode = code.this_mode == ThisMode::Lexical;
 
@@ -1013,8 +1010,6 @@ fn function_construct(
 
     let new_target = context.vm.pop();
 
-    let at = context.vm.stack.len() - argument_count;
-
     let this = if code.is_derived_constructor() {
         None
     } else {
@@ -1040,9 +1035,12 @@ fn function_construct(
         .with_env_fp(env_fp)
         .with_flags(CallFrameFlags::CONSTRUCT);
 
+    let len = context.vm.stack.len();
+
     context.vm.push_frame(frame);
 
-    context.vm.frame_mut().fp = at as u32 - 1;
+    // NOTE(HalidOdat): +1 because we insert `this` value below.
+    context.vm.frame_mut().rp = len as u32 + 1;
 
     let mut last_env = 0;
 
@@ -1075,10 +1073,10 @@ fn function_construct(
     );
 
     // Insert `this` value
-    context
-        .vm
-        .stack
-        .insert(at - 1, this.map(JsValue::new).unwrap_or_default());
+    context.vm.stack.insert(
+        len - argument_count - 1,
+        this.map(JsValue::new).unwrap_or_default(),
+    );
 
     Ok(CallValue::Ready)
 }

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};
@@ -64,28 +64,24 @@ pub(crate) struct GeneratorContext {
 }
 
 impl GeneratorContext {
-    /// Creates a new `GeneratorContext` from the raw `Context` state components.
-    pub(crate) fn new(stack: Vec<JsValue>, call_frame: CallFrame) -> Self {
-        Self {
-            stack,
-            call_frame: Some(call_frame),
-        }
-    }
-
     /// Creates a new `GeneratorContext` from the current `Context` state.
     pub(crate) fn from_current(context: &mut Context) -> Self {
         let mut frame = context.vm.frame().clone();
         frame.environments = context.vm.environments.clone();
         frame.realm = context.realm().clone();
-        let fp = context.vm.frame().fp as usize;
-        let this = Self {
-            call_frame: Some(frame),
-            stack: context.vm.stack[fp..].to_vec(),
-        };
+        let fp = frame.fp() as usize;
+        let stack = context.vm.stack.split_off(fp);
 
-        context.vm.stack.truncate(fp);
+        frame.rp = CallFrame::FUNCTION_PROLOGUE + frame.argument_count;
 
-        this
+        // 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;
+
+        Self {
+            call_frame: Some(frame),
+            stack,
+        }
     }
 
     /// Resumes execution with `GeneratorContext` as the current execution context.
@@ -96,11 +92,11 @@ impl GeneratorContext {
         context: &mut Context,
     ) -> CompletionRecord {
         std::mem::swap(&mut context.vm.stack, &mut self.stack);
-        context
-            .vm
-            .push_frame(self.call_frame.take().expect("should have a call frame"));
+        let frame = self.call_frame.take().expect("should have a call frame");
+        let rp = frame.rp;
+        context.vm.push_frame(frame);
 
-        context.vm.frame_mut().fp = 0;
+        context.vm.frame_mut().rp = rp;
         context.vm.frame_mut().set_exit_early(true);
 
         if let Some(value) = value {
@@ -115,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,8 +329,8 @@ impl<'ctx> ByteCompiler<'ctx> {
             params: FormalParameterList::default(),
             current_open_environments_count: 0,
 
-            // This starts at two because the first value is the `this` value, then function object.
-            current_stack_value_count: 2,
+            locals_count: 0,
+            current_stack_value_count: 0,
             code_block_flags,
             handlers: ThinVec::default(),
             ic: Vec::default(),
@@ -1521,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/module/synthetic.rs

@@ -324,7 +324,7 @@ impl SyntheticModule {
         // 11. Suspend moduleContext and remove it from the execution context stack.
         // 12. Resume the context that is now on the top of the execution context stack as the running execution context.
         let frame = context.vm.pop_frame().expect("there should be a frame");
-        context.vm.stack.truncate(frame.fp as usize);
+        frame.restore_stack(&mut context.vm);
 
         // 13. Let pc be ! NewPromiseCapability(%Promise%).
         let (promise, ResolvingFunctions { resolve, reject }) = JsPromise::new_pending(context);

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;
     }
 }
 
@@ -33,17 +36,16 @@ bitflags::bitflags! {
 pub struct CallFrame {
     pub(crate) code_block: Gc<CodeBlock>,
     pub(crate) pc: u32,
-    pub(crate) fp: u32,
-    pub(crate) env_fp: u32,
-    // Tracks the number of environments in environment entry.
-    // On abrupt returns this is used to decide how many environments need to be pop'ed.
+    /// The register pointer, points to the first register in the stack.
+    ///
+    // TODO: Check if storing the frame pointer instead of argument count and computing the
+    //       argument count based on the pointers would be better for accessing the arguments
+    //       and the elements before the register pointer.
+    pub(crate) rp: u32,
     pub(crate) argument_count: u32,
+    pub(crate) env_fp: 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>,
 
@@ -81,22 +83,56 @@ impl CallFrame {
 impl CallFrame {
     /// This is the size of the function prologue.
     ///
-    /// The position of the elements are relative to the [`CallFrame::fp`].
+    /// The position of the elements are relative to the [`CallFrame::fp`] (register pointer).
+    ///
+    /// ```text
+    ///                      Setup by the caller
+    ///   ┌─────────────────────────────────────────────────────────┐ ┌───── register pointer
+    ///   ▼                                                         ▼ ▼
+    /// | -(2 + N): this | -(1 + N): func | -N: arg1 | ... | -1: argN | 0: local1 | ... | K: localK |
+    ///   ▲                              ▲   ▲                      ▲   ▲                         ▲
+    ///   └──────────────────────────────┘   └──────────────────────┘   └─────────────────────────┘
+    ///         function prologue                    arguments              Setup by the callee
+    ///   ▲
+    ///   └─ Frame pointer
+    /// ```
+    ///
+    /// ### Example
+    ///
+    /// The following function calls, generate the following stack:
+    ///
+    /// ```JavaScript
+    /// function x(a) {
+    /// }
+    /// function y(b, c) {
+    ///     return x(b + c)
+    /// }
+    ///
+    /// y(1, 2)
+    /// ```
     ///
     /// ```text
-    ///   --- frame pointer                             arguments
-    ///  /                      __________________________/
-    /// /                      /                          \
-    /// | 0: this | 1: func | 2: arg1 | ... | (2 + N): argN |
-    ///    \            /
-    ///     ------------
-    ///     |
-    /// function prolugue
+    ///     caller prologue    caller arguments   callee prologue   callee arguments
+    ///   ┌─────────────────┐   ┌─────────┐   ┌─────────────────┐  ┌──────┐
+    ///   ▼                 ▼   ▼         ▼   │                 ▼  ▼      ▼
+    /// | 0: undefined | 1: y | 2: 1 | 3: 2 | 4: undefined | 5: x | 6:  3  |
+    /// ▲                                   ▲                            ▲
+    /// │       caller register pointer ────┤                            │
+    /// │                                   │                callee register pointer
+    /// │                             callee frame pointer
+    /// │
+    /// └─────  caller frame pointer
+    ///
     /// ```
-    pub(crate) const FUNCTION_PROLOGUE: usize = 2;
-    pub(crate) const THIS_POSITION: usize = 0;
-    pub(crate) const FUNCTION_POSITION: usize = 1;
-    pub(crate) const FIRST_ARGUMENT_POSITION: usize = 2;
+    ///
+    /// Some questions:
+    ///
+    /// - Who is responsible for cleaning up the stack after a call? The rust caller.
+    ///
+    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;
 
     /// Creates a new `CallFrame` with the provided `CodeBlock`.
     pub(crate) fn new(
@@ -108,11 +144,10 @@ impl CallFrame {
         Self {
             code_block,
             pc: 0,
-            fp: 0,
+            rp: 0,
             env_fp: 0,
             argument_count: 0,
             promise_capability: None,
-            async_generator: None,
             iterators: ThinVec::new(),
             binding_stack: Vec::new(),
             loop_iteration_count: 0,
@@ -142,19 +177,61 @@ impl CallFrame {
     }
 
     pub(crate) fn this(&self, vm: &Vm) -> JsValue {
-        let this_index = self.fp as usize + Self::THIS_POSITION;
-        vm.stack[this_index].clone()
+        let this_index = self.rp - self.argument_count - Self::THIS_POSITION;
+        vm.stack[this_index as usize].clone()
     }
 
     pub(crate) fn function(&self, vm: &Vm) -> Option<JsObject> {
-        let function_index = self.fp as usize + Self::FUNCTION_POSITION;
-        if let Some(object) = vm.stack[function_index].as_object() {
+        let function_index = self.rp - self.argument_count - Self::FUNCTION_POSITION;
+        if let Some(object) = vm.stack[function_index as usize].as_object() {
             return Some(object.clone());
         }
 
         None
     }
 
+    pub(crate) fn arguments<'stack>(&self, vm: &'stack Vm) -> &'stack [JsValue] {
+        let rp = self.rp as usize;
+        let argument_count = self.argument_count as usize;
+        let arguments_start = rp - argument_count;
+        &vm.stack[arguments_start..rp]
+    }
+
+    pub(crate) fn argument<'stack>(&self, index: usize, vm: &'stack Vm) -> Option<&'stack JsValue> {
+        self.arguments(vm).get(index)
+    }
+
+    pub(crate) fn fp(&self) -> u32 {
+        self.rp - self.argument_count - Self::FUNCTION_PROLOGUE
+    }
+
+    pub(crate) fn restore_stack(&self, vm: &mut Vm) {
+        let fp = self.fp();
+        vm.stack.truncate(fp as usize);
+    }
+
+    /// Returns the async generator object, if the function that this [`CallFrame`] 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_REGISTER_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.rp + 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)
@@ -167,12 +244,16 @@ 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 ----
 impl CallFrame {
-    pub(crate) fn set_frame_pointer(&mut self, pointer: u32) {
-        self.fp = pointer;
+    pub(crate) fn set_register_pointer(&mut self, pointer: u32) {
+        self.rp = pointer;
     }
 }