CVE-2024-43367

@@ -138,21 +138,19 @@ impl AsyncGenerator {
                 .with_message("generator resumed on non generator object")
                 .into()
         });
-        let generator_object = if_abrupt_reject_promise!(result, promise_capability, context);
-        let result: JsResult<_> = generator_object.downcast_mut::<Self>().ok_or_else(|| {
+        let generator = if_abrupt_reject_promise!(result, promise_capability, context);
+        let result: JsResult<_> = generator.clone().downcast::<Self>().map_err(|_| {
             JsNativeError::typ()
                 .with_message("generator resumed on non generator object")
                 .into()
         });
-        let mut generator = if_abrupt_reject_promise!(result, promise_capability, context);
+        let generator = if_abrupt_reject_promise!(result, promise_capability, context);
 
         // 5. Let state be generator.[[AsyncGeneratorState]].
-        let state = generator.state;
+        let state = generator.borrow().data.state;
 
         // 6. If state is completed, then
         if state == AsyncGeneratorState::Completed {
-            drop(generator);
-
             // a. Let iteratorResult be CreateIterResultObject(undefined, true).
             let iterator_result = create_iter_result_object(JsValue::undefined(), true, context);
 
@@ -170,28 +168,7 @@ impl AsyncGenerator {
         let completion = CompletionRecord::Normal(args.get_or_undefined(0).clone());
 
         // 8. Perform AsyncGeneratorEnqueue(generator, completion, promiseCapability).
-        generator.enqueue(completion.clone(), promise_capability.clone());
-
-        // 9. If state is either suspendedStart or suspendedYield, then
-        if state == AsyncGeneratorState::SuspendedStart
-            || state == AsyncGeneratorState::SuspendedYield
-        {
-            // a. Perform AsyncGeneratorResume(generator, completion).
-            let generator_context = generator
-                .context
-                .take()
-                .expect("generator context cannot be empty here");
-
-            drop(generator);
-
-            Self::resume(
-                generator_object,
-                state,
-                generator_context,
-                completion,
-                context,
-            );
-        }
+        Self::enqueue(&generator, completion, promise_capability.clone(), context);
 
         // 11. Return promiseCapability.[[Promise]].
         Ok(promise_capability.promise().clone().into())
@@ -226,49 +203,19 @@ impl AsyncGenerator {
                 .into()
         });
         let generator_object = if_abrupt_reject_promise!(result, promise_capability, context);
-        let result: JsResult<_> = generator_object.downcast_mut::<Self>().ok_or_else(|| {
+        let result: JsResult<_> = generator_object.clone().downcast::<Self>().map_err(|_| {
             JsNativeError::typ()
                 .with_message("generator resumed on non generator object")
                 .into()
         });
-        let mut generator = if_abrupt_reject_promise!(result, promise_capability, context);
+        let generator = if_abrupt_reject_promise!(result, promise_capability, context);
 
         // 5. Let completion be Completion Record { [[Type]]: return, [[Value]]: value, [[Target]]: empty }.
         let return_value = args.get_or_undefined(0).clone();
-        let completion = CompletionRecord::Return(return_value.clone());
+        let completion = CompletionRecord::Return(return_value);
 
         // 6. Perform AsyncGeneratorEnqueue(generator, completion, promiseCapability).
-        generator.enqueue(completion.clone(), promise_capability.clone());
-
-        // 7. Let state be generator.[[AsyncGeneratorState]].
-        let state = generator.state;
-
-        // 8. If state is either suspendedStart or completed, then
-        if state == AsyncGeneratorState::SuspendedStart || state == AsyncGeneratorState::Completed {
-            // a. Set generator.[[AsyncGeneratorState]] to awaiting-return.
-            generator.state = AsyncGeneratorState::AwaitingReturn;
-
-            // b. Perform ! AsyncGeneratorAwaitReturn(generator).
-            drop(generator);
-            Self::await_return(generator_object.clone(), return_value, context);
-        }
-        // 9. Else if state is suspendedYield, then
-        else if state == AsyncGeneratorState::SuspendedYield {
-            // a. Perform AsyncGeneratorResume(generator, completion).
-            let generator_context = generator
-                .context
-                .take()
-                .expect("generator context cannot be empty here");
-
-            drop(generator);
-            Self::resume(
-                generator_object,
-                state,
-                generator_context,
-                completion,
-                context,
-            );
-        }
+        Self::enqueue(&generator, completion, promise_capability.clone(), context);
 
         // 11. Return promiseCapability.[[Promise]].
         Ok(promise_capability.promise().clone().into())
@@ -303,30 +250,31 @@ impl AsyncGenerator {
                 .into()
         });
         let generator_object = if_abrupt_reject_promise!(result, promise_capability, context);
-        let result: JsResult<_> = generator_object.downcast_mut::<Self>().ok_or_else(|| {
+        let result: JsResult<_> = generator_object.clone().downcast::<Self>().map_err(|_| {
             JsNativeError::typ()
                 .with_message("generator resumed on non generator object")
                 .into()
         });
-        let mut generator = if_abrupt_reject_promise!(result, promise_capability, context);
+        let generator = if_abrupt_reject_promise!(result, promise_capability, context);
+        let mut gen = generator.borrow_mut();
 
         // 5. Let state be generator.[[AsyncGeneratorState]].
-        let mut state = generator.state;
+        let mut state = gen.data.state;
 
         // 6. If state is suspendedStart, then
         if state == AsyncGeneratorState::SuspendedStart {
             // a. Set generator.[[AsyncGeneratorState]] to completed.
-            generator.state = AsyncGeneratorState::Completed;
-            generator.context = None;
+            gen.data.state = AsyncGeneratorState::Completed;
+            gen.data.context = None;
 
             // b. Set state to completed.
             state = AsyncGeneratorState::Completed;
         }
 
+        drop(gen);
+
         // 7. If state is completed, then
         if state == AsyncGeneratorState::Completed {
-            drop(generator);
-
             // a. Perform ! Call(promiseCapability.[[Reject]], undefined, « exception »).
             promise_capability
                 .reject()
@@ -346,25 +294,12 @@ impl AsyncGenerator {
             CompletionRecord::Throw(JsError::from_opaque(args.get_or_undefined(0).clone()));
 
         // 9. Perform AsyncGeneratorEnqueue(generator, completion, promiseCapability).
-        generator.enqueue(completion.clone(), promise_capability.clone());
-
-        // 10. If state is suspendedYield, then
-        if state == AsyncGeneratorState::SuspendedYield {
-            let generator_context = generator
-                .context
-                .take()
-                .expect("generator context cannot be empty here");
-            drop(generator);
-
-            // a. Perform AsyncGeneratorResume(generator, completion).
-            Self::resume(
-                generator_object,
-                state,
-                generator_context,
-                completion,
-                context,
-            );
-        }
+        Self::enqueue(
+            &generator,
+            completion.clone(),
+            promise_capability.clone(),
+            context,
+        );
 
         // 12. Return promiseCapability.[[Promise]].
         Ok(promise_capability.promise().clone().into())
@@ -377,18 +312,27 @@ impl AsyncGenerator {
     ///
     /// [spec]: https://tc39.es/ecma262/#sec-asyncgeneratorenqueue
     pub(crate) fn enqueue(
-        &mut self,
+        generator: &JsObject<AsyncGenerator>,
         completion: CompletionRecord,
         promise_capability: PromiseCapability,
+        context: &mut Context,
     ) {
+        let mut gen = generator.borrow_mut();
         // 1. Let request be AsyncGeneratorRequest { [[Completion]]: completion, [[Capability]]: promiseCapability }.
         let request = AsyncGeneratorRequest {
             completion,
             capability: promise_capability,
         };
 
         // 2. Append request to the end of generator.[[AsyncGeneratorQueue]].
-        self.queue.push_back(request);
+        gen.data.queue.push_back(request);
+
+        // Patch that mirrors https://262.ecma-international.org/12.0/#sec-asyncgeneratorenqueue
+        // This resolves the return bug.
+        if gen.data.state != AsyncGeneratorState::Executing {
+            drop(gen);
+            AsyncGenerator::resume_next(generator, context);
+        }
     }
 
     /// `AsyncGeneratorCompleteStep ( generator, completion, done [ , realm ] )`
@@ -453,67 +397,17 @@ impl AsyncGenerator {
         }
     }
 
-    /// `AsyncGeneratorResume ( generator, completion )`
-    ///
-    /// More information:
-    ///  - [ECMAScript reference][spec]
-    ///
-    /// [spec]: https://tc39.es/ecma262/#sec-asyncgeneratorresume
-    pub(crate) fn resume(
-        generator: &JsObject,
-        state: AsyncGeneratorState,
-        mut generator_context: GeneratorContext,
-        completion: CompletionRecord,
-        context: &mut Context,
-    ) {
-        // 1. Assert: generator.[[AsyncGeneratorState]] is either suspendedStart or suspendedYield.
-        assert!(
-            state == AsyncGeneratorState::SuspendedStart
-                || state == AsyncGeneratorState::SuspendedYield
-        );
-
-        // 2. Let genContext be generator.[[AsyncGeneratorContext]].
-
-        // 3. Let callerContext be the running execution context.
-        // 4. Suspend callerContext.
-
-        // 5. Set generator.[[AsyncGeneratorState]] to executing.
-        generator
-            .downcast_mut::<Self>()
-            .expect("already checked before")
-            .state = AsyncGeneratorState::Executing;
-
-        let (value, resume_kind) = match completion {
-            CompletionRecord::Normal(val) => (val, GeneratorResumeKind::Normal),
-            CompletionRecord::Return(val) => (val, GeneratorResumeKind::Return),
-            CompletionRecord::Throw(err) => (err.to_opaque(context), GeneratorResumeKind::Throw),
-        };
-        // 6. Push genContext onto the execution context stack; genContext is now the running execution context.
-
-        let result = generator_context.resume(Some(value), resume_kind, context);
-
-        // 7. Resume the suspended evaluation of genContext using completion as the result of the operation that suspended it. Let result be the Completion Record returned by the resumed computation.
-
-        generator
-            .downcast_mut::<Self>()
-            .expect("already checked before")
-            .context = Some(generator_context);
-
-        // 8. Assert: result is never an abrupt completion.
-        assert!(!result.is_throw_completion());
-
-        // 9. Assert: When we return here, genContext has already been removed from the execution context stack and
-        // callerContext is the currently running execution context.
-        // 10. Return unused.
-    }
-
     /// `AsyncGeneratorAwaitReturn ( generator )`
     ///
     /// More information:
     ///  - [ECMAScript reference][spec]
     ///
     /// [spec]: https://tc39.es/ecma262/#sec-asyncgeneratorawaitreturn
-    pub(crate) fn await_return(generator: JsObject, value: JsValue, context: &mut Context) {
+    pub(crate) fn await_return(
+        generator: JsObject<AsyncGenerator>,
+        value: JsValue,
+        context: &mut Context,
+    ) {
         // 1. Let queue be generator.[[AsyncGeneratorQueue]].
         // 2. Assert: queue is not empty.
         // 3. Let next be the first element of queue.
@@ -532,15 +426,14 @@ impl AsyncGenerator {
         let promise = match promise_completion {
             Ok(value) => value,
             Err(value) => {
-                let mut gen = generator
-                    .downcast_mut::<Self>()
-                    .expect("already checked before");
-                gen.state = AsyncGeneratorState::Completed;
-                gen.context = None;
-                let next = gen.queue.pop_front().expect("queue must not be empty");
-                drop(gen);
+                let next = {
+                    let mut gen = generator.borrow_mut();
+                    gen.data.state = AsyncGeneratorState::Completed;
+                    gen.data.context = None;
+                    gen.data.queue.pop_front().expect("queue must not be empty")
+                };
                 Self::complete_step(&next, Err(value), true, None, context);
-                Self::drain_queue(&generator, context);
+                Self::resume_next(&generator, context);
                 return;
             }
         };
@@ -552,15 +445,13 @@ impl AsyncGenerator {
             NativeFunction::from_copy_closure_with_captures(
                 |_this, args, generator, context| {
                     let next = {
-                        let mut gen = generator
-                            .downcast_mut::<Self>()
-                            .expect("already checked before");
+                        let mut gen = generator.borrow_mut();
 
                         // a. Set generator.[[AsyncGeneratorState]] to completed.
-                        gen.state = AsyncGeneratorState::Completed;
-                        gen.context = None;
+                        gen.data.state = AsyncGeneratorState::Completed;
+                        gen.data.context = None;
 
-                        gen.queue.pop_front().expect("must have one entry")
+                        gen.data.queue.pop_front().expect("must have one entry")
                     };
 
                     // b. Let result be NormalCompletion(value).
@@ -570,7 +461,7 @@ impl AsyncGenerator {
                     Self::complete_step(&next, result, true, None, context);
 
                     // d. Perform AsyncGeneratorDrainQueue(generator).
-                    Self::drain_queue(generator, context);
+                    Self::resume_next(generator, context);
 
                     // e. Return undefined.
                     Ok(JsValue::undefined())
@@ -588,24 +479,24 @@ impl AsyncGenerator {
             context.realm(),
             NativeFunction::from_copy_closure_with_captures(
                 |_this, args, generator, context| {
-                    let mut gen = generator
-                        .downcast_mut::<Self>()
-                        .expect("already checked before");
+                    let next = {
+                        let mut gen = generator.borrow_mut();
 
-                    // a. Set generator.[[AsyncGeneratorState]] to completed.
-                    gen.state = AsyncGeneratorState::Completed;
-                    gen.context = None;
+                        // a. Set generator.[[AsyncGeneratorState]] to completed.
+                        gen.data.state = AsyncGeneratorState::Completed;
+                        gen.data.context = None;
+
+                        gen.data.queue.pop_front().expect("must have one entry")
+                    };
 
                     // b. Let result be ThrowCompletion(reason).
                     let result = Err(JsError::from_opaque(args.get_or_undefined(0).clone()));
 
                     // c. Perform AsyncGeneratorCompleteStep(generator, result, true).
-                    let next = gen.queue.pop_front().expect("must have one entry");
-                    drop(gen);
                     Self::complete_step(&next, result, true, None, context);
 
                     // d. Perform AsyncGeneratorDrainQueue(generator).
-                    Self::drain_queue(generator, context);
+                    Self::resume_next(generator, context);
 
                     // e. Return undefined.
                     Ok(JsValue::undefined())
@@ -627,64 +518,135 @@ impl AsyncGenerator {
         );
     }
 
-    /// `AsyncGeneratorDrainQueue ( generator )`
+    /// [`AsyncGeneratorResumeNext ( generator )`][spec]
     ///
-    /// More information:
-    ///  - [ECMAScript reference][spec]
-    ///
-    /// [spec]: https://tc39.es/ecma262/#sec-asyncgeneratordrainqueue
-    pub(crate) fn drain_queue(generator: &JsObject, context: &mut Context) {
-        let mut gen = generator
-            .downcast_mut::<Self>()
-            .expect("already checked before");
-
-        // 1. Assert: generator.[[AsyncGeneratorState]] is completed.
-        assert_eq!(gen.state, AsyncGeneratorState::Completed);
-
-        // 2. Let queue be generator.[[AsyncGeneratorQueue]].
-        let queue = &mut gen.queue;
-
-        // 3. If queue is empty, return unused.
-        if queue.is_empty() {
-            return;
+    /// [spec]: https://262.ecma-international.org/12.0/#sec-asyncgeneratorresumenext
+    pub(crate) fn resume_next(generator: &JsObject<AsyncGenerator>, context: &mut Context) {
+        // 1. Assert: generator is an AsyncGenerator instance.
+        let mut gen = generator.borrow_mut();
+        // 2. Let state be generator.[[AsyncGeneratorState]].
+        match gen.data.state {
+            // 3. Assert: state is not executing.
+            AsyncGeneratorState::Executing => panic!("3. Assert: state is not executing."),
+            // 4. If state is awaiting-return, return undefined.
+            AsyncGeneratorState::AwaitingReturn => return,
+            _ => {}
         }
 
-        // 4. Let done be false.
-        // 5. Repeat, while done is false,
-        loop {
-            // a. Let next be the first element of queue.
-            let next = queue.front().expect("must have entry");
-
-            // b. Let completion be Completion(next.[[Completion]]).
-            match next.completion.clone() {
-                // c. If completion.[[Type]] is return, then
-                CompletionRecord::Return(val) => {
-                    // i. Set generator.[[AsyncGeneratorState]] to awaiting-return.
-                    gen.state = AsyncGeneratorState::AwaitingReturn;
+        // 5. Let queue be generator.[[AsyncGeneratorQueue]].
+        // 6. If queue is an empty List, return undefined.
+        // 7. Let next be the value of the first element of queue.
+        // 8. Assert: next is an AsyncGeneratorRequest record.
+        let Some(next) = gen.data.queue.front() else {
+            return;
+        };
+        // 9. Let completion be next.[[Completion]].
+        let completion = &next.completion;
+
+        match (completion, gen.data.state) {
+            // 11. Else if state is completed, return ! AsyncGeneratorResolve(generator, undefined, true).
+            (CompletionRecord::Normal(_), s) => {
+                if s == AsyncGeneratorState::Completed {
+                    let next = gen
+                        .data
+                        .queue
+                        .pop_front()
+                        .expect("already have a reference to the front");
                     drop(gen);
+                    AsyncGenerator::complete_step(
+                        &next,
+                        Ok(JsValue::undefined()),
+                        true,
+                        None,
+                        context,
+                    );
+                    return AsyncGenerator::resume_next(generator, context);
+                }
+            }
+            // b. If state is completed, then
+            //    i. If completion.[[Type]] is return, then
+            (
+                CompletionRecord::Return(val),
+                AsyncGeneratorState::SuspendedStart | AsyncGeneratorState::Completed,
+            ) => {
+                let val = val.clone();
+                // 1. Set generator.[[AsyncGeneratorState]] to awaiting-return.
+                gen.data.state = AsyncGeneratorState::AwaitingReturn;
+                drop(gen);
 
-                    // ii. Perform ! AsyncGeneratorAwaitReturn(generator).
-                    Self::await_return(generator.clone(), val, context);
+                // Steps 2-11 are superseeded by `AsyncGeneratorAwaitReturn`
+                AsyncGenerator::await_return(generator.clone(), val, context);
 
-                    // iii. Set done to true.
-                    break;
-                }
-                // d. Else,
-                completion => {
-                    // i. If completion.[[Type]] is normal, then
-                    // 1. Set completion to NormalCompletion(undefined).
-                    let completion = completion.consume().map(|_| JsValue::undefined());
-
-                    // ii. Perform AsyncGeneratorCompleteStep(generator, completion, true).
-                    let next = queue.pop_front().expect("must have entry");
-                    Self::complete_step(&next, completion, true, None, context);
-
-                    // iii. If queue is empty, set done to true.
-                    if queue.is_empty() {
-                        break;
-                    }
-                }
+                // 12. Return undefined.
+                return;
             }
+            // ii. Else,
+            (
+                CompletionRecord::Throw(e),
+                AsyncGeneratorState::SuspendedStart | AsyncGeneratorState::Completed,
+            ) => {
+                let e = e.clone();
+                // 1. Assert: completion.[[Type]] is throw.
+                // 2. Perform ! AsyncGeneratorReject(generator, completion.[[Value]]).
+                gen.data.state = AsyncGeneratorState::Completed;
+
+                let next = gen
+                    .data
+                    .queue
+                    .pop_front()
+                    .expect("already have a reference to the front");
+                drop(gen);
+                AsyncGenerator::complete_step(&next, Err(e), true, None, context);
+                // 3. Return undefined.
+                return AsyncGenerator::resume_next(generator, context);
+            }
+            _ => {}
         }
+
+        // 12. Assert: state is either suspendedStart or suspendedYield.
+        assert!(matches!(
+            gen.data.state,
+            AsyncGeneratorState::SuspendedStart | AsyncGeneratorState::SuspendedYield
+        ));
+
+        let completion = completion.clone();
+
+        // 16. Set generator.[[AsyncGeneratorState]] to executing.
+        gen.data.state = AsyncGeneratorState::Executing;
+
+        // 13. Let genContext be generator.[[AsyncGeneratorContext]].
+        let mut generator_context = gen
+            .data
+            .context
+            .take()
+            .expect("generator context cannot be empty here");
+
+        drop(gen);
+
+        let (value, resume_kind) = match completion {
+            CompletionRecord::Normal(val) => (val, GeneratorResumeKind::Normal),
+            CompletionRecord::Return(val) => (val, GeneratorResumeKind::Return),
+            CompletionRecord::Throw(err) => (err.to_opaque(context), GeneratorResumeKind::Throw),
+        };
+
+        // 14. Let callerContext be the running execution context.
+        // 15. Suspend callerContext.
+        // 17. Push genContext onto the execution context stack; genContext is now the running execution context.
+        // 18. Resume the suspended evaluation of genContext using completion as the result of the operation that suspended it.
+        //     Let result be the completion record returned by the resumed computation.
+        let result = generator_context.resume(Some(value), resume_kind, context);
+
+        // 19. Assert: result is never an abrupt completion.
+        assert!(!matches!(result, CompletionRecord::Throw(_)));
+
+        generator
+            .borrow_mut()
+            .data
+            .context
+            .get_or_insert(generator_context);
+
+        // 20. Assert: When we return here, genContext has already been removed from the execution context stack and
+        //     callerContext is the currently running execution context.
+        // 21. Return undefined.
     }
 }

@@ -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)),
+    ]);
+}

@@ -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;

@@ -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;

@@ -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):

@@ -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();

@@ -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)
     }

@@ -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)
     }