Repository: dai-shi/waku (ts)

Issue:
Uncovered server action patterns
https://github.com/dai-shi/waku/pull/774#issuecomment-2208076409

## Default export function with top-level "use server"

```js
"use server";

export default async function foo() {
  // ...
}
```

## Default export function with inline "use server"

```js
export default async function foo() {
  "use server";
  // ...
}
```

----

There might be some more.

Pull request description:
fix: support default export server action
close #777

Patch:
diff --git a/packages/waku/src/lib/plugins/vite-plugin-rsc-transform.ts b/packages/waku/src/lib/plugins/vite-plugin-rsc-transform.ts
index 3d66a038..9584f1ad 100644
--- a/packages/waku/src/lib/plugins/vite-plugin-rsc-transform.ts
+++ b/packages/waku/src/lib/plugins/vite-plugin-rsc-transform.ts
@@ -106,6 +106,13 @@ const findLastImportIndex = (mod: swc.Module) => {
   return lastImportIndex === -1 ? 0 : lastImportIndex;
 };
 
+const replaceNode = <T extends swc.Node>(origNode: swc.Node, newNode: T): T => {
+  Object.keys(origNode).forEach((key) => {
+    delete origNode[key as never];
+  });
+  return Object.assign(origNode, newNode);
+};
+
 const transformExportedServerActions = (
   mod: swc.Module,
   getActionId: () => string,
@@ -113,15 +120,9 @@ const transformExportedServerActions = (
   let changed = false;
   for (let i = 0; i < mod.body.length; ++i) {
     const item = mod.body[i]!;
-    const addRegisterStatement = (
-      name: string,
-      fn:
-        | swc.FunctionDeclaration
-        | swc.FunctionExpression
-        | swc.ArrowFunctionExpression,
-    ) => {
+    const handleDeclaration = (name: string, fn: swc.FunctionDeclaration) => {
       changed = true;
-      if (fn.body?.type === 'BlockStatement') {
+      if (fn.body) {
         fn.body.stmts = fn.body.stmts.filter(
           (stmt) => !isUseServerDirective(stmt),
         );
@@ -140,12 +141,29 @@ const transformExportedServerActions = (
       };
       mod.body.splice(++i, 0, stmt);
     };
+    const handleExpression = (
+      name: string,
+      fn: swc.FunctionExpression | swc.ArrowFunctionExpression,
+    ) => {
+      changed = true;
+      if (fn.body?.type === 'BlockStatement') {
+        fn.body.stmts = fn.body.stmts.filter(
+          (stmt) => !isUseServerDirective(stmt),
+        );
+      }
+      const callExp = createCallExpression(
+        createIdentifier('__waku_registerServerReference'),
+        [
+          Object.assign({}, fn),
+          createStringLiteral(getActionId()),
+          createStringLiteral(name),
+        ],
+      );
+      replaceNode(fn, callExp);
+    };
     if (item.type === 'ExportDeclaration') {
       if (item.declaration.type === 'FunctionDeclaration') {
-        addRegisterStatement(
-          item.declaration.identifier.value,
-          item.declaration,
-        );
+        handleDeclaration(item.declaration.identifier.value, item.declaration);
       } else if (item.declaration.type === 'VariableDeclaration') {
         for (const d of item.declaration.declarations) {
           if (
@@ -153,10 +171,28 @@ const transformExportedServerActions = (
             (d.init?.type === 'FunctionExpression' ||
               d.init?.type === 'ArrowFunctionExpression')
           ) {
-            addRegisterStatement(d.id.value, d.init);
+            handleExpression(d.id.value, d.init);
           }
         }
       }
+    } else if (item.type === 'ExportDefaultDeclaration') {
+      if (item.decl.type === 'FunctionExpression') {
+        handleExpression('default', item.decl);
+        const callExp = item.decl;
+        const decl: swc.ExportDefaultExpression = {
+          type: 'ExportDefaultExpression',
+          expression: callExp,
+          span: { start: 0, end: 0, ctxt: 0 },
+        };
+        replaceNode(item, decl);
+      }
+    } else if (item.type === 'ExportDefaultExpression') {
+      if (
+        item.expression.type === 'FunctionExpression' ||
+        item.expression.type === 'ArrowFunctionExpression'
+      ) {
+        handleExpression('default', item.expression);
+      }
     }
   }
   return changed;
@@ -336,28 +372,40 @@ const transformInlineServerActions = (
         ],
         span: { start: 0, end: 0, ctxt: 0 },
       };
-      Object.keys(decl).forEach((key) => {
-        delete decl[key as keyof typeof decl];
-      });
-      Object.assign(decl, newDecl);
+      replaceNode(decl, newDecl);
     }
   };
   const handleExpression = (
     parentFn: swc.Fn | swc.ArrowFunctionExpression | undefined,
     exp: swc.Expression,
-  ) => {
+  ): swc.CallExpression | undefined => {
     if (isInlineServerAction(exp)) {
       const callExp = registerServerAction(parentFn, Object.assign({}, exp));
-      Object.keys(exp).forEach((key) => {
-        delete exp[key as keyof typeof exp];
-      });
-      Object.assign(exp, callExp);
+      return replaceNode(exp, callExp);
     }
   };
   const walk = (
     parentFn: swc.Fn | swc.ArrowFunctionExpression | undefined,
     node: swc.Node,
   ) => {
+    if (node.type === 'ExportDefaultDeclaration') {
+      const item = node as swc.ExportDefaultDeclaration;
+      if (item.decl.type === 'FunctionExpression') {
+        const callExp = handleExpression(
+          parentFn,
+          item.decl as swc.FunctionExpression,
+        );
+        if (callExp) {
+          const decl: swc.ExportDefaultExpression = {
+            type: 'ExportDefaultExpression',
+            expression: callExp,
+            span: { start: 0, end: 0, ctxt: 0 },
+          };
+          replaceNode(item, decl);
+          return;
+        }
+      }
+    }
     // FIXME do we need to walk the entire tree? feels inefficient
     Object.values(node).forEach((value) => {
       const fn =


Test patch:
diff --git a/packages/waku/tests/vite-plugin-rsc-transform-internals.test.ts b/packages/waku/tests/vite-plugin-rsc-transform-internals.test.ts
index 7e01a659..c348c53c 100644
--- a/packages/waku/tests/vite-plugin-rsc-transform-internals.test.ts
+++ b/packages/waku/tests/vite-plugin-rsc-transform-internals.test.ts
@@ -71,24 +71,24 @@ export async function greet(name) {
   return 'Hello ' + name;
 }
 
-// TODO https://github.com/dai-shi/waku/issues/777
-// support default export
-// export default async function() {
-//   return Date.now();
-// }
+export default async function() {
+  return Date.now();
+}
 `;
     expect(await transform(code, '/src/App.tsx', { ssr: true }))
       .toMatchInlineSnapshot(`
         "import { registerServerReference as __waku_registerServerReference } from 'react-server-dom-webpack/server.edge';
         const privateFunction = ()=>"Secret";
-        export const log = async (mesg)=>{
+        export const log = __waku_registerServerReference(async (mesg)=>{
             console.log(mesg);
-        };
-        __waku_registerServerReference(log, "/src/App.tsx", "log");
+        }, "/src/App.tsx", "log");
         export async function greet(name) {
             return 'Hello ' + name;
         }
         __waku_registerServerReference(greet, "/src/App.tsx", "greet");
+        export default __waku_registerServerReference(async function() {
+            return Date.now();
+        }, "/src/App.tsx", "default");
         "
       `);
   });
@@ -153,6 +153,8 @@ export async function exportedAction() {
   'use server';
   return null;
 }
+
+export default async () => null;
 `;
     expect(await transform(code, '/src/App.tsx', { ssr: true }))
       .toMatchInlineSnapshot(`
@@ -173,6 +175,7 @@ export async function exportedAction() {
             return null;
         }
         __waku_registerServerReference(exportedAction, "/src/App.tsx", "exportedAction");
+        export default __waku_registerServerReference(async ()=>null, "/src/App.tsx", "default");
         "
       `);
   });
@@ -282,12 +285,10 @@ const log4 = async (mesg) => {
   console.log(mesg);
 };
 
-// TODO https://github.com/dai-shi/waku/issues/777
-// default export anonymous function
-// export default async function(mesg) {
-//   'use server';
-//   console.log(mesg);
-// }
+export default async function(mesg) {
+  'use server';
+  console.log(mesg);
+}
 `;
     expect(await transform(code, '/src/App.tsx', { ssr: true }))
       .toMatchInlineSnapshot(`
@@ -305,12 +306,16 @@ const log4 = async (mesg) => {
         export const __waku_action4 = __waku_registerServerReference(async (mesg)=>{
             console.log(mesg);
         }, "/src/App.tsx", "__waku_action4");
+        export const __waku_action5 = __waku_registerServerReference(async function(mesg) {
+            console.log(mesg);
+        }, "/src/App.tsx", "__waku_action5");
         const actions = {
             log: __waku_action1.bind(null)
         };
         const log2 = __waku_action2.bind(null);
         const log3 = __waku_action3.bind(null);
         const log4 = __waku_action4.bind(null);
+        export default __waku_action5.bind(null);
         "
       `);
   });