handle async always

unyt-org · Oct 2, 2024 · fbf4f3d · fbf4f3d
1 parent 2252334
commit fbf4f3d
Show file tree

Hide file tree

Showing 4 changed files with 92 additions and 21 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -535,3 +535,34 @@ test!(
     "#
 );
 
+test!(
+    Syntax::Es(EsSyntax {
+        jsx: true,
+        ..Default::default()
+    },),
+    |_| TransformVisitor,
+    t39,
+    r#"
+    <div>
+        {
+            await x(y)
+        }
+    </div>
+    "#
+);
+
+test!(
+    Syntax::Es(EsSyntax {
+        jsx: true,
+        ..Default::default()
+    },),
+    |_| TransformVisitor,
+    t40,
+    r#"
+    <div>
+        {
+            (async () => await x(y))()
+        }
+    </div>
+    "#
+);
diff --git a/src/visitor.rs b/src/visitor.rs
@@ -3,9 +3,9 @@ use swc_core::{
     common::{util::take::Take, SyntaxContext, DUMMY_SP},
     ecma::{
         ast::{
-            ArrowExpr, BlockStmt, BlockStmtOrExpr, CallExpr, Callee, Expr, ExprOrSpread, ExprStmt, FnDecl, Ident, JSXAttr, JSXAttrName, JSXAttrValue, JSXElement, JSXElementChild, JSXElementName, JSXEmptyExpr, JSXExpr, JSXExprContainer, JSXSpreadChild, Lit, MemberProp, Null, ObjectPatProp, Pat, ReturnStmt, Stmt, Str, VarDecl
+            ArrowExpr, AwaitExpr, BlockStmt, BlockStmtOrExpr, CallExpr, Callee, Expr, ExprOrSpread, ExprStmt, FnDecl, Ident, JSXAttr, JSXAttrName, JSXAttrValue, JSXElement, JSXElementChild, JSXElementName, JSXEmptyExpr, JSXExpr, JSXExprContainer, JSXSpreadChild, Lit, MemberProp, Null, ObjectPatProp, Pat, ReturnStmt, Stmt, Str, VarDecl
         },
-        visit::{Fold, FoldWith, Visit, VisitWith},
+        visit::{Fold, FoldWith, Visit, VisitMutWith, VisitWith},
     },
 };
 
@@ -127,6 +127,22 @@ impl Visit for VariableCollector {
 }
 
 
+struct AsyncChecker {
+    is_async: bool,
+}
+
+impl Visit for AsyncChecker {
+    fn visit_await_expr(&mut self, _node: &AwaitExpr) {
+        self.is_async = true;
+    }
+
+    // ignore if async inside function
+    fn visit_fn_decl(&mut self, _node: &FnDecl) {}
+    fn visit_arrow_expr(&mut self, _node: &ArrowExpr) {}
+    fn visit_fn_expr(&mut self, _node: &swc_core::ecma::ast::FnExpr) {}
+}
+
+
 const GLOBAL_THIS_ALIASES: [&'static str; 3] = [
     "globalThis",
     "self",
@@ -308,27 +324,45 @@ impl TransformVisitor {
             }
 
             // default: wrap in always
-            _ => Box::new(Expr::Call(CallExpr {
-                span: DUMMY_SP,
-                callee: Callee::Expr(Box::new(Expr::Ident(Ident::new(
-                    always_fn_name.into(),
-                    DUMMY_SP,
-                    Default::default(),
-                )))),
-                args: vec![Expr::Arrow(ArrowExpr {
+            _ => {
+                // check if body contains await -> is_async
+                let mut async_checker = AsyncChecker { is_async: false };
+                e.visit_with(&mut async_checker);
+                let is_async = async_checker.is_async;
+
+                let call_expr = Expr::Call(CallExpr {
                     span: DUMMY_SP,
-                    params: Take::dummy(),
-                    body: Box::new(BlockStmtOrExpr::Expr(e)),
-                    is_async: false,
-                    is_generator: false,
-                    type_params: Take::dummy(),
-                    return_type: Take::dummy(),
+                    callee: Callee::Expr(Box::new(Expr::Ident(Ident::new(
+                        always_fn_name.into(),
+                        DUMMY_SP,
+                        Default::default(),
+                    )))),
+                    args: vec![Expr::Arrow(ArrowExpr {
+                        span: DUMMY_SP,
+                        params: Take::dummy(),
+                        body: Box::new(BlockStmtOrExpr::Expr(e)),
+                        is_async,
+                        is_generator: false,
+                        type_params: Take::dummy(),
+                        return_type: Take::dummy(),
+                        ctxt: Default::default(),
+                    })
+                    .into()],
+                    type_args: Take::dummy(),
                     ctxt: Default::default(),
-                })
-                .into()],
-                type_args: Take::dummy(),
-                ctxt: Default::default(),
-            })),
+                });
+
+                // add await if async
+                if is_async {
+                    Box::new(Expr::Await(AwaitExpr {
+                        span: DUMMY_SP,
+                        arg: Box::new(call_expr)
+                    }))
+                }
+                else {
+                    Box::new(call_expr)
+                }
+            },
         }
     }
 

diff --git a/tests/__swc_snapshots__/src/lib.rs/t39.js b/tests/__swc_snapshots__/src/lib.rs/t39.js
@@ -0,0 +1,3 @@
+<div>
+        {await _$(async ()=>await x(y))}
+    </div>;
diff --git a/tests/__swc_snapshots__/src/lib.rs/t40.js b/tests/__swc_snapshots__/src/lib.rs/t40.js
@@ -0,0 +1,3 @@
+<div>
+        {_$(()=>(async ()=>await x(y))())}
+    </div>;