Skip to content

Commit ac2325a

Browse files
committed
feat(rsc): generalize server function transforms
1 parent d8e0f4f commit ac2325a

5 files changed

Lines changed: 297 additions & 39 deletions

File tree

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

Lines changed: 78 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import path from 'node:path'
22
import { parseAstAsync } from 'vite'
3-
import { describe, expect, it } from 'vitest'
3+
import { describe, expect, it, vi } from 'vitest'
44
import { transformHoistInlineDirective } from './hoist'
55
import { debugSourceMap } from './test-utils'
66

@@ -466,11 +466,11 @@ export async function kv() {
466466
}),
467467
).toMatchInlineSnapshot(`
468468
"
469-
export const none = /* #__PURE__ */ $$register($$hoist_0_none, "<id>", "$$hoist_0_none", {"directiveMatch":["use cache",null]});
469+
export const none = /* #__PURE__ */ $$register($$hoist_0_none, "<id>", "$$hoist_0_none", {"directiveMatch":["use cache",null],"hasBoundArgs":false});
470470
471-
export const fs = /* #__PURE__ */ $$register($$hoist_1_fs, "<id>", "$$hoist_1_fs", {"directiveMatch":["use cache: fs",": fs"]});
471+
export const fs = /* #__PURE__ */ $$register($$hoist_1_fs, "<id>", "$$hoist_1_fs", {"directiveMatch":["use cache: fs",": fs"],"hasBoundArgs":false});
472472
473-
export const kv = /* #__PURE__ */ $$register($$hoist_2_kv, "<id>", "$$hoist_2_kv", {"directiveMatch":["use cache: kv",": kv"]});
473+
export const kv = /* #__PURE__ */ $$register($$hoist_2_kv, "<id>", "$$hoist_2_kv", {"directiveMatch":["use cache: kv",": kv"],"hasBoundArgs":false});
474474
475475
;async function $$hoist_0_none() {
476476
"use cache";
@@ -508,4 +508,78 @@ export async function test() {
508508
"
509509
`)
510510
})
511+
512+
it('uses stable names and reports bound arguments', async () => {
513+
const input = `
514+
async function outer(value) {
515+
return async function cached() {
516+
"use cache";
517+
return value;
518+
};
519+
}
520+
`
521+
const ast = await parseAstAsync(input)
522+
const runtime = vi.fn((value: string) => value)
523+
const first = transformHoistInlineDirective(input, ast, {
524+
directive: 'use cache',
525+
stableName: true,
526+
runtime,
527+
})
528+
const second = transformHoistInlineDirective(input, ast, {
529+
directive: 'use cache',
530+
stableName: true,
531+
runtime: (value) => value,
532+
})
533+
expect(runtime).toHaveBeenCalledWith(
534+
expect.any(String),
535+
expect.any(String),
536+
expect.objectContaining({ hasBoundArgs: true }),
537+
)
538+
expect(first.names).toEqual(second.names)
539+
expect(first.names[0]).toMatch(/^\$\$hoist_[a-f0-9]{12}_0_cached$/)
540+
})
541+
542+
it('supports object and static class methods', async () => {
543+
const input = `
544+
const object = {
545+
async ["cached"]() { "use cache"; return 1 },
546+
};
547+
class Cache {
548+
static async ["cached"]() { "use cache"; return 2 }
549+
}
550+
`
551+
const transformed = await testTransform(input, {
552+
directive: 'use cache',
553+
})
554+
expect(transformed).toContain('"cached": /* #__PURE__ */')
555+
expect(transformed).toContain('static ["cached"] = /* #__PURE__ */')
556+
await parseAstAsync(transformed!)
557+
})
558+
559+
it('rejects instance, private, getter and setter methods', async () => {
560+
for (const input of [
561+
`class Cache { async cached() { "use cache" } }`,
562+
`class Cache { static async #cached() { "use cache" } }`,
563+
`const cache = { get cached() { "use cache"; return 1 } }`,
564+
`class Cache { static set cached(value) { "use cache" } }`,
565+
]) {
566+
await expect(
567+
testTransform(input, { directive: 'use cache' }),
568+
).rejects.toThrow(/not allowed/)
569+
}
570+
})
571+
572+
it('supports stateful directive regexes repeatedly', async () => {
573+
const directive = /^use cache(?:: .+)?$/gy
574+
const input = `
575+
export async function one() { "use cache" }
576+
export async function two() { "use cache: remote" }
577+
`
578+
expect(await testTransform(input, { directive, noExport: true })).toContain(
579+
'use cache: remote',
580+
)
581+
expect(await testTransform(input, { directive, noExport: true })).toContain(
582+
'use cache: remote',
583+
)
584+
})
511585
})

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

Lines changed: 77 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import type {
88
} from 'estree'
99
import { walk } from 'estree-walker'
1010
import MagicString from 'magic-string'
11+
import { hashString } from '../plugins/utils'
1112
import { buildScopeTree, type ScopeTree } from './scope'
1213

1314
export function transformHoistInlineDirective(
@@ -21,13 +22,14 @@ export function transformHoistInlineDirective(
2122
runtime: (
2223
value: string,
2324
name: string,
24-
meta: { directiveMatch: RegExpMatchArray },
25+
meta: { directiveMatch: RegExpMatchArray; hasBoundArgs: boolean },
2526
) => string
2627
directive: string | RegExp
2728
rejectNonAsyncFunction?: boolean
2829
encode?: (value: string) => string
2930
decode?: (value: string) => string
3031
noExport?: boolean
32+
stableName?: boolean
3133
},
3234
): {
3335
output: MagicString
@@ -45,6 +47,7 @@ export function transformHoistInlineDirective(
4547

4648
const scopeTree = buildScopeTree(ast)
4749
const names: string[] = []
50+
const signatureCounts = new Map<string, number>()
4851

4952
walk(ast, {
5053
enter(node, parent) {
@@ -65,12 +68,60 @@ export function transformHoistInlineDirective(
6568
)
6669
}
6770

71+
const isObjectMethod =
72+
node.type === 'FunctionExpression' &&
73+
parent?.type === 'Property' &&
74+
(parent.method || parent.kind !== 'init')
75+
const isClassMethod =
76+
node.type === 'FunctionExpression' &&
77+
parent?.type === 'MethodDefinition'
78+
if (isClassMethod && !parent.static) {
79+
throw Object.assign(
80+
new Error(
81+
`It is not allowed to define inline ${JSON.stringify(match[0])} annotated class instance methods. Use a function, object method property, or static class method instead.`,
82+
),
83+
{ pos: parent.start },
84+
)
85+
}
86+
if (isClassMethod && parent.key.type === 'PrivateIdentifier') {
87+
throw Object.assign(
88+
new Error(
89+
`It is not allowed to define inline ${JSON.stringify(match[0])} annotated private class methods.`,
90+
),
91+
{ pos: parent.start },
92+
)
93+
}
94+
if (
95+
(isObjectMethod && parent.kind !== 'init') ||
96+
(isClassMethod && parent.kind !== 'method')
97+
) {
98+
throw Object.assign(
99+
new Error(
100+
`It is not allowed to define inline ${JSON.stringify(match[0])} annotated getters or setters.`,
101+
),
102+
{ pos: parent.start },
103+
)
104+
}
105+
68106
const declName = node.type === 'FunctionDeclaration' && node.id.name
107+
const expressionName =
108+
node.type === 'FunctionExpression' ? node.id?.name : undefined
109+
const methodName =
110+
(isObjectMethod || isClassMethod) &&
111+
(parent.key.type === 'Identifier' || parent.key.type === 'Literal')
112+
? String(
113+
parent.key.type === 'Identifier'
114+
? parent.key.name
115+
: parent.key.value,
116+
)
117+
: undefined
69118
const originalName =
70119
declName ||
120+
methodName ||
71121
(parent?.type === 'VariableDeclarator' &&
72122
parent.id.type === 'Identifier' &&
73123
parent.id.name) ||
124+
expressionName ||
74125
'anonymous_server_function'
75126

76127
const bindVars = getBindVars(node, scopeTree)
@@ -92,8 +143,15 @@ export function transformHoistInlineDirective(
92143
}
93144

94145
// append a new `FunctionDeclaration` at the end
146+
let nameKey = String(names.length)
147+
if (options.stableName) {
148+
const signature = `${originalName}:${input.slice(node.start, node.end)}`
149+
const signatureCount = signatureCounts.get(signature) ?? 0
150+
signatureCounts.set(signature, signatureCount + 1)
151+
nameKey = `${hashString(signature)}_${signatureCount}`
152+
}
95153
const newName =
96-
`$$hoist_${names.length}` + (originalName ? `_${originalName}` : '')
154+
`$$hoist_${nameKey}` + (originalName ? `_${originalName}` : '')
97155
names.push(newName)
98156
output.update(
99157
node.start,
@@ -113,14 +171,27 @@ export function transformHoistInlineDirective(
113171
// replace original declartion with action register + bind
114172
let newCode = `/* #__PURE__ */ ${runtime(newName, newName, {
115173
directiveMatch: match,
174+
hasBoundArgs: bindVars.length > 0,
116175
})}`
117176
if (bindVars.length > 0) {
118177
const bindArgs = options.encode
119178
? options.encode('[' + bindVars.map((b) => b.expr).join(', ') + ']')
120179
: bindVars.map((b) => b.expr).join(', ')
121180
newCode = `${newCode}.bind(null, ${bindArgs})`
122181
}
123-
if (declName) {
182+
if (isObjectMethod) {
183+
output.update(
184+
parent.start,
185+
node.start,
186+
`${input.slice(parent.key.start, parent.key.end)}: `,
187+
)
188+
} else if (isClassMethod) {
189+
const key = parent.computed
190+
? `[${input.slice(parent.key.start, parent.key.end)}]`
191+
: input.slice(parent.key.start, parent.key.end)
192+
output.update(parent.start, node.start, `static ${key} = `)
193+
newCode += ';'
194+
} else if (declName) {
124195
newCode = `const ${declName} = ${newCode};`
125196
if (parent?.type === 'ExportDefaultDeclaration') {
126197
output.remove(parent.start, node.start)
@@ -132,10 +203,7 @@ export function transformHoistInlineDirective(
132203
},
133204
})
134205

135-
return {
136-
output,
137-
names,
138-
}
206+
return { output, names }
139207
}
140208

141209
const exactRegex = (s: string): RegExp =>
@@ -151,7 +219,9 @@ function matchDirective(
151219
stmt.expression.type === 'Literal' &&
152220
typeof stmt.expression.value === 'string'
153221
) {
222+
directive.lastIndex = 0
154223
const match = stmt.expression.value.match(directive)
224+
directive.lastIndex = 0
155225
if (match) {
156226
return { match, node: stmt.expression }
157227
}
Lines changed: 50 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
1-
import type { Program } from 'estree'
1+
import type { Literal, Program } from 'estree'
22
import type MagicString from 'magic-string'
33
import { transformHoistInlineDirective } from './hoist'
4-
import { hasDirective } from './utils'
5-
import { transformWrapExport } from './wrap-export'
4+
import {
5+
transformWrapExport,
6+
type TransformWrapExportOptions,
7+
} from './wrap-export'
68

79
// TODO
810
// source map for `options.runtime` (registerServerReference) call
@@ -12,9 +14,19 @@ export function transformServerActionServer(
1214
ast: Program,
1315
options: {
1416
runtime: (value: string, name: string) => string
17+
directive?: string | RegExp
18+
moduleDirective?: Literal
19+
moduleRuntime?: TransformWrapExportOptions['runtime']
20+
inlineRuntime?: Parameters<
21+
typeof transformHoistInlineDirective
22+
>[2]['runtime']
23+
filter?: TransformWrapExportOptions['filter']
1524
rejectNonAsyncFunction?: boolean
25+
rejectNonAsyncModule?: boolean
1626
encode?: (value: string) => string
1727
decode?: (value: string) => string
28+
stableName?: boolean
29+
preserveModuleDirective?: boolean
1830
},
1931
):
2032
| {
@@ -25,12 +37,42 @@ export function transformServerActionServer(
2537
output: MagicString
2638
names: string[]
2739
} {
28-
// TODO: unify (generalize transformHoistInlineDirective to support top-level directive cases)
29-
if (hasDirective(ast.body, 'use server')) {
30-
return transformWrapExport(input, ast, options)
40+
const useServerStatement = ast.body.find(
41+
(statement) =>
42+
statement.type === 'ExpressionStatement' &&
43+
statement.expression.type === 'Literal' &&
44+
statement.expression.value === 'use server',
45+
)
46+
const moduleDirective =
47+
options.moduleDirective ??
48+
(useServerStatement?.type === 'ExpressionStatement' &&
49+
useServerStatement.expression.type === 'Literal'
50+
? useServerStatement.expression
51+
: undefined)
52+
53+
if (moduleDirective?.type === 'Literal') {
54+
const result = transformWrapExport(input, ast, {
55+
runtime: options.moduleRuntime ?? options.runtime,
56+
filter: options.filter,
57+
rejectNonAsyncFunction:
58+
options.rejectNonAsyncModule ?? options.rejectNonAsyncFunction,
59+
})
60+
if (!options.preserveModuleDirective && options.moduleDirective) {
61+
result.output.overwrite(
62+
moduleDirective.start,
63+
moduleDirective.end,
64+
`/* ${JSON.stringify(moduleDirective.value)} */`,
65+
)
66+
}
67+
return result
3168
}
69+
3270
return transformHoistInlineDirective(input, ast, {
33-
...options,
34-
directive: 'use server',
71+
runtime: options.inlineRuntime ?? options.runtime,
72+
directive: options.directive ?? 'use server',
73+
rejectNonAsyncFunction: options.rejectNonAsyncFunction,
74+
encode: options.encode,
75+
decode: options.decode,
76+
stableName: options.stableName,
3577
})
3678
}

packages/plugin-rsc/src/transforms/wrap-export.test.ts

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -314,4 +314,32 @@ export default Page;
314314
"
315315
`)
316316
})
317+
318+
test('filtered exports are not validated or reported', async () => {
319+
const input = `
320+
export const revalidate = 1;
321+
export default async function Page() {}
322+
`
323+
const ast = await parseAstAsync(input)
324+
const result = transformWrapExport(input, ast, {
325+
runtime: (value, name) => `$$wrap(${value}, ${JSON.stringify(name)})`,
326+
rejectNonAsyncFunction: true,
327+
filter: (name) => name !== 'revalidate',
328+
})
329+
expect(result.exportNames).toEqual(['default'])
330+
expect(result.output.toString()).toContain('export { revalidate };')
331+
})
332+
333+
test('unknown identifier exports remain eligible for wrapping', async () => {
334+
const input = `
335+
const cached = async () => 1;
336+
export default cached;
337+
`
338+
const ast = await parseAstAsync(input)
339+
const result = transformWrapExport(input, ast, {
340+
runtime: (value, name) => `$$wrap(${value}, ${JSON.stringify(name)})`,
341+
filter: (_name, meta) => meta.isFunction !== false,
342+
})
343+
expect(result.exportNames).toEqual(['default'])
344+
})
317345
})

0 commit comments

Comments
 (0)