Node RSC on-demand transform / resolution for RSC

loader.ts

import { readFile } from 'node:fs/promises'
import * as path from 'node:path'
import { type LoadHook, type ResolveHook, type ResolveHookContext } from 'node:module'
import { fileURLToPath } from 'node:url'

import * as oxy from '@oxidation-compiler/napi'

import type { ModuleExport } from './module-info.js'

import * as clientTransforms from './transform-client.js'
import * as serverTransforms from './transform-server.js'

export const resolve: ResolveHook = async (specifier, context, nextResolve) => {
	if (!context.parentURL) return nextResolve(specifier, context)

	const parentURL = new URL(context.parentURL)
	const specifierURL = new URL(specifier, context.parentURL)
	const graph =
		specifierURL.searchParams.get('graph') ||
		parentURL.searchParams.get('graph') ||
		(specifierURL.searchParams.has('server') && 'server') ||
		(specifierURL.searchParams.has('client') && 'client')

	if (!graph) return nextResolve(specifier, context)

	const cleanParentURL = new URL(context.parentURL)
	cleanParentURL.hash = ''
	cleanParentURL.searchParams.delete('graph')
	cleanParentURL.searchParams.delete('client')
	cleanParentURL.searchParams.delete('server')
	const resolveContext: ResolveHookContext = {
		conditions: [...context.conditions],
		importAssertions: context.importAssertions,
		parentURL: cleanParentURL.href,
	}

	if (graph === 'server') {
		resolveContext.conditions.unshift('react-server')
	}

	const resolved = await nextResolve(specifier, resolveContext)

	if (specifier === '#conditional-fixture') {
		console.log({ specifier, resolved, resolveContext })
	}

	const url = new URL(resolved.url)
	url.searchParams.set('graph', graph)

	const filePath = fileURLToPath(resolved.url)
	const source = await readFile(filePath, 'utf8')

	const parseResult = await oxy.parseAsync(source, {
		sourceFilename: filePath,
		sourceType: 'module',
	})

	const program = JSON.parse(parseResult.program)
	const identifier = path.relative(process.cwd(), filePath)

	let useClient = false,
		useServer = false
	for (const { directive } of program.directives) {
		if (directive === 'use client') {
			useClient = true
			url.searchParams.set('client-module', '')
		} else if (directive === 'use server') {
			useServer = true
			url.searchParams.set('server-module', '')
		}
	}

	resolved.url = url.href

	if (!useClient && !useServer) return resolved
	if (useClient && useServer) {
		throw new Error(`Cannot use both "use client" and "use server" in the same module`)
	}

	for (const node of program.body) {
		// Handle FunctionDeclaration exports
		if (
			node.type === 'ExportNamedDeclaration' &&
			node.declaration?.type === 'FunctionDeclaration'
		) {
			const name = node.declaration.id.name
			url.searchParams.append(
				'exports',
				JSON.stringify({
					identifier,
					localName: name,
					publicName: name,
				} satisfies ModuleExport)
			)
		}

		// Handle export specifiers
		if (node.type === 'ExportNamedDeclaration' && node.specifiers.length > 0) {
			for (const specifier of node.specifiers) {
				const localName = specifier.local.name
				const publicName = specifier.exported.name
				url.searchParams.append(
					'exports',
					JSON.stringify({
						identifier,
						localName,
						publicName,
					} satisfies ModuleExport)
				)
			}
		}

		// Handle default exports
		if (node.type === 'ExportDefaultDeclaration') {
			if (!node.declaration.id) throw new Error(`Cannot export anonymous default export`)
			const name = node.declaration.id.name

			url.searchParams.append(
				'exports',
				JSON.stringify({
					identifier,
					localName: name,
					publicName: 'default',
				} satisfies ModuleExport)
			)
		}
	}

	resolved.url = url.href

	return resolved
}

export const load: LoadHook = async (urlString, context, defaultLoad) => {
	const url = new URL(urlString)

	const urlWithoutCustomSearchParams = new URL(urlString)
	urlWithoutCustomSearchParams.searchParams.delete('graph')
	urlWithoutCustomSearchParams.searchParams.delete('client-module')
	urlWithoutCustomSearchParams.searchParams.delete('server-module')
	const loaded = await defaultLoad(urlWithoutCustomSearchParams.href, context)

	const graph = url.searchParams.get('graph')
	const clientModule = url.searchParams.has('client-module')
	const serverModule = url.searchParams.has('server-module')
	if (!graph || (!clientModule && !serverModule) || !loaded.source || loaded.format !== 'module')
		return loaded

	let source: string | undefined
	switch (typeof loaded.source) {
		case 'string':
			source = loaded.source
			break
		case 'object':
			source = Buffer.from(loaded.source).toString('utf8')
			break
	}
	if (typeof source !== 'string') throw new Error(`Unexpected source type: ${typeof source}`)

	const moduleExports: ModuleExport[] = []
	for (const exportString of url.searchParams.getAll('exports')) {
		moduleExports.push(JSON.parse(exportString))
	}

	switch (graph) {
		case 'client':
			if (serverModule) {
				loaded.source = clientTransforms.createServerModule(moduleExports)
			}
			break
		case 'server': {
			if (clientModule) {
				loaded.source = serverTransforms.createClientModule(moduleExports)
			} else if (serverModule) {
				loaded.source += '\n;' + serverTransforms.createServerModuleFooter(moduleExports)
			}
			break
		}
	}

	return loaded
}

module-info.ts

export interface ModuleExport {
	identifier: string
  localName: string
	publicName: string
}

transform-client.ts

import { js } from './template-strings.js'
import { type ModuleExport } from './module-info.js'

export function createServerModule(moduleExports: Iterable<ModuleExport>): string {
	let code = ''

	let seen = new Set<string>()
	for (const { identifier, publicName } of moduleExports) {
		if (seen.has(publicName)) {
			throw new Error(`Duplicate export name: ${publicName}`)
		}

		const serverReference = js`{
      $$typeof: Symbol.for('react.server.reference'),
      $$id: ${JSON.stringify(identifier)},
    }`

		if (publicName === 'default') {
			code += js`
        export default ${serverReference};
      `
		} else {
			code += js`
        export const ${publicName} = ${serverReference};
      `
		}
	}

	return code
}

transform-server.ts

import { js } from './template-strings.js'
import { type ModuleExport } from './module-info.js'

export function createClientModule(moduleExports: Iterable<ModuleExport>): string {
	let code = ''

	let seen = new Set<string>()
	for (const { publicName, identifier } of moduleExports) {
		if (seen.has(publicName)) {
			throw new Error(`Duplicate export name: ${publicName}`)
		}

		const serverReference = js`{
      $$typeof: Symbol.for('react.client.reference'),
      $$id: ${JSON.stringify(identifier)},
    }`

		if (publicName === 'default') {
			code += js`
        export default ${serverReference};
      `
		} else {
			code += js`
        export const ${publicName} = ${serverReference};
      `
		}
	}

	return code
}

export function createServerModuleFooter(moduleExports: Iterable<ModuleExport>): string {
	let code = ''

	const seenPublicNames = new Set<string>()
	const seenLocalNames = new Set<string>()
	for (const { identifier, localName, publicName } of moduleExports) {
		if (seenPublicNames.has(publicName)) {
			throw new Error(`Duplicate export name: ${publicName}`)
		}
		if (seenLocalNames.has(localName)) {
			continue
		}

		code += js`
      if (typeof ${localName} === 'function') {
        Object.defineProperties(${localName}, {
          $$typeof: { value: Symbol.for('react.server.reference') },
          $$id: { value: ${JSON.stringify(identifier)} },
        })
      }
    `
	}

	return code
}