Improve swc transforms (#45083)

x-ref: https://vercel.slack.com/archives/C02HY34AKME/p1674048645326239

## Bug

- [ ] Related issues linked using `fixes #number`
- [ ] Integration tests added
- [ ] Errors have a helpful link attached, see [`contributing.md`](https://github.com/vercel/next.js/blob/canary/contributing.md)

## Feature

- [ ] Implements an existing feature request or RFC. Make sure the feature request has been accepted for implementation before opening a PR.
- [ ] Related issues linked using `fixes #number`
- [ ] [e2e](https://github.com/vercel/next.js/blob/canary/contributing/core/testing.md#writing-tests-for-nextjs) tests added
- [ ] Documentation added
- [ ] Telemetry added. In case of a feature if it's used or not.
- [ ] Errors have a helpful link attached, see [`contributing.md`](https://github.com/vercel/next.js/blob/canary/contributing.md)

## Documentation / Examples

- [ ] Make sure the linting passes by running `pnpm build && pnpm lint`
- [ ] The "examples guidelines" are followed from [our contributing doc](https://github.com/vercel/next.js/blob/canary/contributing/examples/adding-examples.md)

packages/next-swc/crates/core/src/lib.rs

@@ -59,6 +59,7 @@ pub mod react_server_components;
 #[cfg(not(target_arch = "wasm32"))]
 pub mod relay;
 pub mod remove_console;
+pub mod server_actions;
 pub mod shake_exports;
 mod top_level_binding_collector;
 
@@ -122,6 +123,9 @@ pub struct TransformOptions {
 
     #[serde(default)]
     pub font_loaders: Option<next_font_loaders::Config>,
+
+    #[serde(default)]
+    pub server_actions: Option<server_actions::Config>,
 }
 
 pub fn custom_before_pass<'a, C: Comments + 'a>(
@@ -252,6 +256,11 @@ where
             Some(config) => Either::Left(next_font_loaders::next_font_loaders(config.clone())),
             None => Either::Right(noop()),
         },
+        match &opts.server_actions {
+            Some(config) =>
+                Either::Left(server_actions::server_actions(&file.name, config.clone())),
+            None => Either::Right(noop()),
+        },
     )
 }
 

packages/next-swc/crates/core/src/server_actions.rs

@@ -0,0 +1,392 @@
+use next_binding::swc::core::{
+    common::{errors::HANDLER, util::take::Take, FileName, DUMMY_SP},
+    ecma::{
+        ast::{
+            op, ArrayLit, AssignExpr, BlockStmt, CallExpr, ComputedPropName, Decl, ExportDecl,
+            Expr, ExprStmt, FnDecl, Function, Id, Ident, KeyValueProp, Lit, MemberExpr, MemberProp,
+            ModuleDecl, ModuleItem, PatOrExpr, Prop, PropName, ReturnStmt, Stmt, Str, VarDecl,
+            VarDeclKind, VarDeclarator,
+        },
+        atoms::JsWord,
+        utils::{find_pat_ids, private_ident, quote_ident, ExprFactory},
+        visit::{
+            as_folder, noop_visit_mut_type, noop_visit_type, visit_obj_and_computed, Fold, Visit,
+            VisitMut, VisitMutWith, VisitWith,
+        },
+    },
+};
+use serde::Deserialize;
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(deny_unknown_fields, rename_all = "camelCase")]
+pub struct Config {}
+
+pub fn server_actions(file_name: &FileName, config: Config) -> impl VisitMut + Fold {
+    as_folder(ServerActions {
+        config,
+        file_name: file_name.clone(),
+        in_action_file: false,
+        in_export_decl: false,
+        top_level: false,
+        closure_candidates: Default::default(),
+        annotations: Default::default(),
+        extra_items: Default::default(),
+    })
+}
+
+struct ServerActions {
+    #[allow(unused)]
+    config: Config,
+    file_name: FileName,
+
+    in_action_file: bool,
+    in_export_decl: bool,
+    top_level: bool,
+
+    closure_candidates: Vec<Id>,
+
+    annotations: Vec<Stmt>,
+    extra_items: Vec<ModuleItem>,
+}
+
+impl VisitMut for ServerActions {
+    fn visit_mut_export_decl(&mut self, decl: &mut ExportDecl) {
+        let old = self.in_export_decl;
+        self.in_export_decl = true;
+        decl.decl.visit_mut_with(self);
+        self.in_export_decl = old;
+    }
+
+    fn visit_mut_fn_decl(&mut self, f: &mut FnDecl) {
+        {
+            let old_len = self.closure_candidates.len();
+            self.closure_candidates
+                .extend(find_pat_ids(&f.function.params));
+
+            f.visit_mut_children_with(self);
+
+            self.closure_candidates.truncate(old_len);
+        }
+
+        if !(self.in_action_file && self.in_export_decl) {
+            // Check if the first item is `"use action"`;
+            if let Some(body) = &mut f.function.body {
+                if let Some(Stmt::Expr(first)) = body.stmts.first() {
+                    match &*first.expr {
+                        Expr::Lit(Lit::Str(Str { value, .. })) if value == "use action" => {}
+                        _ => return,
+                    }
+                } else {
+                    return;
+                }
+
+                body.stmts.remove(0);
+            } else {
+                return;
+            }
+        }
+
+        if !f.function.is_async {
+            HANDLER.with(|handler| {
+                handler
+                    .struct_span_err(f.ident.span, "Server actions must be async")
+                    .emit();
+            });
+        }
+
+        let action_name: JsWord = format!("$ACTION_{}", f.ident.sym).into();
+        let action_ident = private_ident!(action_name.clone());
+
+        // myAction.$$typeof = Symbol.for('react.action.reference');
+        self.annotations.push(annotate(
+            &f.ident,
+            "$$typeof",
+            CallExpr {
+                span: DUMMY_SP,
+                callee: quote_ident!("Symbol")
+                    .make_member(quote_ident!("for"))
+                    .as_callee(),
+                args: vec!["react.action.reference".as_arg()],
+                type_args: Default::default(),
+            }
+            .into(),
+        ));
+
+        // myAction.$$filepath = '/app/page.tsx';
+        self.annotations.push(annotate(
+            &f.ident,
+            "$$filepath",
+            self.file_name.to_string().into(),
+        ));
+
+        // myAction.$$name = '$ACTION_myAction';
+        self.annotations
+            .push(annotate(&f.ident, "$$name", action_name.into()));
+
+        if self.top_level {
+            // export const $ACTION_myAction = myAction;
+            self.extra_items
+                .push(ModuleItem::ModuleDecl(ModuleDecl::ExportDecl(ExportDecl {
+                    span: DUMMY_SP,
+                    decl: Decl::Var(Box::new(VarDecl {
+                        span: DUMMY_SP,
+                        kind: VarDeclKind::Const,
+                        declare: Default::default(),
+                        decls: vec![VarDeclarator {
+                            span: DUMMY_SP,
+                            name: action_ident.into(),
+                            init: Some(f.ident.clone().into()),
+                            definite: Default::default(),
+                        }],
+                    })),
+                })));
+        } else {
+            // Hoist the function to the top level.
+
+            let mut used_ids = idents_used_by(&f.function.body);
+
+            used_ids.retain(|id| self.closure_candidates.contains(id));
+
+            let closure_arg = private_ident!("closure");
+
+            f.function.body.visit_mut_with(&mut ClosureReplacer {
+                closure_arg: &closure_arg,
+                used_ids: &used_ids,
+            });
+
+            // myAction.$$closure = [id1, id2]
+            self.annotations.push(annotate(
+                &f.ident,
+                "$$closure",
+                ArrayLit {
+                    span: DUMMY_SP,
+                    elems: used_ids
+                        .iter()
+                        .cloned()
+                        .map(|id| Some(id.as_arg()))
+                        .collect(),
+                }
+                .into(),
+            ));
+
+            let call = CallExpr {
+                span: DUMMY_SP,
+                callee: action_ident.clone().as_callee(),
+                args: vec![f
+                    .ident
+                    .clone()
+                    .make_member(quote_ident!("$$closure"))
+                    .as_arg()],
+                type_args: Default::default(),
+            };
+
+            let new_fn = Box::new(Function {
+                params: f.function.params.clone(),
+                decorators: f.function.decorators.take(),
+                span: f.function.span,
+                body: Some(BlockStmt {
+                    span: DUMMY_SP,
+                    stmts: vec![Stmt::Return(ReturnStmt {
+                        span: DUMMY_SP,
+                        arg: Some(call.into()),
+                    })],
+                }),
+                is_generator: f.function.is_generator,
+                is_async: f.function.is_async,
+                type_params: Default::default(),
+                return_type: Default::default(),
+            });
+
+            self.extra_items
+                .push(ModuleItem::ModuleDecl(ModuleDecl::ExportDecl(ExportDecl {
+                    span: DUMMY_SP,
+                    decl: FnDecl {
+                        ident: action_ident,
+                        function: Box::new(Function {
+                            params: vec![closure_arg.into()],
+                            ..*f.function.take()
+                        }),
+                        declare: Default::default(),
+                    }
+                    .into(),
+                })));
+
+            f.function = new_fn;
+        }
+    }
+
+    fn visit_mut_module_item(&mut self, s: &mut ModuleItem) {
+        s.visit_mut_children_with(self);
+
+        if self.in_action_file {
+            if let ModuleItem::ModuleDecl(ModuleDecl::ExportDecl(ExportDecl {
+                decl: decl @ Decl::Fn(..),
+                ..
+            })) = s
+            {
+                *s = ModuleItem::Stmt(Stmt::Decl(decl.take()));
+            }
+        }
+    }
+
+    fn visit_mut_module_items(&mut self, stmts: &mut Vec<ModuleItem>) {
+        if let Some(ModuleItem::Stmt(Stmt::Expr(first))) = stmts.first() {
+            match &*first.expr {
+                Expr::Lit(Lit::Str(Str { value, .. })) if value == "use action" => {
+                    self.in_action_file = true;
+                }
+                _ => {}
+            }
+        }
+
+        if self.in_action_file {
+            stmts.remove(0);
+        }
+
+        let old_annotations = self.annotations.take();
+
+        let mut new = Vec::with_capacity(stmts.len());
+        for mut stmt in stmts.take() {
+            self.top_level = true;
+            stmt.visit_mut_with(self);
+
+            new.push(stmt);
+            new.extend(self.annotations.drain(..).map(ModuleItem::Stmt));
+            new.append(&mut self.extra_items);
+        }
+
+        *stmts = new;
+
+        self.annotations = old_annotations;
+    }
+
+    fn visit_mut_stmts(&mut self, stmts: &mut Vec<Stmt>) {
+        let old_top_level = self.top_level;
+        let old_annotations = self.annotations.take();
+
+        let mut new = Vec::with_capacity(stmts.len());
+        for mut stmt in stmts.take() {
+            self.top_level = false;
+            stmt.visit_mut_with(self);
+
+            new.push(stmt);
+            new.append(&mut self.annotations);
+        }
+
+        *stmts = new;
+
+        self.annotations = old_annotations;
+        self.top_level = old_top_level;
+    }
+
+    noop_visit_mut_type!();
+}
+
+fn annotate(fn_name: &Ident, field_name: &str, value: Box<Expr>) -> Stmt {
+    Stmt::Expr(ExprStmt {
+        span: DUMMY_SP,
+        expr: AssignExpr {
+            span: DUMMY_SP,
+            op: op!("="),
+            left: PatOrExpr::Expr(fn_name.clone().make_member(quote_ident!(field_name)).into()),
+            right: value,
+        }
+        .into(),
+    })
+}
+
+fn idents_used_by<N>(n: &N) -> Vec<Id>
+where
+    N: VisitWith<IdentUsageCollector>,
+{
+    let mut v = IdentUsageCollector {
+        ..Default::default()
+    };
+    n.visit_with(&mut v);
+    v.ids
+}
+
+#[derive(Default)]
+pub(crate) struct IdentUsageCollector {
+    ids: Vec<Id>,
+}
+
+impl Visit for IdentUsageCollector {
+    noop_visit_type!();
+
+    visit_obj_and_computed!();
+
+    fn visit_ident(&mut self, n: &Ident) {
+        if self.ids.contains(&n.to_id()) {
+            return;
+        }
+        self.ids.push(n.to_id());
+    }
+
+    fn visit_member_prop(&mut self, n: &MemberProp) {
+        if let MemberProp::Computed(..) = n {
+            n.visit_children_with(self);
+        }
+    }
+
+    fn visit_prop_name(&mut self, n: &PropName) {
+        if let PropName::Computed(..) = n {
+            n.visit_children_with(self);
+        }
+    }
+}
+
+pub(crate) struct ClosureReplacer<'a> {
+    closure_arg: &'a Ident,
+    used_ids: &'a [Id],
+}
+
+impl ClosureReplacer<'_> {
+    fn index(&self, i: &Ident) -> Option<usize> {
+        self.used_ids
+            .iter()
+            .position(|used_id| i.sym == used_id.0 && i.span.ctxt == used_id.1)
+    }
+}
+
+impl VisitMut for ClosureReplacer<'_> {
+    fn visit_mut_expr(&mut self, e: &mut Expr) {
+        e.visit_mut_children_with(self);
+
+        if let Expr::Ident(i) = e {
+            if let Some(index) = self.index(i) {
+                *e = Expr::Member(MemberExpr {
+                    span: DUMMY_SP,
+                    obj: self.closure_arg.clone().into(),
+                    prop: MemberProp::Computed(ComputedPropName {
+                        span: DUMMY_SP,
+                        expr: index.into(),
+                    }),
+                });
+            }
+        }
+    }
+
+    fn visit_mut_prop(&mut self, p: &mut Prop) {
+        p.visit_mut_children_with(self);
+
+        if let Prop::Shorthand(i) = p {
+            if let Some(index) = self.index(i) {
+                *p = Prop::KeyValue(KeyValueProp {
+                    key: PropName::Ident(i.clone()),
+                    value: MemberExpr {
+                        span: DUMMY_SP,
+                        obj: self.closure_arg.clone().into(),
+                        prop: MemberProp::Computed(ComputedPropName {
+                            span: DUMMY_SP,
+                            expr: index.into(),
+                        }),
+                    }
+                    .into(),
+                });
+            }
+        }
+    }
+
+    noop_visit_mut_type!();
+}

packages/next-swc/crates/core/tests/fixture.rs

@@ -17,6 +17,7 @@ use next_swc::{
     react_server_components::server_components,
     relay::{relay, Config as RelayConfig, RelayLanguageConfig},
     remove_console::remove_console,
+    server_actions::{self, server_actions},
     shake_exports::{shake_exports, Config as ShakeExportsConfig},
 };
 use std::path::PathBuf;
@@ -295,3 +296,20 @@ fn next_font_loaders_fixture(input: PathBuf) {
         Default::default(),
     );
 }
+
+#[fixture("tests/fixture/server-actions/**/input.js")]
+fn server_actions_fixture(input: PathBuf) {
+    let output = input.parent().unwrap().join("output.js");
+    test_fixture(
+        syntax(),
+        &|_tr| {
+            server_actions(
+                &FileName::Real("/app/item.js".into()),
+                server_actions::Config {},
+            )
+        },
+        &input,
+        &output,
+        Default::default(),
+    );
+}

packages/next-swc/crates/core/tests/fixture/server-actions/1/input.js

@@ -0,0 +1,8 @@
+export function Item({ id1, id2 }) {
+    async function deleteItem() {
+        "use action";
+        await deleteFromDb(id1);
+        await deleteFromDb(id2);
+    }
+    return <Button action={deleteItem}>Delete</Button>;
+}

packages/next-swc/crates/core/tests/fixture/server-actions/1/output.js

@@ -0,0 +1,17 @@
+export function Item({ id1 , id2  }) {
+    async function deleteItem() {
+        return $ACTION_deleteItem(deleteItem.$$closure);
+    }
+    deleteItem.$$typeof = Symbol.for("react.action.reference");
+    deleteItem.$$filepath = "/app/item.js";
+    deleteItem.$$name = "$ACTION_deleteItem";
+    deleteItem.$$closure = [
+        id1,
+        id2
+    ];
+    return <Button action={deleteItem}>Delete</Button>;
+}
+export async function $ACTION_deleteItem(closure) {
+    await deleteFromDb(closure[0]);
+    await deleteFromDb(closure[1]);
+}

packages/next-swc/crates/core/tests/fixture/server-actions/2/input.js

@@ -0,0 +1,8 @@
+async function myAction(a, b, c) {
+  "use action";
+  console.log('a')
+}
+
+export default function Page() {
+   return <Button action={myAction}>Delete</Button>;
+}

packages/next-swc/crates/core/tests/fixture/server-actions/2/output.js

@@ -0,0 +1,10 @@
+async function myAction(a, b, c) {
+    console.log('a');
+}
+myAction.$$typeof = Symbol.for("react.action.reference");
+myAction.$$filepath = "/app/item.js";
+myAction.$$name = "$ACTION_myAction";
+export const $ACTION_myAction = myAction;
+export default function Page() {
+    return <Button action={myAction}>Delete</Button>;
+}

packages/next-swc/crates/core/tests/fixture/server-actions/3/input.js

@@ -0,0 +1,5 @@
+// app/send.ts
+"use action";
+export async function myAction(a, b, c) {
+  console.log('a')
+}

packages/next-swc/crates/core/tests/fixture/server-actions/3/output.js

@@ -0,0 +1,8 @@
+// app/send.ts
+async function myAction(a, b, c) {
+    console.log('a');
+}
+myAction.$$typeof = Symbol.for("react.action.reference");
+myAction.$$filepath = "/app/item.js";
+myAction.$$name = "$ACTION_myAction";
+export const $ACTION_myAction = myAction;

packages/next-swc/crates/core/tests/full.rs

@@ -70,6 +70,7 @@ fn test(input: &Path, minify: bool) {
                 emotion: Some(assert_json("{}")),
                 modularize_imports: None,
                 font_loaders: None,
+                server_actions: None,
             };
 
             let options = options.patch(&fm);