feat(transformer): async-to-generator plugin. (#5590)

Tests are still not passed. A lot need to do yet.

---------

Co-authored-by: Dunqing <dengqing0821@gmail.com>

crates/oxc_transformer/src/es2017/async_to_generator.rs

@@ -0,0 +1,233 @@
+//! ES2017: Async / Await \[WIP\]
+//!
+//! This plugin transforms async functions to generator functions.
+//!
+//! ## Example
+//!
+//! Input:
+//! ```js
+//! async function foo() {
+//!   await bar();
+//! }
+//! const foo2 = async () => {
+//!   await bar();
+//! };
+//! async () => {
+//!   await bar();
+//! }
+//! ```
+//!
+//! Output (Currently):
+//! ```js
+//! function foo() {
+//!   return _asyncToGenerator(function* () {
+//!     yield bar();
+//!   })
+//! }
+//! const foo2 = () => _asyncToGenerator(function* () {
+//!   yield bar();
+//! }
+//! ```
+//!
+//! ## Implementation
+//!
+//! Implementation based on [@babel/plugin-transform-async-to-generator](https://babel.dev/docs/babel-plugin-transform-async-to-generator).
+//!
+//!
+//! Reference:
+//! * Babel docs: <https://babeljs.io/docs/en/babel-plugin-transform-async-to-generator>
+//! * Esbuild implementation: <https://github.com/evanw/esbuild/blob/main/internal/js_parser/js_parser_lower.go#L392>
+//! * Babel implementation: <https://github.com/babel/babel/blob/main/packages/babel-plugin-transform-async-to-generator>
+//! * Babel helper implementation: <https://github.com/babel/babel/blob/main/packages/babel-helper-remap-async-to-generator>
+//! * Async / Await TC39 proposal: <https://github.com/tc39/proposal-async-await>
+//!
+
+use oxc_ast::ast::{
+    ArrowFunctionExpression, Expression, Function, FunctionType, Statement, VariableDeclarationKind,
+};
+use oxc_ast::NONE;
+use oxc_span::{Atom, SPAN};
+use oxc_syntax::reference::ReferenceFlags;
+use oxc_syntax::symbol::SymbolId;
+use oxc_traverse::{Ancestor, Traverse, TraverseCtx};
+
+pub struct AsyncToGenerator;
+
+impl AsyncToGenerator {
+    fn get_helper_callee<'a>(
+        symbol_id: Option<SymbolId>,
+        ctx: &mut TraverseCtx<'a>,
+    ) -> Expression<'a> {
+        let ident = ctx.create_reference_id(
+            SPAN,
+            Atom::from("babelHelpers"),
+            symbol_id,
+            ReferenceFlags::Read,
+        );
+        let object = ctx.ast.expression_from_identifier_reference(ident);
+        let property = ctx.ast.identifier_name(SPAN, Atom::from("asyncToGenerator"));
+        Expression::from(ctx.ast.member_expression_static(SPAN, object, property, false))
+    }
+
+    fn transform_function<'a>(func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) -> Function<'a> {
+        let babel_helpers_id = ctx.scopes().find_binding(ctx.current_scope_id(), "babelHelpers");
+        let callee = Self::get_helper_callee(babel_helpers_id, ctx);
+        let target = ctx.ast.function(
+            func.r#type,
+            SPAN,
+            None,
+            true,
+            false,
+            false,
+            func.type_parameters.take(),
+            func.this_param.take(),
+            ctx.ast.alloc(ctx.ast.formal_parameters(
+                SPAN,
+                func.params.kind,
+                ctx.ast.move_vec(&mut func.params.items),
+                func.params.rest.take(),
+            )),
+            func.return_type.take(),
+            func.body.take(),
+        );
+        let parameters =
+            ctx.ast.vec1(ctx.ast.argument_expression(ctx.ast.expression_from_function(target)));
+        let call = ctx.ast.expression_call(SPAN, callee, NONE, parameters, false);
+        let returns = ctx.ast.return_statement(SPAN, Some(call));
+        let body = Statement::ReturnStatement(ctx.ast.alloc(returns));
+        let body = ctx.ast.function_body(SPAN, ctx.ast.vec(), ctx.ast.vec1(body));
+        let body = ctx.ast.alloc(body);
+        let params = ctx.ast.formal_parameters(SPAN, func.params.kind, ctx.ast.vec(), NONE);
+        ctx.ast.function(
+            FunctionType::FunctionExpression,
+            SPAN,
+            None,
+            false,
+            false,
+            false,
+            func.type_parameters.take(),
+            func.this_param.take(),
+            params,
+            func.return_type.take(),
+            Some(body),
+        )
+    }
+}
+
+impl<'a> Traverse<'a> for AsyncToGenerator {
+    fn exit_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
+        if let Expression::AwaitExpression(await_expr) = expr {
+            // Do not transform top-level await, or in async generator functions.
+            let in_async_function = ctx
+                .ancestry
+                .ancestors()
+                .find_map(|ance| {
+                    // We need to check if there's async generator or async function.
+                    // If it is async generator, we should not transform the await expression here.
+                    if let Ancestor::FunctionBody(body) = ance {
+                        if *body.r#async() {
+                            Some(!body.generator())
+                        } else {
+                            None
+                        }
+                    } else if let Ancestor::ArrowFunctionExpressionBody(_) = ance {
+                        // Arrow function is never generator.
+                        Some(true)
+                    } else {
+                        None
+                    }
+                })
+                .unwrap_or(false);
+            if in_async_function {
+                // Move the expression to yield.
+                *expr = ctx.ast.expression_yield(
+                    SPAN,
+                    false,
+                    Some(ctx.ast.move_expression(&mut await_expr.argument)),
+                );
+            }
+        } else if let Expression::FunctionExpression(func) = expr {
+            if !func.r#async || func.generator {
+                return;
+            }
+            let new_function = Self::transform_function(func, ctx);
+            *expr = ctx.ast.expression_from_function(new_function);
+        }
+    }
+
+    fn exit_statement(&mut self, stmt: &mut Statement<'a>, ctx: &mut TraverseCtx<'a>) {
+        if let Statement::FunctionDeclaration(func) = stmt {
+            if !func.r#async || func.generator {
+                return;
+            }
+            let new_function = Self::transform_function(func, ctx);
+            if let Some(id) = func.id.take() {
+                *stmt = ctx.ast.statement_declaration(ctx.ast.declaration_variable(
+                    SPAN,
+                    VariableDeclarationKind::Const,
+                    ctx.ast.vec1(ctx.ast.variable_declarator(
+                        SPAN,
+                        VariableDeclarationKind::Const,
+                        ctx.ast.binding_pattern(
+                            ctx.ast.binding_pattern_kind_from_binding_identifier(id),
+                            NONE,
+                            false,
+                        ),
+                        Some(ctx.ast.expression_from_function(new_function)),
+                        false,
+                    )),
+                    false,
+                ));
+            } else {
+                *stmt =
+                    ctx.ast.statement_declaration(ctx.ast.declaration_from_function(new_function));
+            }
+        }
+    }
+
+    fn exit_arrow_function_expression(
+        &mut self,
+        arrow: &mut ArrowFunctionExpression<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        if !arrow.r#async {
+            return;
+        }
+        let babel_helpers_id = ctx.scopes().find_binding(ctx.current_scope_id(), "babelHelpers");
+        let callee = Self::get_helper_callee(babel_helpers_id, ctx);
+        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 = ctx.ast.expression_call(SPAN, callee, NONE, parameters, false);
+        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;
+    }
+}

crates/oxc_transformer/src/es2017/mod.rs

@@ -0,0 +1,45 @@
+mod async_to_generator;
+pub mod options;
+
+use crate::es2017::async_to_generator::AsyncToGenerator;
+use crate::es2017::options::ES2017Options;
+use oxc_ast::ast::{ArrowFunctionExpression, Expression, Statement};
+use oxc_traverse::{Traverse, TraverseCtx};
+
+#[allow(dead_code)]
+pub struct ES2017 {
+    options: ES2017Options,
+
+    // Plugins
+    async_to_generator: AsyncToGenerator,
+}
+
+impl ES2017 {
+    pub fn new(options: ES2017Options) -> ES2017 {
+        ES2017 { async_to_generator: AsyncToGenerator, options }
+    }
+}
+
+impl<'a> Traverse<'a> for ES2017 {
+    fn exit_expression(&mut self, node: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
+        if self.options.async_to_generator {
+            self.async_to_generator.exit_expression(node, ctx);
+        }
+    }
+
+    fn exit_statement(&mut self, stmt: &mut Statement<'a>, ctx: &mut TraverseCtx<'a>) {
+        if self.options.async_to_generator {
+            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/es2017/options.rs

@@ -0,0 +1,27 @@
+use crate::env::{can_enable_plugin, Versions};
+use serde::Deserialize;
+
+#[derive(Debug, Default, Clone, Deserialize)]
+#[serde(default, rename_all = "camelCase", deny_unknown_fields)]
+pub struct ES2017Options {
+    #[serde(skip)]
+    pub async_to_generator: bool,
+}
+
+impl ES2017Options {
+    pub fn with_async_to_generator(&mut self, enable: bool) -> &mut Self {
+        self.async_to_generator = enable;
+        self
+    }
+
+    #[must_use]
+    pub fn from_targets_and_bugfixes(targets: Option<&Versions>, bugfixes: bool) -> Self {
+        Self {
+            async_to_generator: can_enable_plugin(
+                "transform-async-to-generator",
+                targets,
+                bugfixes,
+            ),
+        }
+    }
+}

crates/oxc_transformer/src/lib.rs

@@ -19,6 +19,7 @@ mod options;
 mod env;
 mod es2015;
 mod es2016;
+mod es2017;
 mod es2018;
 mod es2019;
 mod es2020;
@@ -33,6 +34,7 @@ use std::path::Path;
 
 use common::Common;
 use es2016::ES2016;
+use es2017::ES2017;
 use es2018::ES2018;
 use es2019::ES2019;
 use es2020::ES2020;
@@ -95,6 +97,7 @@ impl<'a> Transformer<'a> {
             x2_es2019: ES2019::new(self.options.es2019),
             x2_es2018: ES2018::new(self.options.es2018, &self.ctx),
             x2_es2016: ES2016::new(self.options.es2016, &self.ctx),
+            x2_es2017: ES2017::new(self.options.es2017),
             x3_es2015: ES2015::new(self.options.es2015),
             x4_regexp: RegExp::new(self.options.regexp, &self.ctx),
             common: Common::new(&self.ctx),
@@ -113,6 +116,7 @@ struct TransformerImpl<'a, 'ctx> {
     x2_es2020: ES2020<'a, 'ctx>,
     x2_es2019: ES2019,
     x2_es2018: ES2018<'a, 'ctx>,
+    x2_es2017: ES2017,
     x2_es2016: ES2016<'a, 'ctx>,
     x3_es2015: ES2015<'a>,
     x4_regexp: RegExp<'a, 'ctx>,
@@ -196,6 +200,7 @@ impl<'a, 'ctx> Traverse<'a> for TransformerImpl<'a, 'ctx> {
 
     fn exit_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
         self.x1_react.exit_expression(expr, ctx);
+        self.x2_es2017.exit_expression(expr, ctx);
         self.x3_es2015.exit_expression(expr, ctx);
     }
 
@@ -230,6 +235,7 @@ impl<'a, 'ctx> Traverse<'a> for TransformerImpl<'a, 'ctx> {
     fn exit_function(&mut self, func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.exit_function(func, ctx);
         self.x1_react.exit_function(func, ctx);
+        self.x2_es2017.exit_function(func, ctx);
         self.x3_es2015.exit_function(func, ctx);
     }
 
@@ -326,6 +332,7 @@ 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>) {
@@ -334,6 +341,10 @@ impl<'a, 'ctx> Traverse<'a> for TransformerImpl<'a, 'ctx> {
         self.common.exit_statements(stmts, ctx);
     }
 
+    fn exit_statement(&mut self, stmt: &mut Statement<'a>, ctx: &mut TraverseCtx<'a>) {
+        self.x2_es2017.exit_statement(stmt, ctx);
+    }
+
     fn enter_tagged_template_expression(
         &mut self,
         expr: &mut TaggedTemplateExpression<'a>,