feat(core): new transform for client references (#1201)

This enables an experimental `allowServer`. I'm not sure if this is going to be finalized.
dai-shi · Jan 30, 2025 · c0e1f3f · c0e1f3f
1 parent b3fe435
commit c0e1f3f
Show file tree

Hide file tree

Showing 4 changed files with 244 additions and 19 deletions.
diff --git a/packages/waku/src/client.ts b/packages/waku/src/client.ts
@@ -1,2 +1,4 @@
-// empty for now
-export default {};
+/**
+ * Highly experimental, the name might change.
+ */
+export const allowServer = <T>(x: T) => x;
diff --git a/packages/waku/src/lib/plugins/vite-plugin-rsc-transform.ts b/packages/waku/src/lib/plugins/vite-plugin-rsc-transform.ts
@@ -158,6 +158,189 @@ const replaceNode = <T extends swc.Node>(origNode: swc.Node, newNode: T): T => {
   return Object.assign(origNode, newNode);
 };
 
+const transformExportedClientThings = (
+  mod: swc.Module,
+  getFuncId: () => string,
+): Set<string> => {
+  const exportNames = new Set<string>();
+  // HACK this doesn't cover all cases
+  const allowServerItems = new Map<string, swc.Expression>();
+  const allowServerDependencies = new Set<string>();
+  const visited = new WeakSet<swc.Node>();
+  const findDependencies = (node: swc.Node) => {
+    if (visited.has(node)) {
+      return;
+    }
+    visited.add(node);
+    if (node.type === 'Identifier') {
+      const id = node as swc.Identifier;
+      if (!allowServerItems.has(id.value) && !exportNames.has(id.value)) {
+        allowServerDependencies.add(id.value);
+      }
+    }
+    Object.values(node).forEach((value) => {
+      (Array.isArray(value) ? value : [value]).forEach((v) => {
+        if (typeof v?.type === 'string') {
+          findDependencies(v);
+        } else if (typeof v?.expression?.type === 'string') {
+          findDependencies(v.expression);
+        }
+      });
+    });
+  };
+  // Pass 1: find allowServer identifier
+  let allowServer = 'allowServer';
+  for (const item of mod.body) {
+    if (item.type === 'ImportDeclaration') {
+      if (item.source.value === 'waku/client') {
+        for (const specifier of item.specifiers) {
+          if (specifier.type === 'ImportSpecifier') {
+            if (specifier.local.value === allowServer && specifier.imported) {
+              allowServer = specifier.imported.value;
+              break;
+            }
+          }
+        }
+        break;
+      }
+    }
+  }
+  // Pass 2: collect export names and allowServer names
+  for (const item of mod.body) {
+    if (item.type === 'ExportDeclaration') {
+      if (item.declaration.type === 'FunctionDeclaration') {
+        exportNames.add(item.declaration.identifier.value);
+      } else if (item.declaration.type === 'ClassDeclaration') {
+        exportNames.add(item.declaration.identifier.value);
+      } else if (item.declaration.type === 'VariableDeclaration') {
+        for (const d of item.declaration.declarations) {
+          if (d.id.type === 'Identifier') {
+            if (
+              d.init?.type === 'CallExpression' &&
+              d.init.callee.type === 'Identifier' &&
+              d.init.callee.value === allowServer
+            ) {
+              if (d.init.arguments.length !== 1) {
+                throw new Error('allowServer should have exactly one argument');
+              }
+              allowServerItems.set(d.id.value, d.init.arguments[0]!.expression);
+              findDependencies(d.init);
+            } else {
+              exportNames.add(d.id.value);
+            }
+          }
+        }
+      }
+    } else if (item.type === 'ExportNamedDeclaration') {
+      for (const s of item.specifiers) {
+        if (s.type === 'ExportSpecifier') {
+          exportNames.add(s.exported ? s.exported.value : s.orig.value);
+        }
+      }
+    } else if (item.type === 'ExportDefaultExpression') {
+      exportNames.add('default');
+    } else if (item.type === 'ExportDefaultDeclaration') {
+      exportNames.add('default');
+    }
+  }
+  // Pass 3: collect dependencies
+  let dependenciesSize: number;
+  do {
+    dependenciesSize = allowServerDependencies.size;
+    for (const item of mod.body) {
+      if (item.type === 'VariableDeclaration') {
+        for (const d of item.declarations) {
+          if (
+            d.id.type === 'Identifier' &&
+            allowServerDependencies.has(d.id.value)
+          ) {
+            findDependencies(d);
+          }
+        }
+      } else if (item.type === 'FunctionDeclaration') {
+        if (allowServerDependencies.has(item.identifier.value)) {
+          findDependencies(item);
+        }
+      } else if (item.type === 'ClassDeclaration') {
+        if (allowServerDependencies.has(item.identifier.value)) {
+          findDependencies(item);
+        }
+      }
+    }
+  } while (dependenciesSize < allowServerDependencies.size);
+  allowServerDependencies.delete(allowServer);
+  // Pass 4: filter with dependencies
+  for (let i = 0; i < mod.body.length; ++i) {
+    const item = mod.body[i]!;
+    if (
+      item.type === 'ImportDeclaration' &&
+      item.specifiers.some(
+        (s) =>
+          s.type === 'ImportSpecifier' &&
+          allowServerDependencies.has(
+            s.imported ? s.imported.value : s.local.value,
+          ),
+      )
+    ) {
+      continue;
+    }
+    if (item.type === 'VariableDeclaration') {
+      item.declarations = item.declarations.filter(
+        (d) =>
+          d.id.type === 'Identifier' && allowServerDependencies.has(d.id.value),
+      );
+      if (item.declarations.length) {
+        continue;
+      }
+    }
+    if (item.type === 'FunctionDeclaration') {
+      if (allowServerDependencies.has(item.identifier.value)) {
+        continue;
+      }
+    }
+    if (item.type === 'ClassDeclaration') {
+      if (allowServerDependencies.has(item.identifier.value)) {
+        continue;
+      }
+    }
+    mod.body.splice(i--, 1);
+  }
+  // Pass 5: add allowServer exports
+  for (const [allowServerName, callExp] of allowServerItems) {
+    const stmt: swc.ExportDeclaration = {
+      type: 'ExportDeclaration',
+      declaration: {
+        type: 'VariableDeclaration',
+        kind: 'const',
+        declare: false,
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-expect-error
+        ctxt: 0,
+        declarations: [
+          {
+            type: 'VariableDeclarator',
+            id: createIdentifier(allowServerName),
+            init: createCallExpression(
+              createIdentifier('__waku_registerClientReference'),
+              [
+                callExp,
+                createStringLiteral(getFuncId()),
+                createStringLiteral(allowServerName),
+              ],
+            ),
+            definite: false,
+            span: createEmptySpan(),
+          },
+        ],
+        span: createEmptySpan(),
+      },
+      span: createEmptySpan(),
+    };
+    mod.body.push(stmt);
+  }
+  return exportNames;
+};
+
 const transformExportedServerFunctions = (
   mod: swc.Module,
   getFuncId: () => string,
@@ -580,13 +763,14 @@ const transformServer = (
     }
   }
   if (hasUseClient) {
-    const exportNames = collectExportNames(mod);
+    const exportNames = transformExportedClientThings(mod, getClientId);
     let newCode = `
-import { registerClientReference } from 'react-server-dom-webpack/server.edge';
+import { registerClientReference as __waku_registerClientReference } from 'react-server-dom-webpack/server.edge';
 `;
+    newCode += swc.printSync(mod).code;
     for (const name of exportNames) {
       newCode += `
-export ${name === 'default' ? name : `const ${name} =`} registerClientReference(() => { throw new Error('It is not possible to invoke a client function from the server: ${getClientId()}#${name}'); }, '${getClientId()}', '${name}');
+export ${name === 'default' ? name : `const ${name} =`} __waku_registerClientReference(() => { throw new Error('It is not possible to invoke a client function from the server: ${getClientId()}#${name}'); }, '${getClientId()}', '${name}');
 `;
     }
     return newCode;

diff --git a/packages/waku/src/lib/utils/treeshake.ts b/packages/waku/src/lib/utils/treeshake.ts
@@ -4,15 +4,16 @@ import * as swc from '@swc/core';
 export const treeshake = async (
   code: string,
   modifyModule?: (mod: swc.Module) => void,
+  tsx = false,
 ): Promise<string> => {
-  const mod = swc.parseSync(code, { syntax: 'typescript' });
+  const mod = swc.parseSync(code, { syntax: 'typescript', tsx });
   modifyModule?.(mod);
-  code = swc.printSync(mod).code;
-  // FIXME can we avoid this and transform with printSync directly?
-  code = swc.transformSync(code, {
-    jsc: { parser: { syntax: 'typescript' } },
+  const jsCode = swc.transformSync(mod, {
+    jsc: {
+      target: 'esnext',
+      parser: { syntax: 'typescript', tsx },
+    },
   }).code;
-
   const bundle = await rollup({
     input: '\0code',
     external: () => true,
@@ -22,8 +23,12 @@ export const treeshake = async (
       }
       defaultHandler(warning);
     },
+    output: {
+      generatedCode: 'es2015',
+    },
     treeshake: {
       moduleSideEffects: 'no-external',
+      propertyReadSideEffects: false,
     },
     plugins: [
       {
@@ -35,7 +40,7 @@ export const treeshake = async (
         },
         load(id) {
           if (id === '\0code') {
-            return code;
+            return jsCode;
           }
         },
         resolveDynamicImport(id) {

diff --git a/packages/waku/tests/vite-plugin-rsc-transform-internals.test.ts b/packages/waku/tests/vite-plugin-rsc-transform-internals.test.ts
@@ -30,13 +30,24 @@ export default function App() {
     const code = `
 'use client';
 
-import { Component } from 'react';
+import { Component, createContext, useContext, memo } from 'react';
+import { atom } from 'jotai/vanilla';
+import { allowServer } from 'waku/client';
+
+const initialCount = 1;
+const TWO = 2;
+function double (x: number) {
+  return x * TWO;
+}
+export const countAtom = allowServer(atom(double(initialCount)));
 
 export const Empty = () => null;
 
 function Private() {
   return "Secret";
 }
+const SecretComponent = () => <p>Secret</p>;
+const SecretFunction = (n: number) => 'Secret' + n;
 
 export function Greet({ name }: { name: string }) {
   return <>Hello {name}</>;
@@ -48,22 +59,45 @@ export class MyComponent extends Component {
   }
 }
 
+const MyContext = createContext();
+
+export const useMyContext = () => useContext(MyContext);
+
+const MyProvider = memo(MyContext.Provider);
+
+export const NAME = 'World';
+
 export default function App() {
-  return <div>Hello World</div>;
+  return (
+    <MyProvider value="Hello">
+      <div>Hello World</div>
+    </MyProvider>
+  );
 }
 `;
     expect(await transform(code, '/src/App.tsx', { ssr: true }))
       .toMatchInlineSnapshot(`
         "
-        import { registerClientReference } from 'react-server-dom-webpack/server.edge';
+        import { registerClientReference as __waku_registerClientReference } from 'react-server-dom-webpack/server.edge';
+        import { atom } from 'jotai/vanilla';
+        const initialCount = 1;
+        const TWO = 2;
+        function double(x: number) {
+            return x * TWO;
+        }
+        export const countAtom = __waku_registerClientReference(atom(double(initialCount)), "/src/App.tsx", "countAtom");
+
+        export const Empty = __waku_registerClientReference(() => { throw new Error('It is not possible to invoke a client function from the server: /src/App.tsx#Empty'); }, '/src/App.tsx', 'Empty');
+
+        export const Greet = __waku_registerClientReference(() => { throw new Error('It is not possible to invoke a client function from the server: /src/App.tsx#Greet'); }, '/src/App.tsx', 'Greet');
 
-        export const Empty = registerClientReference(() => { throw new Error('It is not possible to invoke a client function from the server: /src/App.tsx#Empty'); }, '/src/App.tsx', 'Empty');
+        export const MyComponent = __waku_registerClientReference(() => { throw new Error('It is not possible to invoke a client function from the server: /src/App.tsx#MyComponent'); }, '/src/App.tsx', 'MyComponent');
 
-        export const Greet = registerClientReference(() => { throw new Error('It is not possible to invoke a client function from the server: /src/App.tsx#Greet'); }, '/src/App.tsx', 'Greet');
+        export const useMyContext = __waku_registerClientReference(() => { throw new Error('It is not possible to invoke a client function from the server: /src/App.tsx#useMyContext'); }, '/src/App.tsx', 'useMyContext');
 
-        export const MyComponent = registerClientReference(() => { throw new Error('It is not possible to invoke a client function from the server: /src/App.tsx#MyComponent'); }, '/src/App.tsx', 'MyComponent');
+        export const NAME = __waku_registerClientReference(() => { throw new Error('It is not possible to invoke a client function from the server: /src/App.tsx#NAME'); }, '/src/App.tsx', 'NAME');
 
-        export default registerClientReference(() => { throw new Error('It is not possible to invoke a client function from the server: /src/App.tsx#default'); }, '/src/App.tsx', 'default');
+        export default __waku_registerClientReference(() => { throw new Error('It is not possible to invoke a client function from the server: /src/App.tsx#default'); }, '/src/App.tsx', 'default');
         "
       `);
   });