fix: compile directive replacements

packages/directive-functions-plugin/src/compilers.ts

@@ -68,6 +68,7 @@ export function compileDirectives(opts: CompileDirectivesOpts): {
   const directiveFnsById = findDirectives(ast, {
     ...opts,
     directiveSplitParam,
+    isDirectiveSplitParam,
   })
 
   // Add runtime code if there are directives
@@ -224,6 +225,7 @@ export function findDirectives(
     directiveSplitParam: string
     filename: string
     root: string
+    isDirectiveSplitParam: boolean
   },
 ): Record<string, DirectiveFn> {
   const directiveFnsById: Record<string, DirectiveFn> = {}
@@ -241,6 +243,9 @@ export function findDirectives(
   const hasFileDirective = ast.program.directives.some(
     (directive) => directive.value.value === opts.directive,
   )
+  const compileDirectiveOpts = {
+    isDirectiveSplitParam: opts.isDirectiveSplitParam,
+  }
 
   // If the entire file has a directive, we need to compile all of the functions that are
   // exported by the file.
@@ -250,12 +255,18 @@ export function findDirectives(
     babel.traverse(ast, {
       ExportDefaultDeclaration(path) {
         if (babel.types.isFunctionDeclaration(path.node.declaration)) {
-          compileDirective(path.get('declaration') as SupportedFunctionPath)
+          compileDirective(
+            path.get('declaration') as SupportedFunctionPath,
+            compileDirectiveOpts,
+          )
         }
       },
       ExportNamedDeclaration(path) {
         if (babel.types.isFunctionDeclaration(path.node.declaration)) {
-          compileDirective(path.get('declaration') as SupportedFunctionPath)
+          compileDirective(
+            path.get('declaration') as SupportedFunctionPath,
+            compileDirectiveOpts,
+          )
         }
       },
       ExportDeclaration(path) {
@@ -273,6 +284,7 @@ export function findDirectives(
             path.get(
               'declaration.declarations.0.init',
             ) as SupportedFunctionPath,
+            compileDirectiveOpts,
           )
         }
       },
@@ -327,15 +339,18 @@ export function findDirectives(
             )
           }
 
-          compileDirective(directiveFn)
+          compileDirective(directiveFn, compileDirectiveOpts)
         }
       },
     })
   }
 
   return directiveFnsById
 
-  function compileDirective(directiveFn: SupportedFunctionPath) {
+  function compileDirective(
+    directiveFn: SupportedFunctionPath,
+    compileDirectiveOpts: { isDirectiveSplitParam: boolean },
+  ) {
     // Move the function to program level while preserving its position
     // in the program body
     const programBody = programPath.node.body
@@ -432,7 +447,11 @@ export function findDirectives(
     // If it's an exported named function, we need to swap it with an
     // export const originalFunctionName = functionName
     if (
-      babel.types.isExportNamedDeclaration(directiveFn.parentPath.node) &&
+      (babel.types.isExportNamedDeclaration(directiveFn.parentPath.node) ||
+        (compileDirectiveOpts.isDirectiveSplitParam &&
+          babel.types.isExportDefaultDeclaration(
+            directiveFn.parentPath.node,
+          ))) &&
       (babel.types.isFunctionDeclaration(directiveFn.node) ||
         babel.types.isFunctionExpression(directiveFn.node)) &&
       babel.types.isIdentifier(directiveFn.node.id)
@@ -542,57 +561,32 @@ function codeFrameError(
 }
 
 const safeRemoveExports = (ast: babel.types.File) => {
-  const programBody = ast.program.body
-
-  const removeExport = (
-    path:
-      | babel.NodePath<babel.types.ExportDefaultDeclaration>
-      | babel.NodePath<babel.types.ExportNamedDeclaration>,
-  ) => {
-    // If the value is a function declaration, class declaration, or variable declaration,
-    // That means it has a name and can remain in the file, just unexported.
+  ast.program.body = ast.program.body.flatMap((node) => {
     if (
-      babel.types.isFunctionDeclaration(path.node.declaration) ||
-      babel.types.isClassDeclaration(path.node.declaration) ||
-      babel.types.isVariableDeclaration(path.node.declaration)
+      babel.types.isExportNamedDeclaration(node) ||
+      babel.types.isExportDefaultDeclaration(node)
     ) {
-      // If the value is a function declaration, class declaration, or variable declaration,
-      // That means it has a name and can remain in the file, just unexported.
       if (
-        babel.types.isFunctionDeclaration(path.node.declaration) ||
-        babel.types.isClassDeclaration(path.node.declaration) ||
-        babel.types.isVariableDeclaration(path.node.declaration)
+        babel.types.isFunctionDeclaration(node.declaration) ||
+        babel.types.isClassDeclaration(node.declaration) ||
+        babel.types.isVariableDeclaration(node.declaration)
       ) {
-        // Move the declaration to the top level at the same index
-        const insertIndex = programBody.findIndex(
-          (node) => node === path.node.declaration,
-        )
         // do not remove export if it is an anonymous function / class, otherwise this would produce a syntax error
         if (
-          babel.types.isFunctionDeclaration(path.node.declaration) ||
-          babel.types.isClassDeclaration(path.node.declaration)
+          babel.types.isFunctionDeclaration(node.declaration) ||
+          babel.types.isClassDeclaration(node.declaration)
         ) {
-          if (!path.node.declaration.id) {
-            return
+          if (!node.declaration.id) {
+            return node
           }
         }
-        programBody.splice(insertIndex, 0, path.node.declaration as any)
+        return node.declaration
+      } else if (node.declaration === null) {
+        // remove e.g. `export { RouteComponent as component }`
+        return []
       }
     }
-
-    // Otherwise, remove the export declaration
-    path.remove()
-  }
-
-  // Before we add our export, remove any other exports.
-  // Don't remove the thing they export, just the export declaration
-  babel.traverse(ast, {
-    ExportDefaultDeclaration(path) {
-      removeExport(path)
-    },
-    ExportNamedDeclaration(path) {
-      removeExport(path)
-    },
+    return node
   })
 }
 

packages/directive-functions-plugin/tests/compiler.test.ts

@@ -255,9 +255,11 @@ describe('server function compilation', () => {
       const defaultExportFn_1 = createServerRpc("test_ts--defaultExportFn_1", function defaultExportFn() {
         return 'hello';
       });
+      const defaultExportFn = defaultExportFn_1;
       const namedExportFn_1 = createServerRpc("test_ts--namedExportFn_1", function namedExportFn() {
         return 'hello';
       });
+      const namedExportFn = namedExportFn_1;
       const exportedArrowFunction_wrapper = createServerRpc("test_ts--exportedArrowFunction_wrapper", () => {
         return 'hello';
       });
@@ -274,8 +276,9 @@ describe('server function compilation', () => {
       function unusedFn() {
         return 'hello';
       }
+      const namedDefaultExport = 'namedDefaultExport';
+      export default namedDefaultExport;
       const usedButNotExported = 'usedButNotExported';
-      const namedExportFn = namedExportFn_1;
       export { namedFunction_createServerFn_namedFunction, namedGeneratorFunction_createServerFn_namedGeneratorFunction, arrowFunction_createServerFn, anonymousFunction_createServerFn, anonymousGeneratorFunction_createServerFn, multipleDirectives_multipleDirectives, iife_1, defaultExportFn_1, namedExportFn_1, exportedArrowFunction_wrapper, namedExportConst_1, namedExportConstGenerator_1 };"
     `,
     )
@@ -552,6 +555,7 @@ describe('server function compilation', () => {
       const useServer_1 = createServerRpc("test_ts--useServer_1", function useServer() {
         return usedInUseServer();
       });
+      const useServer = useServer_1;
       function notExported() {
         return 'hello';
       }
@@ -561,7 +565,7 @@ describe('server function compilation', () => {
       const defaultExport_1 = createServerRpc("test_ts--defaultExport_1", function defaultExport() {
         return 'hello';
       });
-      const useServer = useServer_1;
+      const defaultExport = defaultExport_1;
       export { useServer_1, defaultExport_1 };"
     `)
   })
@@ -642,6 +646,7 @@ describe('server function compilation', () => {
       const myServerFn_createServerFn_handler = createServerRpc("test_ts--myServerFn_createServerFn_handler", opts => {
         return myServerFn.__executeServer(opts);
       });
+      const myServerFn = createServerFn().handler(myServerFn_createServerFn_handler, myFunc);
       const myFunc2 = () => {
         return myServerFn({
           data: 'hello 2 from the server'
@@ -650,7 +655,6 @@ describe('server function compilation', () => {
       const myServerFn2_createServerFn_handler = createServerRpc("test_ts--myServerFn2_createServerFn_handler", opts => {
         return myServerFn2.__executeServer(opts);
       });
-      const myServerFn = createServerFn().handler(myServerFn_createServerFn_handler, myFunc);
       const myServerFn2 = createServerFn().handler(myServerFn2_createServerFn_handler, myFunc2);
       export { myServerFn_createServerFn_handler, myServerFn2_createServerFn_handler };"
     `)
@@ -903,4 +907,69 @@ describe('server function compilation', () => {
       export { asyncGenerator_1 };"
     `)
   })
+
+  test('server function references exported value', () => {
+    const code = `
+
+      import z from "zod";
+
+      export const zExampleSchema = z.object({
+        name: z.string(),
+        age: z.number(),
+      });
+
+      export const namedFunction = createServerFn(() => {
+        'use server'
+        zExampleSchema.safeParse({ name: 'test', age: 123 })
+        return 'hello'
+      })
+    `
+
+    const client = compileDirectives({ ...clientConfig, code })
+    const ssr = compileDirectives({ ...ssrConfig, code })
+    const server = compileDirectives({
+      ...serverConfig,
+      code,
+      filename:
+        ssr.directiveFnsById[Object.keys(ssr.directiveFnsById)[0]!]!
+          .extractedFilename,
+    })
+
+    expect(client.compiledResult.code).toMatchInlineSnapshot(`
+      "import { createClientRpc } from "my-rpc-lib-client";
+      import z from "zod";
+      export const zExampleSchema = z.object({
+        name: z.string(),
+        age: z.number()
+      });
+      const namedFunction_createServerFn = createClientRpc("test_ts--namedFunction_createServerFn");
+      export const namedFunction = createServerFn(namedFunction_createServerFn);"
+    `)
+    expect(ssr.compiledResult.code).toMatchInlineSnapshot(`
+      "import { createSsrRpc } from "my-rpc-lib-server";
+      import z from "zod";
+      export const zExampleSchema = z.object({
+        name: z.string(),
+        age: z.number()
+      });
+      const namedFunction_createServerFn = createSsrRpc("test_ts--namedFunction_createServerFn");
+      export const namedFunction = createServerFn(namedFunction_createServerFn);"
+    `)
+    expect(server.compiledResult.code).toMatchInlineSnapshot(`
+      "import { createServerRpc } from "my-rpc-lib-server";
+      import z from "zod";
+      const zExampleSchema = z.object({
+        name: z.string(),
+        age: z.number()
+      });
+      const namedFunction_createServerFn = createServerRpc("test_ts--namedFunction_createServerFn", () => {
+        zExampleSchema.safeParse({
+          name: 'test',
+          age: 123
+        });
+        return 'hello';
+      });
+      export { namedFunction_createServerFn };"
+    `)
+  })
 })