fix(rsc): align export async function checks

packages/plugin-rsc/src/transforms/proxy-export.ts

@@ -1,7 +1,7 @@
 import { tinyassert } from '@hiogawa/utils'
 import type { Node, Program } from 'estree'
 import MagicString from 'magic-string'
-import { extractNames, hasDirective } from './utils'
+import { extractNames, hasDirective, validateNonAsyncFunction } from './utils'
 
 export type TransformProxyExportOptions = {
   /** Required for source map and `keep` options */
@@ -58,14 +58,6 @@ export function transformProxyExport(
     output.update(node.start, node.end, newCode)
   }
 
-  function validateNonAsyncFunction(node: Node, ok?: boolean) {
-    if (options.rejectNonAsyncFunction && !ok) {
-      throw Object.assign(new Error(`unsupported non async function`), {
-        pos: node.start,
-      })
-    }
-  }
-
   for (const node of ast.body) {
     if (node.type === 'ExportNamedDeclaration') {
       if (node.declaration) {
@@ -76,24 +68,15 @@ export function transformProxyExport(
           /**
            * export function foo() {}
            */
-          validateNonAsyncFunction(
-            node,
-            node.declaration.type === 'FunctionDeclaration' &&
-              node.declaration.async,
-          )
+          validateNonAsyncFunction(options, node.declaration)
           createExport(node, [node.declaration.id.name])
         } else if (node.declaration.type === 'VariableDeclaration') {
           /**
            * export const foo = 1, bar = 2
            */
-          validateNonAsyncFunction(
-            node,
-            node.declaration.declarations.every(
-              (decl) =>
-                decl.init?.type === 'ArrowFunctionExpression' &&
-                decl.init.async,
-            ),
-          )
+          for (const decl of node.declaration.declarations) {
+            if (decl.init) validateNonAsyncFunction(options, decl.init)
+          }
           if (options.keep && options.code) {
             if (node.declaration.declarations.length === 1) {
               const decl = node.declaration.declarations[0]!
@@ -148,12 +131,7 @@ export function transformProxyExport(
      * export default () => {}
      */
     if (node.type === 'ExportDefaultDeclaration') {
-      validateNonAsyncFunction(
-        node,
-        node.declaration.type === 'Identifier' ||
-          (node.declaration.type === 'FunctionDeclaration' &&
-            node.declaration.async),
-      )
+      validateNonAsyncFunction(options, node.declaration)
       createExport(node, ['default'])
       continue
     }

packages/plugin-rsc/src/transforms/utils.test.ts

@@ -0,0 +1,83 @@
+import { parseAstAsync } from 'vite'
+import { describe, expect, test } from 'vitest'
+import { transformProxyExport } from './proxy-export'
+import { validateNonAsyncFunction } from './utils'
+import { transformWrapExport } from './wrap-export'
+
+describe(validateNonAsyncFunction, () => {
+  // next.js's validation isn't entirely consistent.
+  // for now we aim to make it at least as forgiving as next.js.
+
+  const accepted = [
+    `export async function f() {}`,
+    `export default async function f() {}`,
+    `export const fn = async function fn() {}`,
+    `export const fn = async () => {}`,
+    `export const fn = async () => {}, fn2 = x`,
+    `export const fn = x`,
+    `export const fn = x({ x: y })`,
+    `export const fn = x(async () => {})`,
+    `export default x`,
+    `const y = x; export { y }`,
+    `export const fn = x(() => {})`, // rejected by next.js
+    `export const testAction = actionClient.action(async () => { return { message: "Hello, world!" }; });`,
+  ]
+
+  const rejected = [
+    `export function f() {}`,
+    `export default function f() {}`,
+    `export const fn = function fn() {}`,
+    `export const fn = () => {}`,
+    `export const fn = x, fn2 = () => {}`,
+    `export class Cls {}`,
+    `export const Cls = class {}`,
+    `export const Cls = class Foo {}`,
+  ]
+
+  test(transformWrapExport, async () => {
+    const testTransform = async (input: string) => {
+      const ast = await parseAstAsync(input)
+      const result = transformWrapExport(input, ast, {
+        runtime: (value, name) =>
+          `$$wrap(${value}, "<id>", ${JSON.stringify(name)})`,
+        ignoreExportAllDeclaration: true,
+        rejectNonAsyncFunction: true,
+      })
+      return result.output.hasChanged()
+    }
+
+    for (const code of accepted) {
+      await expect.soft(testTransform(code)).resolves.toBe(true)
+    }
+    for (const code of rejected) {
+      await expect
+        .soft(testTransform(code))
+        .rejects.toMatchInlineSnapshot(
+          `[Error: unsupported non async function]`,
+        )
+    }
+  })
+
+  test(transformProxyExport, async () => {
+    const testTransform = async (input: string) => {
+      const ast = await parseAstAsync(input)
+      const result = transformProxyExport(ast, {
+        code: input,
+        rejectNonAsyncFunction: true,
+        runtime: (name) => `$$proxy("<id>", ${JSON.stringify(name)})`,
+      })
+      return result.output.hasChanged()
+    }
+
+    for (const code of accepted) {
+      await expect.soft(testTransform(code)).resolves.toBe(true)
+    }
+    for (const code of rejected) {
+      await expect
+        .soft(testTransform(code))
+        .rejects.toMatchInlineSnapshot(
+          `[Error: unsupported non async function]`,
+        )
+    }
+  })
+})

packages/plugin-rsc/src/transforms/utils.ts

@@ -1,5 +1,6 @@
 import { tinyassert } from '@hiogawa/utils'
-import type { Identifier, Pattern, Program } from 'estree'
+import type { ExportDefaultDeclaration } from 'estree'
+import type { Identifier, Node, Pattern, Program } from 'estree'
 
 export function hasDirective(
   body: Program['body'],
@@ -136,3 +137,23 @@ export function extractIdentifiers(
   }
   return nodes
 }
+
+export function validateNonAsyncFunction(
+  opts: { rejectNonAsyncFunction?: boolean },
+  // export default function/class can be unnamed
+  node: Node | ExportDefaultDeclaration['declaration'],
+): void {
+  if (!opts.rejectNonAsyncFunction) return
+  if (
+    node.type === 'ClassDeclaration' ||
+    node.type === 'ClassExpression' ||
+    ((node.type === 'FunctionDeclaration' ||
+      node.type === 'FunctionExpression' ||
+      node.type === 'ArrowFunctionExpression') &&
+      !node.async)
+  ) {
+    throw Object.assign(new Error(`unsupported non async function`), {
+      pos: node.start,
+    })
+  }
+}

packages/plugin-rsc/src/transforms/wrap-export.test.ts

@@ -304,55 +304,4 @@ export default Page;
       "
     `)
   })
-
-  test('reject non async function', async () => {
-    // next.js's validataion isn't entirely consisten.
-    // for now we aim to make it at least as forgiving as next.js.
-
-    const accepted = [
-      `export async function f() {}`,
-      `export default async function f() {}`,
-      `export const fn = async function fn() {}`,
-      `export const fn = async () => {}`,
-      `export const fn = async () => {}, fn2 = x`,
-      `export const fn = x`,
-      `export const fn = x({ x: y })`,
-      `export const fn = x(async () => {})`,
-      `export default x`,
-      `const y = x; export { y }`,
-      `export const fn = x(() => {})`, // rejected by next.js
-    ]
-
-    const rejected = [
-      `export function f() {}`,
-      `export default function f() {}`,
-      `export const fn = function fn() {}`,
-      `export const fn = () => {}`,
-      `export const fn = x, fn2 = () => {}`,
-      `export class Cls {}`,
-    ]
-
-    async function toActual(input: string) {
-      try {
-        await testTransform(input, {
-          rejectNonAsyncFunction: true,
-        })
-        return [input, true]
-      } catch (e) {
-        return [input, e instanceof Error ? e.message : e]
-      }
-    }
-
-    const actual = [
-      ...(await Promise.all(accepted.map((e) => toActual(e)))),
-      ...(await Promise.all(rejected.map((e) => toActual(e)))),
-    ]
-
-    const expected = [
-      ...accepted.map((e) => [e, true]),
-      ...rejected.map((e) => [e, 'unsupported non async function']),
-    ]
-
-    expect(actual).toEqual(expected)
-  })
 })

packages/plugin-rsc/src/transforms/wrap-export.ts

@@ -1,7 +1,7 @@
 import { tinyassert } from '@hiogawa/utils'
-import type { Node, Program } from 'estree'
+import type { Program } from 'estree'
 import MagicString from 'magic-string'
-import { extractNames } from './utils'
+import { extractNames, validateNonAsyncFunction } from './utils'
 
 type ExportMeta = {
   declName?: string
@@ -83,21 +83,6 @@ export function transformWrapExport(
     )
   }
 
-  function validateNonAsyncFunction(node: Node) {
-    if (!options.rejectNonAsyncFunction) return
-    if (
-      node.type === 'ClassDeclaration' ||
-      ((node.type === 'FunctionDeclaration' ||
-        node.type === 'FunctionExpression' ||
-        node.type === 'ArrowFunctionExpression') &&
-        !node.async)
-    ) {
-      throw Object.assign(new Error(`unsupported non async function`), {
-        pos: node.start,
-      })
-    }
-  }
-
   for (const node of ast.body) {
     // named exports
     if (node.type === 'ExportNamedDeclaration') {
@@ -109,7 +94,7 @@ export function transformWrapExport(
           /**
            * export function foo() {}
            */
-          validateNonAsyncFunction(node.declaration)
+          validateNonAsyncFunction(options, node.declaration)
           const name = node.declaration.id.name
           wrapSimple(node.start, node.declaration.start, [
             { name, meta: { isFunction: true, declName: name } },
@@ -120,7 +105,7 @@ export function transformWrapExport(
            */
           for (const decl of node.declaration.declarations) {
             if (decl.init) {
-              validateNonAsyncFunction(decl.init)
+              validateNonAsyncFunction(options, decl.init)
             }
           }
           if (node.declaration.kind === 'const') {
@@ -203,7 +188,7 @@ export function transformWrapExport(
      * export default () => {}
      */
     if (node.type === 'ExportDefaultDeclaration') {
-      validateNonAsyncFunction(node.declaration as Node)
+      validateNonAsyncFunction(options, node.declaration)
       let localName: string
       let isFunction = false
       let declName: string | undefined