Fix AsyncGenerator to correctly handle return inside then (#3879)

* Fix invalid assumptions for `AsyncGenerator`

* Patch AsyncGenerator functions

* Progress

* Progress 2

* Finish implementation

* cargo fmt

* cargo clippy

* Fix await resumptions

* Assert count in queue test

core/engine/src/tests/async_generator.rs

@@ -0,0 +1,122 @@
+use crate::{
+    builtins::promise::PromiseState, object::JsPromise, run_test_actions, Context, JsValue,
+    TestAction,
+};
+use boa_macros::js_str;
+use indoc::indoc;
+
+#[track_caller]
+fn assert_promise_iter_value(
+    promise: &JsValue,
+    target: &JsValue,
+    done: bool,
+    context: &mut Context,
+) {
+    let promise = JsPromise::from_object(promise.as_object().unwrap().clone()).unwrap();
+    let PromiseState::Fulfilled(v) = promise.state() else {
+        panic!("promise was not fulfilled");
+    };
+    let o = v.as_object().unwrap();
+    let value = o.get(js_str!("value"), context).unwrap();
+    let d = o
+        .get(js_str!("done"), context)
+        .unwrap()
+        .as_boolean()
+        .unwrap();
+    assert_eq!(&value, target);
+    assert_eq!(d, done);
+}
+
+#[test]
+fn return_on_then_infinite_loop() {
+    // Checks that calling `return` inside `then` only enters an infinite loop without
+    // crashing the engine.
+    run_test_actions([
+        TestAction::run(indoc! {r#"
+            async function* f() {}
+            const g = f();
+            let count = 0;
+            Object.defineProperty(Object.prototype, "then", {
+                get: function() {
+                    if (count < 100) {
+                        count++;
+                        g.return();
+                    }
+                    return;
+                },
+            });
+            g.return();
+        "#}),
+        TestAction::inspect_context(Context::run_jobs),
+        TestAction::assert_eq("count", 100),
+    ]);
+}
+
+#[test]
+fn return_on_then_single() {
+    // Checks that calling `return` inside `then` once runs without panicking.
+    run_test_actions([
+        TestAction::run(indoc! {r#"
+            async function* f() {}
+            const g = f();
+            let first = true;
+            Object.defineProperty(Object.prototype, "then", {
+                get: function() {
+                    if (first) {
+                        first = false;
+                        g.return();
+                    }
+                    return;
+                },
+            });
+            let ret = g.return()
+        "#}),
+        TestAction::inspect_context(Context::run_jobs),
+        TestAction::assert_eq("first", false),
+        TestAction::assert_with_op("ret", |ret, context| {
+            assert_promise_iter_value(&ret, &JsValue::undefined(), true, context);
+            true
+        }),
+    ]);
+}
+
+#[test]
+fn return_on_then_queue() {
+    // Checks that calling `return` inside `then` doesn't mess with the request queue.
+    run_test_actions([
+        TestAction::run(indoc! {r#"
+            async function* f() {
+                yield 1;
+                yield 2;
+            }
+            const g = f();
+            let count = 0;
+            Object.defineProperty(Object.prototype, "then", {
+                get: function() {
+                    if (count < 2) {
+                        count++;
+                        g.return();
+                    }
+                    return;
+                },
+            });
+            let first = g.next();
+            let second = g.next();
+            let ret = g.return();
+        "#}),
+        TestAction::inspect_context(Context::run_jobs),
+        TestAction::assert_with_op("first", |first, context| {
+            assert_promise_iter_value(&first, &JsValue::from(1), false, context);
+            true
+        }),
+        TestAction::assert_with_op("second", |second, context| {
+            assert_promise_iter_value(&second, &JsValue::from(2), false, context);
+            true
+        }),
+        TestAction::assert_with_op("ret", |ret, context| {
+            assert_promise_iter_value(&ret, &JsValue::undefined(), true, context);
+            true
+        }),
+        TestAction::assert_eq("count", JsValue::from(2)),
+    ]);
+}

core/engine/src/tests/mod.rs

@@ -1,6 +1,9 @@
+#![cfg(any(not(feature = "intl"), feature = "intl_bundled"))]
+
 use boa_macros::js_str;
 use indoc::indoc;
 
+mod async_generator;
 mod control_flow;
 mod env;
 mod function;

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

@@ -292,6 +292,7 @@ impl CallFrame {
     /// # Panics
     ///
     /// If the index is out of bounds.
+    #[track_caller]
     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;

core/engine/src/vm/completion_record.rs

@@ -29,10 +29,6 @@ unsafe impl Trace for CompletionRecord {
 
 // ---- `CompletionRecord` methods ----
 impl CompletionRecord {
-    pub(crate) const fn is_throw_completion(&self) -> bool {
-        matches!(self, Self::Throw(_))
-    }
-
     /// This function will consume the current `CompletionRecord` and return a `JsResult<JsValue>`
     // NOTE: rustc bug around evaluating destructors that prevents this from being a const function.
     // Related issue(s):

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

@@ -1,4 +1,6 @@
-use boa_gc::{Gc, GcRefCell};
+use std::cell::Cell;
+
+use boa_gc::Gc;
 use boa_macros::js_str;
 
 use crate::{
@@ -46,7 +48,17 @@ impl Operation for Await {
 
         let gen = GeneratorContext::from_current(context);
 
-        let captures = Gc::new(GcRefCell::new(Some(gen)));
+        // Even though it would be great to avoid cloning, we need to ensure
+        // the original async generator has a copy of the context in case it is resumed
+        // by a `return` or `throw` call instead of a continuation.
+        if let Some(async_generator) = gen.async_generator_object() {
+            async_generator
+                .downcast_mut::<AsyncGenerator>()
+                .expect("must be async generator")
+                .context = Some(gen.clone());
+        }
+
+        let captures = Gc::new(Cell::new(Some(gen)));
 
         // 3. Let fulfilledClosure be a new Abstract Closure with parameters (value) that captures asyncContext and performs the following steps when called:
         // 4. Let onFulfilled be CreateBuiltinFunction(fulfilledClosure, 1, "", « »).
@@ -58,7 +70,7 @@ impl Operation for Await {
                     // b. Suspend prevContext.
                     // c. Push asyncContext onto the execution context stack; asyncContext is now the running execution context.
                     // 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");
+                    let mut gen = captures.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();
@@ -100,7 +112,7 @@ impl Operation for Await {
                     // e. Assert: When we reach this step, asyncContext has already been removed from the execution context stack and prevContext is the currently running execution context.
                     // f. Return undefined.
 
-                    let mut gen = captures.borrow_mut().take().expect("should only run once");
+                    let mut gen = captures.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();

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

@@ -118,31 +118,35 @@ impl Operation for AsyncGeneratorClose {
 
     fn execute(context: &mut Context) -> JsResult<CompletionType> {
         // Step 3.e-g in [AsyncGeneratorStart](https://tc39.es/ecma262/#sec-asyncgeneratorstart)
-        let async_generator_object = context
+        let generator = context
             .vm
             .frame()
             .async_generator_object(&context.vm.stack)
-            .expect("There should be a object");
-
-        let mut gen = async_generator_object
-            .downcast_mut::<AsyncGenerator>()
+            .expect("There should be a object")
+            .downcast::<AsyncGenerator>()
             .expect("must be async generator");
 
-        gen.state = AsyncGeneratorState::Completed;
-        gen.context = None;
+        let mut gen = generator.borrow_mut();
 
-        let next = gen.queue.pop_front().expect("must have item in queue");
-        drop(gen);
+        gen.data.state = AsyncGeneratorState::Completed;
+        gen.data.context = None;
+
+        let next = gen.data.queue.pop_front().expect("must have item in queue");
 
         let return_value = context.vm.get_return_value();
         context.vm.set_return_value(JsValue::undefined());
 
-        if let Some(error) = context.vm.pending_exception.take() {
-            AsyncGenerator::complete_step(&next, Err(error), true, None, context);
-        } else {
-            AsyncGenerator::complete_step(&next, Ok(return_value), true, None, context);
-        }
-        AsyncGenerator::drain_queue(&async_generator_object, context);
+        let completion = context
+            .vm
+            .pending_exception
+            .take()
+            .map_or(Ok(return_value), Err);
+
+        drop(gen);
+
+        AsyncGenerator::complete_step(&next, completion, true, None, context);
+        // TODO: Upgrade to the latest spec when the problem is fixed.
+        AsyncGenerator::resume_next(&generator, context);
 
         Ok(CompletionType::Normal)
     }

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

@@ -44,22 +44,28 @@ impl Operation for AsyncGeneratorYield {
             .async_generator_object(&context.vm.stack)
             .expect("`AsyncGeneratorYield` must only be called inside async generators");
         let completion = Ok(value);
+        let async_generator_object = async_generator_object
+            .downcast::<AsyncGenerator>()
+            .expect("must be async generator object");
         let next = async_generator_object
-            .downcast_mut::<AsyncGenerator>()
-            .expect("must be async generator object")
+            .borrow_mut()
+            .data
             .queue
             .pop_front()
             .expect("must have item in queue");
 
         // 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_generator_object.borrow_mut();
-        let gen = generator_object_mut
-            .downcast_mut::<AsyncGenerator>()
-            .expect("must be async generator object");
+        // TODO: Upgrade to the latest spec when the problem is fixed.
+        let mut gen = async_generator_object.borrow_mut();
+        if gen.data.state == AsyncGeneratorState::Executing {
+            let Some(next) = gen.data.queue.front() else {
+                gen.data.state = AsyncGeneratorState::SuspendedYield;
+                context.vm.set_return_value(JsValue::undefined());
+                return Ok(CompletionType::Yield);
+            };
 
-        if let Some(next) = gen.queue.front() {
             let resume_kind = match next.completion.clone() {
                 CompletionRecord::Normal(val) => {
                     context.vm.push(val);
@@ -81,7 +87,14 @@ impl Operation for AsyncGeneratorYield {
             return Ok(CompletionType::Normal);
         }
 
-        gen.state = AsyncGeneratorState::SuspendedYield;
+        assert!(matches!(
+            gen.data.state,
+            AsyncGeneratorState::AwaitingReturn | AsyncGeneratorState::Completed
+        ));
+
+        AsyncGenerator::resume_next(&async_generator_object, context);
+
+        async_generator_object.borrow_mut().data.state = AsyncGeneratorState::SuspendedYield;
         context.vm.set_return_value(JsValue::undefined());
         Ok(CompletionType::Yield)
     }