refactor(transformer/async-generator-functions): do not transform await expression where inside ArrowFunctionExpression (#7132)

The `await expr` should be transformed in AsyncToGenerator rather than AsyncGeneratorFunctions

crates/oxc_transformer/src/es2018/async_generator_functions/mod.rs

@@ -68,7 +68,6 @@ mod for_await;
 
 use oxc_allocator::GetAddress;
 use oxc_ast::ast::*;
-use oxc_data_structures::stack::Stack;
 use oxc_span::SPAN;
 use oxc_traverse::{Ancestor, Traverse, TraverseCtx};
 
@@ -76,17 +75,12 @@ use crate::{common::helper_loader::Helper, context::TransformCtx, es2017::AsyncG
 
 pub struct AsyncGeneratorFunctions<'a, 'ctx> {
     ctx: &'ctx TransformCtx<'a>,
-    stack: Stack<bool>,
     executor: AsyncGeneratorExecutor<'a, 'ctx>,
 }
 
 impl<'a, 'ctx> AsyncGeneratorFunctions<'a, 'ctx> {
     pub fn new(ctx: &'ctx TransformCtx<'a>) -> Self {
-        Self {
-            ctx,
-            executor: AsyncGeneratorExecutor::new(Helper::WrapAsyncGenerator, ctx),
-            stack: Stack::new(),
-        }
+        Self { ctx, executor: AsyncGeneratorExecutor::new(Helper::WrapAsyncGenerator, ctx) }
     }
 }
 
@@ -164,10 +158,6 @@ impl<'a, 'ctx> Traverse<'a> for AsyncGeneratorFunctions<'a, 'ctx> {
         }
     }
 
-    fn enter_function(&mut self, func: &mut Function<'a>, _ctx: &mut TraverseCtx<'a>) {
-        self.stack.push(func.r#async && func.generator);
-    }
-
     fn exit_function(&mut self, func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) {
         if func.r#async
             && func.generator
@@ -180,11 +170,29 @@ impl<'a, 'ctx> Traverse<'a> for AsyncGeneratorFunctions<'a, 'ctx> {
         {
             self.executor.transform_function_for_method_definition(func, ctx);
         }
-        self.stack.pop();
     }
 }
 
 impl<'a, 'ctx> AsyncGeneratorFunctions<'a, 'ctx> {
+    /// Check whether the current node is inside an async generator function.
+    fn is_inside_async_generator_function(ctx: &mut TraverseCtx<'a>) -> bool {
+        // Early return if current scope is top because we don't need to transform top-level await expression.
+        if ctx.current_scope_flags().is_top() {
+            return false;
+        }
+
+        for ancestor in ctx.ancestors() {
+            match ancestor {
+                Ancestor::FunctionBody(func) => return *func.r#async() && *func.generator(),
+                Ancestor::ArrowFunctionExpressionBody(_) => {
+                    return false;
+                }
+                _ => {}
+            }
+        }
+        false
+    }
+
     /// Transform `yield * argument` expression to `yield asyncGeneratorDelegate(asyncIterator(argument))`.
     #[allow(clippy::unused_self)]
     fn transform_yield_expression(
@@ -214,11 +222,7 @@ impl<'a, 'ctx> AsyncGeneratorFunctions<'a, 'ctx> {
         expr: &mut AwaitExpression<'a>,
         ctx: &mut TraverseCtx<'a>,
     ) -> Option<Expression<'a>> {
-        // We don't need to handle top-level await.
-        if ctx.parent().is_program() ||
-        // Check the function is async generator function
-        !self.stack.last().copied().unwrap_or(false)
-        {
+        if !Self::is_inside_async_generator_function(ctx) {
             return None;
         }
 

crates/oxc_transformer/src/es2018/mod.rs

@@ -62,12 +62,6 @@ impl<'a, 'ctx> Traverse<'a> for ES2018<'a, 'ctx> {
         }
     }
 
-    fn enter_function(&mut self, node: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) {
-        if self.options.async_generator_functions {
-            self.async_generator_functions.enter_function(node, ctx);
-        }
-    }
-
     fn exit_function(&mut self, node: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) {
         if self.options.async_generator_functions {
             self.async_generator_functions.exit_function(node, ctx);

crates/oxc_transformer/src/lib.rs

@@ -266,7 +266,6 @@ impl<'a, 'ctx> Traverse<'a> for TransformerImpl<'a, 'ctx> {
     }
 
     fn enter_function(&mut self, func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) {
-        self.x2_es2018.enter_function(func, ctx);
         self.common.enter_function(func, ctx);
     }
 

tasks/transform_conformance/snapshots/babel.snap.md

@@ -1,11 +1,12 @@
 commit: d20b314c
 
-Passed: 381/1078
+Passed: 382/1078
 
 # All Passed:
 * babel-plugin-transform-class-static-block
 * babel-plugin-transform-logical-assignment-operators
 * babel-plugin-transform-optional-catch-binding
+* babel-plugin-transform-async-generator-functions
 * babel-preset-react
 * babel-plugin-transform-react-display-name
 * babel-plugin-transform-react-jsx-self
@@ -1450,11 +1451,6 @@ x Output mismatch
 x Output mismatch
 
 
-# babel-plugin-transform-async-generator-functions (18/19)
-* nested/arrows-in-declaration/input.js
-x Output mismatch
-
-
 # babel-plugin-transform-object-rest-spread (5/59)
 * assumption-ignoreFunctionLength/parameters-object-rest-used-in-default/input.js
 x Output mismatch