@tanstack/start-plugin-core 1.132.0-alpha.7 → 1.132.0-alpha.9

package/src/create-server-fn-plugin/compiler.ts

@@ -0,0 +1,456 @@
+/* eslint-disable import/no-commonjs */
+import * as t from '@babel/types'
+import { generateFromAst, parseAst } from '@tanstack/router-utils'
+import babel from '@babel/core'
+import {
+  deadCodeElimination,
+  findReferencedIdentifiers,
+} from 'babel-dead-code-elimination'
+import { handleCreateServerFn } from './handleCreateServerFn'
+
+type Binding =
+  | {
+      type: 'import'
+      source: string
+      importedName: string
+      resolvedKind?: Kind
+    }
+  | {
+      type: 'var'
+      init: t.Expression | null
+      resolvedKind?: Kind
+    }
+
+type ExportEntry =
+  | { tag: 'Normal'; name: string }
+  | { tag: 'Default'; name: string }
+  | { tag: 'Namespace'; name: string; targetId: string } // for `export * as ns from './x'`
+
+type Kind = 'None' | 'Root' | 'Builder' | 'ServerFn'
+
+interface ModuleInfo {
+  id: string
+  code: string
+  ast: ReturnType<typeof parseAst>
+  bindings: Map<string, Binding>
+  exports: Map<string, ExportEntry>
+}
+
+export class ServerFnCompiler {
+  private moduleCache = new Map<string, ModuleInfo>()
+  private resolvedLibId!: string
+  private initialized = false
+  constructor(
+    private options: {
+      env: 'client' | 'server'
+      libName: string
+      rootExport: string
+      loadModule: (id: string) => Promise<void>
+      resolveId: (id: string, importer?: string) => Promise<string | null>
+    },
+  ) {}
+
+  private async init(id: string) {
+    const libId = await this.options.resolveId(this.options.libName, id)
+    if (!libId) {
+      throw new Error(`could not resolve "${this.options.libName}"`)
+    }
+    // insert root binding
+    const rootModule = {
+      ast: null as any,
+      bindings: new Map(),
+      exports: new Map(),
+      code: '',
+      id: libId,
+    }
+    rootModule.exports.set(this.options.rootExport, {
+      tag: 'Normal',
+      name: this.options.rootExport,
+    })
+    rootModule.bindings.set(this.options.rootExport, {
+      type: 'var',
+      init: t.identifier(this.options.rootExport),
+      resolvedKind: 'Root',
+    })
+    this.moduleCache.set(libId, rootModule)
+    this.initialized = true
+    this.resolvedLibId = libId
+  }
+
+  public ingestModule({ code, id }: { code: string; id: string }) {
+    const ast = parseAst({ code })
+
+    const bindings = new Map<string, Binding>()
+    const exports = new Map<string, ExportEntry>()
+
+    // we are only interested in top-level bindings, hence we don't traverse the AST
+    // instead we only iterate over the program body
+    for (const node of ast.program.body) {
+      if (t.isImportDeclaration(node)) {
+        const source = node.source.value
+        for (const s of node.specifiers) {
+          if (t.isImportSpecifier(s)) {
+            const importedName = t.isIdentifier(s.imported)
+              ? s.imported.name
+              : s.imported.value
+            bindings.set(s.local.name, { type: 'import', source, importedName })
+          } else if (t.isImportDefaultSpecifier(s)) {
+            bindings.set(s.local.name, {
+              type: 'import',
+              source,
+              importedName: 'default',
+            })
+          } else if (t.isImportNamespaceSpecifier(s)) {
+            bindings.set(s.local.name, {
+              type: 'import',
+              source,
+              importedName: '*',
+            })
+          }
+        }
+      } else if (t.isVariableDeclaration(node)) {
+        for (const decl of node.declarations) {
+          if (t.isIdentifier(decl.id)) {
+            bindings.set(decl.id.name, {
+              type: 'var',
+              init: decl.init ?? null,
+            })
+          }
+        }
+      } else if (t.isExportNamedDeclaration(node)) {
+        // export const foo = ...
+        if (node.declaration) {
+          if (t.isVariableDeclaration(node.declaration)) {
+            for (const d of node.declaration.declarations) {
+              if (t.isIdentifier(d.id)) {
+                exports.set(d.id.name, { tag: 'Normal', name: d.id.name })
+                bindings.set(d.id.name, { type: 'var', init: d.init ?? null })
+              }
+            }
+          }
+        }
+        for (const sp of node.specifiers) {
+          if (t.isExportNamespaceSpecifier(sp)) {
+            exports.set(sp.exported.name, {
+              tag: 'Namespace',
+              name: sp.exported.name,
+              targetId: node.source?.value || '',
+            })
+          }
+          // export { local as exported }
+          else if (t.isExportSpecifier(sp)) {
+            const local = sp.local.name
+            const exported = t.isIdentifier(sp.exported)
+              ? sp.exported.name
+              : sp.exported.value
+            exports.set(exported, { tag: 'Normal', name: local })
+          }
+        }
+      } else if (t.isExportDefaultDeclaration(node)) {
+        const d = node.declaration
+        if (t.isIdentifier(d)) {
+          exports.set('default', { tag: 'Default', name: d.name })
+        } else {
+          const synth = '__default_export__'
+          bindings.set(synth, { type: 'var', init: d as t.Expression })
+          exports.set('default', { tag: 'Default', name: synth })
+        }
+      }
+    }
+
+    const info: ModuleInfo = { code, id, ast, bindings, exports }
+    this.moduleCache.set(id, info)
+    return info
+  }
+
+  public invalidateModule(id: string) {
+    return this.moduleCache.delete(id)
+  }
+
+  public async compile({ code, id }: { code: string; id: string }) {
+    if (!this.initialized) {
+      await this.init(id)
+    }
+    const { bindings, ast } = this.ingestModule({ code, id })
+    const candidates = this.collectHandlerCandidates(bindings)
+    if (candidates.length === 0) {
+      // this hook will only be invoked if there is `.handler(` in the code,
+      // so not discovering a handler candidate is rather unlikely, but maybe possible?
+      return null
+    }
+
+    // let's find out which of the candidates are actually server functions
+    const toRewrite: Array<t.CallExpression> = []
+    for (const handler of candidates) {
+      const kind = await this.resolveExprKind(handler, id)
+      if (kind === 'ServerFn') {
+        toRewrite.push(handler)
+      }
+    }
+    if (toRewrite.length === 0) {
+      return null
+    }
+    const pathsToRewrite: Array<babel.NodePath<t.CallExpression>> = []
+    babel.traverse(ast, {
+      CallExpression(path) {
+        const found = toRewrite.findIndex((h) => path.node === h)
+        if (found !== -1) {
+          pathsToRewrite.push(path)
+          // delete from toRewrite
+          toRewrite.splice(found, 1)
+        }
+      },
+    })
+
+    if (toRewrite.length > 0) {
+      throw new Error(
+        `Internal error: could not find all paths to rewrite. please file an issue`,
+      )
+    }
+
+    const refIdents = findReferencedIdentifiers(ast)
+
+    pathsToRewrite.map((p) =>
+      handleCreateServerFn(p, { env: this.options.env, code }),
+    )
+
+    deadCodeElimination(ast, refIdents)
+
+    return generateFromAst(ast, {
+      sourceMaps: true,
+      sourceFileName: id,
+      filename: id,
+    })
+  }
+
+  // collects all `.handler(...)` CallExpressions at top-level
+  private collectHandlerCandidates(bindings: Map<string, Binding>) {
+    const candidates: Array<t.CallExpression> = []
+
+    for (const binding of bindings.values()) {
+      if (binding.type === 'var') {
+        const handler = isHandlerCall(binding.init)
+        if (handler) {
+          candidates.push(handler)
+        }
+      }
+    }
+    return candidates
+  }
+
+  private async resolveIdentifierKind(
+    ident: string,
+    id: string,
+    visited = new Set<string>(),
+  ): Promise<Kind> {
+    const info = await this.getModuleInfo(id)
+
+    const binding = info.bindings.get(ident)
+    if (!binding) {
+      return 'None'
+    }
+    if (binding.resolvedKind) {
+      return binding.resolvedKind
+    }
+
+    // TODO improve cycle detection? should we throw here instead of returning 'None'?
+    // prevent cycles
+    const vKey = `${id}:${ident}`
+    if (visited.has(vKey)) {
+      return 'None'
+    }
+    visited.add(vKey)
+
+    const resolvedKind = await this.resolveBindingKind(binding, id, visited)
+    binding.resolvedKind = resolvedKind
+    return resolvedKind
+  }
+
+  private async resolveBindingKind(
+    binding: Binding,
+    fileId: string,
+    visited = new Set<string>(),
+  ): Promise<Kind> {
+    if (binding.resolvedKind) {
+      return binding.resolvedKind
+    }
+    if (binding.type === 'import') {
+      const target = await this.options.resolveId(binding.source, fileId)
+      if (!target) {
+        return 'None'
+      }
+
+      if (binding.importedName === '*') {
+        throw new Error(
+          `should never get here, namespace imports are handled in resolveCalleeKind`,
+        )
+      }
+
+      const importedModule = await this.getModuleInfo(target)
+
+      const moduleExport = importedModule.exports.get(binding.importedName)
+      if (!moduleExport) {
+        return 'None'
+      }
+      const importedBinding = importedModule.bindings.get(moduleExport.name)
+      if (!importedBinding) {
+        return 'None'
+      }
+      if (importedBinding.resolvedKind) {
+        return importedBinding.resolvedKind
+      }
+
+      const resolvedKind = await this.resolveBindingKind(
+        importedBinding,
+        importedModule.id,
+        visited,
+      )
+      importedBinding.resolvedKind = resolvedKind
+      return resolvedKind
+    }
+
+    const resolvedKind = await this.resolveExprKind(
+      binding.init,
+      fileId,
+      visited,
+    )
+    binding.resolvedKind = resolvedKind
+    return resolvedKind
+  }
+
+  private async resolveExprKind(
+    expr: t.Expression | null,
+    fileId: string,
+    visited = new Set<string>(),
+  ): Promise<Kind> {
+    if (!expr) {
+      return 'None'
+    }
+
+    let result: Kind = 'None'
+
+    if (t.isCallExpression(expr)) {
+      if (!t.isExpression(expr.callee)) {
+        return 'None'
+      }
+      const calleeKind = await this.resolveCalleeKind(
+        expr.callee,
+        fileId,
+        visited,
+      )
+      if (calleeKind === 'Root' || calleeKind === 'Builder') {
+        return 'Builder'
+      }
+      if (calleeKind === 'ServerFn') {
+        return 'ServerFn'
+      }
+    } else if (t.isMemberExpression(expr) && t.isIdentifier(expr.property)) {
+      result = await this.resolveCalleeKind(expr.object, fileId, visited)
+    }
+
+    if (result === 'None' && t.isIdentifier(expr)) {
+      result = await this.resolveIdentifierKind(expr.name, fileId, visited)
+    }
+
+    if (result === 'None' && t.isTSAsExpression(expr)) {
+      result = await this.resolveExprKind(expr.expression, fileId, visited)
+    }
+    if (result === 'None' && t.isTSNonNullExpression(expr)) {
+      result = await this.resolveExprKind(expr.expression, fileId, visited)
+    }
+    if (result === 'None' && t.isParenthesizedExpression(expr)) {
+      result = await this.resolveExprKind(expr.expression, fileId, visited)
+    }
+
+    return result
+  }
+
+  private async resolveCalleeKind(
+    callee: t.Expression,
+    fileId: string,
+    visited = new Set<string>(),
+  ): Promise<Kind> {
+    if (t.isIdentifier(callee)) {
+      return this.resolveIdentifierKind(callee.name, fileId, visited)
+    }
+
+    if (t.isMemberExpression(callee) && t.isIdentifier(callee.property)) {
+      const prop = callee.property.name
+
+      if (prop === 'handler') {
+        const base = await this.resolveExprKind(callee.object, fileId, visited)
+        if (base === 'Root' || base === 'Builder') {
+          return 'ServerFn'
+        }
+        return 'None'
+      }
+      // Check if the object is a namespace import
+      if (t.isIdentifier(callee.object)) {
+        const info = await this.getModuleInfo(fileId)
+        const binding = info.bindings.get(callee.object.name)
+        if (
+          binding &&
+          binding.type === 'import' &&
+          binding.importedName === '*'
+        ) {
+          // resolve the property from the target module
+          const targetModuleId = await this.options.resolveId(
+            binding.source,
+            fileId,
+          )
+          if (targetModuleId) {
+            const targetModule = await this.getModuleInfo(targetModuleId)
+            const exportEntry = targetModule.exports.get(callee.property.name)
+            if (exportEntry) {
+              const exportedBinding = targetModule.bindings.get(
+                exportEntry.name,
+              )
+              if (exportedBinding) {
+                return await this.resolveBindingKind(
+                  exportedBinding,
+                  targetModule.id,
+                  visited,
+                )
+              }
+            }
+          }
+        }
+      }
+      return this.resolveExprKind(callee.object, fileId, visited)
+    }
+
+    // handle nested expressions
+    return this.resolveExprKind(callee, fileId, visited)
+  }
+
+  private async getModuleInfo(id: string) {
+    let cached = this.moduleCache.get(id)
+    if (cached) {
+      return cached
+    }
+
+    await this.options.loadModule(id)
+
+    cached = this.moduleCache.get(id)
+    if (!cached) {
+      throw new Error(`could not load module info for ${id}`)
+    }
+    return cached
+  }
+}
+
+function isHandlerCall(
+  node: t.Node | null | undefined,
+): undefined | t.CallExpression {
+  if (!t.isCallExpression(node)) return undefined
+
+  const callee = node.callee
+  if (
+    !t.isMemberExpression(callee) ||
+    !t.isIdentifier(callee.property, { name: 'handler' })
+  ) {
+    return undefined
+  }
+
+  return node
+}

package/src/{start-compiler-plugin/serverFn.ts → create-server-fn-plugin/handleCreateServerFn.ts}

@@ -1,39 +1,32 @@
 import * as t from '@babel/types'
-import { codeFrameError, getRootCallExpression } from './utils'
+import {
+  codeFrameError,
+  getRootCallExpression,
+} from '../start-compiler-plugin/utils'
 import type * as babel from '@babel/core'
 
-import type { CompileOptions } from './compilers'
-
-export function handleCreateServerFnCallExpression(
+export function handleCreateServerFn(
   path: babel.NodePath<t.CallExpression>,
-  opts: CompileOptions,
+  opts: {
+    env: 'client' | 'server'
+    code: string
+  },
 ) {
   // Traverse the member expression and find the call expressions for
   // the validator, handler, and middleware methods. Check to make sure they
   // are children of the createServerFn call expression.
 
-  const calledOptions = path.node.arguments[0]
-    ? (path.get('arguments.0') as babel.NodePath<t.ObjectExpression>)
-    : null
-
-  const shouldValidateClient = !!calledOptions?.node.properties.find((prop) => {
-    return (
-      t.isObjectProperty(prop) &&
-      t.isIdentifier(prop.key) &&
-      prop.key.name === 'validateClient' &&
-      t.isBooleanLiteral(prop.value) &&
-      prop.value.value === true
-    )
-  })
-
-  const callExpressionPaths = {
-    middleware: null as babel.NodePath<t.CallExpression> | null,
-    validator: null as babel.NodePath<t.CallExpression> | null,
-    handler: null as babel.NodePath<t.CallExpression> | null,
+  const validMethods = ['middleware', 'validator', 'handler'] as const
+  type ValidMethods = (typeof validMethods)[number]
+  const callExpressionPaths: Record<
+    ValidMethods,
+    babel.NodePath<t.CallExpression> | null
+  > = {
+    middleware: null,
+    validator: null,
+    handler: null,
   }
 
-  const validMethods = Object.keys(callExpressionPaths)
-
   const rootCallExpression = getRootCallExpression(path)
 
   // if (debug)
@@ -54,8 +47,7 @@ export function handleCreateServerFnCallExpression(
   rootCallExpression.traverse({
     MemberExpression(memberExpressionPath) {
       if (t.isIdentifier(memberExpressionPath.node.property)) {
-        const name = memberExpressionPath.node.property
-          .name as keyof typeof callExpressionPaths
+        const name = memberExpressionPath.node.property.name as ValidMethods
 
         if (
           validMethods.includes(name) &&
@@ -76,15 +68,13 @@ export function handleCreateServerFnCallExpression(
       )
     }
 
-    // If we're on the client, and we're not validating the client, remove the validator call expression
-    if (
-      opts.env === 'client' &&
-      !shouldValidateClient &&
-      t.isMemberExpression(callExpressionPaths.validator.node.callee)
-    ) {
-      callExpressionPaths.validator.replaceWith(
-        callExpressionPaths.validator.node.callee.object,
-      )
+    // If we're on the client, remove the validator call expression
+    if (opts.env === 'client') {
+      if (t.isMemberExpression(callExpressionPaths.validator.node.callee)) {
+        callExpressionPaths.validator.replaceWith(
+          callExpressionPaths.validator.node.callee.object,
+        )
+      }
     }
   }
 

package/src/create-server-fn-plugin/plugin.ts

@@ -0,0 +1,138 @@
+import { VITE_ENVIRONMENT_NAMES } from '../constants'
+import { ServerFnCompiler } from './compiler'
+import type { ViteEnvironmentNames } from '../constants'
+import type { PluginOption } from 'vite'
+import type { CompileStartFrameworkOptions } from '../start-compiler-plugin/compilers'
+
+function cleanId(id: string): string {
+  return id.split('?')[0]!
+}
+
+export function createServerFnPlugin(
+  framework: CompileStartFrameworkOptions,
+): PluginOption {
+  const libName = `@tanstack/${framework}-start`
+  const rootExport = 'createServerFn'
+
+  const SERVER_FN_LOOKUP = 'server-fn-module-lookup'
+
+  const compilers: Partial<Record<ViteEnvironmentNames, ServerFnCompiler>> = {}
+  return [
+    {
+      name: 'tanstack-start-core:capture-server-fn-module-lookup',
+      // we only need this plugin in dev mode
+      apply: 'serve',
+      applyToEnvironment(env) {
+        return [
+          VITE_ENVIRONMENT_NAMES.client,
+          VITE_ENVIRONMENT_NAMES.server,
+        ].includes(env.name as ViteEnvironmentNames)
+      },
+      transform: {
+        filter: {
+          id: new RegExp(`${SERVER_FN_LOOKUP}$`),
+        },
+        handler(code, id) {
+          const compiler =
+            compilers[this.environment.name as ViteEnvironmentNames]
+          compiler?.ingestModule({ code, id: cleanId(id) })
+        },
+      },
+    },
+    {
+      name: 'tanstack-start-core::server-fn',
+      enforce: 'pre',
+
+      applyToEnvironment(env) {
+        return [
+          VITE_ENVIRONMENT_NAMES.client,
+          VITE_ENVIRONMENT_NAMES.server,
+        ].includes(env.name as ViteEnvironmentNames)
+      },
+      transform: {
+        filter: {
+          id: {
+            exclude: new RegExp(`${SERVER_FN_LOOKUP}$`),
+          },
+          code: {
+            // only scan files that mention `.handler(`
+            include: [/\.handler\(/],
+          },
+        },
+        async handler(code, id) {
+          let compiler =
+            compilers[this.environment.name as ViteEnvironmentNames]
+          if (!compiler) {
+            const env =
+              this.environment.name === VITE_ENVIRONMENT_NAMES.client
+                ? 'client'
+                : this.environment.name === VITE_ENVIRONMENT_NAMES.server
+                  ? 'server'
+                  : (() => {
+                      throw new Error(
+                        `Environment ${this.environment.name} not configured`,
+                      )
+                    })()
+
+            compiler = new ServerFnCompiler({
+              env,
+              libName,
+              rootExport,
+              loadModule: async (id: string) => {
+                if (this.environment.mode === 'build') {
+                  const loaded = await this.load({ id })
+                  if (!loaded.code) {
+                    throw new Error(`could not load module ${id}`)
+                  }
+                  compiler!.ingestModule({ code: loaded.code, id })
+                } else if (this.environment.mode === 'dev') {
+                  /**
+                   * in dev, vite does not return code from `ctx.load()`
+                   * so instead, we need to take a different approach
+                   * we must force vite to load the module and run it through the vite plugin pipeline
+                   * we can do this by using the `fetchModule` method
+                   * the `captureServerFnModuleLookupPlugin` captures the module code via its transform hook and invokes analyzeModuleAST
+                   */
+                  await this.environment.fetchModule(
+                    id + '?' + SERVER_FN_LOOKUP,
+                  )
+                } else {
+                  throw new Error(
+                    `could not load module ${id}: unknown environment mode ${this.environment.mode}`,
+                  )
+                }
+              },
+              resolveId: async (source: string, importer?: string) => {
+                const r = await this.resolve(source, importer)
+                return r ? cleanId(r.id) : null
+              },
+            })
+            compilers[this.environment.name as ViteEnvironmentNames] = compiler
+          }
+
+          id = cleanId(id)
+          const result = await compiler.compile({ id, code })
+          return result
+        },
+      },
+
+      hotUpdate(ctx) {
+        const compiler =
+          compilers[this.environment.name as ViteEnvironmentNames]
+
+        ctx.modules.forEach((m) => {
+          if (m.id) {
+            const deleted = compiler?.invalidateModule(m.id)
+            if (deleted) {
+              m.importers.forEach((importer) => {
+                if (importer.id) {
+                  compiler?.invalidateModule(importer.id)
+                }
+              })
+            }
+          }
+        })
+      },
+    },
+  ]
+}

package/src/plugin.ts

@@ -18,6 +18,7 @@ import {
   getServerOutputDirectory,
 } from './output-directory'
 import { postServerBuild } from './post-server-build'
+import { createServerFnPlugin } from './create-server-fn-plugin/plugin'
 import type { ViteEnvironmentNames } from './constants'
 import type { TanStackStartInputConfig } from './schema'
 import type { PluginOption } from 'vite'
@@ -289,11 +290,9 @@ export function TanStackStartVitePluginCore(
         }
       },
     },
+    createServerFnPlugin(corePluginOpts.framework),
     // N.B. TanStackStartCompilerPlugin must be before the TanStackServerFnPluginEnv
-    startCompilerPlugin(corePluginOpts.framework, {
-      client: { envName: VITE_ENVIRONMENT_NAMES.client },
-      server: { envName: VITE_ENVIRONMENT_NAMES.server },
-    }),
+    startCompilerPlugin(corePluginOpts.framework),
     TanStackServerFnPluginEnv({
       // This is the ID that will be available to look up and import
       // our server function manifest and resolve its module