feat(transformer/async-to-generator): handle arrow-function correctly (#6640)

Author: Dunqing

crates/oxc_transformer/src/es2017/async_to_generator.rs

@@ -49,7 +49,8 @@ use oxc_ast::{
     },
     NONE,
 };
-use oxc_span::SPAN;
+use oxc_semantic::ScopeFlags;
+use oxc_span::{GetSpan, SPAN};
 use oxc_traverse::{Ancestor, Traverse, TraverseCtx};
 
 use crate::{common::helper_loader::Helper, context::TransformCtx};
@@ -102,6 +103,11 @@ impl<'a, 'ctx> Traverse<'a> for AsyncToGenerator<'a, 'ctx> {
             }
             let new_function = self.transform_function(func, ctx);
             *expr = ctx.ast.expression_from_function(new_function);
+        } else if let Expression::ArrowFunctionExpression(arrow) = expr {
+            if !arrow.r#async {
+                return;
+            }
+            *expr = self.transform_arrow_function(arrow, ctx);
         }
     }
 
@@ -134,50 +140,6 @@ impl<'a, 'ctx> Traverse<'a> for AsyncToGenerator<'a, 'ctx> {
             }
         }
     }
-
-    fn exit_arrow_function_expression(
-        &mut self,
-        arrow: &mut ArrowFunctionExpression<'a>,
-        ctx: &mut TraverseCtx<'a>,
-    ) {
-        if !arrow.r#async {
-            return;
-        }
-        let body = ctx.ast.function_body(
-            SPAN,
-            ctx.ast.move_vec(&mut arrow.body.directives),
-            ctx.ast.move_vec(&mut arrow.body.statements),
-        );
-        let target = ctx.ast.function(
-            FunctionType::FunctionExpression,
-            SPAN,
-            None,
-            true,
-            false,
-            false,
-            arrow.type_parameters.take(),
-            NONE,
-            ctx.ast.alloc(ctx.ast.formal_parameters(
-                SPAN,
-                arrow.params.kind,
-                ctx.ast.move_vec(&mut arrow.params.items),
-                arrow.params.rest.take(),
-            )),
-            arrow.return_type.take(),
-            Some(body),
-        );
-        let parameters =
-            ctx.ast.vec1(ctx.ast.argument_expression(ctx.ast.expression_from_function(target)));
-        let call = self.ctx.helper_call_expr(Helper::AsyncToGenerator, parameters, ctx);
-        let body = ctx.ast.function_body(
-            SPAN,
-            ctx.ast.vec(),
-            ctx.ast.vec1(ctx.ast.statement_expression(SPAN, call)),
-        );
-        arrow.body = ctx.ast.alloc(body);
-        arrow.r#async = false;
-        arrow.expression = true;
-    }
 }
 
 impl<'a, 'ctx> AsyncToGenerator<'a, 'ctx> {
@@ -226,4 +188,46 @@ impl<'a, 'ctx> AsyncToGenerator<'a, 'ctx> {
             Some(body),
         )
     }
+
+    fn transform_arrow_function(
+        &self,
+        arrow: &mut ArrowFunctionExpression<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) -> Expression<'a> {
+        let mut body = ctx.ast.function_body(
+            SPAN,
+            ctx.ast.move_vec(&mut arrow.body.directives),
+            ctx.ast.move_vec(&mut arrow.body.statements),
+        );
+
+        // If the arrow's expression is true, we need to wrap the only one expression with return statement.
+        if arrow.expression {
+            let statement = body.statements.first_mut().unwrap();
+            let expression = match statement {
+                Statement::ExpressionStatement(es) => ctx.ast.move_expression(&mut es.expression),
+                _ => unreachable!(),
+            };
+            *statement = ctx.ast.statement_return(expression.span(), Some(expression));
+        }
+
+        let r#type = FunctionType::FunctionExpression;
+        let parameters = ctx.ast.alloc(ctx.ast.formal_parameters(
+            SPAN,
+            arrow.params.kind,
+            ctx.ast.move_vec(&mut arrow.params.items),
+            arrow.params.rest.take(),
+        ));
+        let body = Some(body);
+        let mut function = ctx
+            .ast
+            .function(r#type, SPAN, None, true, false, false, NONE, NONE, parameters, NONE, body);
+        function.scope_id = arrow.scope_id.clone();
+        if let Some(scope_id) = function.scope_id.get() {
+            ctx.scopes_mut().get_flags_mut(scope_id).remove(ScopeFlags::Arrow);
+        }
+
+        let arguments =
+            ctx.ast.vec1(ctx.ast.argument_expression(ctx.ast.expression_from_function(function)));
+        self.ctx.helper_call_expr(Helper::AsyncToGenerator, arguments, ctx)
+    }
 }

crates/oxc_transformer/src/es2017/mod.rs

@@ -1,4 +1,4 @@
-use oxc_ast::ast::{ArrowFunctionExpression, Expression, Statement};
+use oxc_ast::ast::{Expression, Statement};
 use oxc_traverse::{Traverse, TraverseCtx};
 
 use crate::context::TransformCtx;
@@ -34,14 +34,4 @@ impl<'a, 'ctx> Traverse<'a> for ES2017<'a, 'ctx> {
             self.async_to_generator.exit_statement(stmt, ctx);
         }
     }
-
-    fn exit_arrow_function_expression(
-        &mut self,
-        node: &mut ArrowFunctionExpression<'a>,
-        ctx: &mut TraverseCtx<'a>,
-    ) {
-        if self.options.async_to_generator {
-            self.async_to_generator.exit_arrow_function_expression(node, ctx);
-        }
-    }
 }

crates/oxc_transformer/src/lib.rs

@@ -333,7 +333,6 @@ impl<'a, 'ctx> Traverse<'a> for TransformerImpl<'a, 'ctx> {
                 .push(ctx.ast.statement_return(SPAN, Some(statement.unbox().expression)));
             arrow.expression = false;
         }
-        self.x2_es2017.exit_arrow_function_expression(arrow, ctx);
     }
 
     fn exit_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) {

tasks/transform_conformance/snapshots/babel.snap.md

@@ -1,6 +1,6 @@
 commit: 3bcfee23
 
-Passed: 341/1051
+Passed: 342/1051
 
 # All Passed:
 * babel-plugin-transform-logical-assignment-operators
@@ -1628,7 +1628,7 @@ x Output mismatch
 x Output mismatch
 
 
-# babel-plugin-transform-async-to-generator (1/24)
+# babel-plugin-transform-async-to-generator (2/24)
 * assumption-ignoreFunctionLength-true/basic/input.mjs
 x Output mismatch
 
@@ -1677,9 +1677,6 @@ x Output mismatch
 * regression/15978/input.js
 x Output mismatch
 
-* regression/4599/input.js
-x Output mismatch
-
 * regression/8783/input.js
 x Output mismatch
 

tasks/transform_conformance/snapshots/babel_exec.snap.md

@@ -1,6 +1,6 @@
 commit: 3bcfee23
 
-Passed: 35/62
+Passed: 34/62
 
 # All Passed:
 * babel-plugin-transform-logical-assignment-operators
@@ -73,7 +73,7 @@ exec failed
 exec failed
 
 
-# babel-plugin-transform-async-to-generator (2/6)
+# babel-plugin-transform-async-to-generator (1/6)
 * regression/15978/exec.js
 exec failed
 
@@ -83,6 +83,9 @@ exec failed
 * regression/T6882/exec.js
 exec failed
 
+* regression/fn-name/exec.js
+exec failed
+
 * regression/test262-fn-length/exec.js
 exec failed