feat(rsc): support regex directive for transformHoistInlineDirective

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

@@ -6,12 +6,21 @@ import { debugSourceMap } from './test-utils'
 describe(transformHoistInlineDirective, () => {
   async function testTransform(
     input: string,
-    options?: { encode?: boolean; noExport?: boolean; directive?: string },
+    options?: {
+      encode?: boolean
+      noExport?: boolean
+      directive?: string | RegExp
+    },
   ) {
     const ast = await parseAstAsync(input)
     const { output } = transformHoistInlineDirective(input, ast, {
-      runtime: (value, name) =>
-        `$$register(${value}, "<id>", ${JSON.stringify(name)})`,
+      runtime: (value, name, meta) =>
+        `$$register(${value}, "<id>", ${JSON.stringify(name)}` +
+        `${
+          options?.directive instanceof RegExp
+            ? `, ${JSON.stringify(meta)}`
+            : ''
+        })`,
       directive: options?.directive ?? 'use server',
       encode: options?.encode ? (v) => `__enc(${v})` : undefined,
       decode: options?.encode ? (v) => `__dec(${v})` : undefined,
@@ -369,4 +378,55 @@ export async function test() {
       "
     `)
   })
+
+  it('directive pattern', async () => {
+    const input = `
+export async function none() {
+  "use cache";
+  return "test";
+}
+
+export async function fs() {
+  "use cache: fs";
+  return "test";
+}
+
+export async function kv() {
+  "use cache: kv";
+  return "test";
+}
+`
+    expect(
+      await testTransform(input, {
+        directive: /^use cache(: .+)?$/,
+        noExport: true,
+      }),
+    ).toMatchInlineSnapshot(`
+      "
+      export const none = /* #__PURE__ */ $$register($$hoist_0_none, "<id>", "$$hoist_0_none", {"directiveMatch":["use cache",null]});
+
+      export const fs = /* #__PURE__ */ $$register($$hoist_1_fs, "<id>", "$$hoist_1_fs", {"directiveMatch":["use cache: fs",": fs"]});
+
+      export const kv = /* #__PURE__ */ $$register($$hoist_2_kv, "<id>", "$$hoist_2_kv", {"directiveMatch":["use cache: kv",": kv"]});
+
+      ;async function $$hoist_0_none() {
+        "use cache";
+        return "test";
+      };
+      /* #__PURE__ */ Object.defineProperty($$hoist_0_none, "name", { value: "none" });
+
+      ;async function $$hoist_1_fs() {
+        "use cache: fs";
+        return "test";
+      };
+      /* #__PURE__ */ Object.defineProperty($$hoist_1_fs, "name", { value: "fs" });
+
+      ;async function $$hoist_2_kv() {
+        "use cache: kv";
+        return "test";
+      };
+      /* #__PURE__ */ Object.defineProperty($$hoist_2_kv, "name", { value: "kv" });
+      "
+    `)
+  })
 })

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

@@ -3,19 +3,21 @@ import type { Program } from 'estree'
 import { walk } from 'estree-walker'
 import MagicString from 'magic-string'
 import { analyze } from 'periscopic'
-import { hasDirective } from './utils'
 
 export function transformHoistInlineDirective(
   input: string,
   ast: Program,
   {
     runtime,
-    directive,
     rejectNonAsyncFunction,
     ...options
   }: {
-    runtime: (value: string, name: string) => string
-    directive: string
+    runtime: (
+      value: string,
+      name: string,
+      meta: { directiveMatch: RegExpMatchArray },
+    ) => string
+    directive: string | RegExp
     rejectNonAsyncFunction?: boolean
     encode?: (value: string) => string
     decode?: (value: string) => string
@@ -26,6 +28,10 @@ export function transformHoistInlineDirective(
   names: string[]
 } {
   const output = new MagicString(input)
+  const directive =
+    typeof options.directive === 'string'
+      ? exactRegex(options.directive)
+      : options.directive
 
   // re-export somehow confuses periscopic scopes so remove them before analysis
   walk(ast, {
@@ -48,9 +54,10 @@ export function transformHoistInlineDirective(
         (node.type === 'FunctionExpression' ||
           node.type === 'FunctionDeclaration' ||
           node.type === 'ArrowFunctionExpression') &&
-        node.body.type === 'BlockStatement' &&
-        hasDirective(node.body.body, directive)
+        node.body.type === 'BlockStatement'
       ) {
+        const match = matchDirective(node.body.body, directive)
+        if (!match) return
         if (!node.async && rejectNonAsyncFunction) {
           throw Object.assign(
             new Error(`"${directive}" doesn't allow non async function`),
@@ -116,7 +123,9 @@ export function transformHoistInlineDirective(
         output.move(node.start, node.end, input.length)
 
         // replace original declartion with action register + bind
-        let newCode = `/* #__PURE__ */ ${runtime(newName, newName)}`
+        let newCode = `/* #__PURE__ */ ${runtime(newName, newName, {
+          directiveMatch: match,
+        })}`
         if (bindVars.length > 0) {
           const bindArgs = options.encode
             ? options.encode('[' + bindVars.join(', ') + ']')
@@ -140,3 +149,24 @@ export function transformHoistInlineDirective(
     names,
   }
 }
+
+const exactRegex = (s: string): RegExp =>
+  new RegExp('^' + s.replace(/[-[\]{}()*+?.,\\^$|#\s]/g, '\\$&') + '$')
+
+function matchDirective(
+  body: Program['body'],
+  directive: RegExp,
+): RegExpMatchArray | undefined {
+  for (const stable of body) {
+    if (
+      stable.type === 'ExpressionStatement' &&
+      stable.expression.type === 'Literal' &&
+      typeof stable.expression.value === 'string'
+    ) {
+      const match = stable.expression.value.match(directive)
+      if (match) {
+        return match
+      }
+    }
+  }
+}