1
+/**
2
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
3
+ *
4
+ * This source code is licensed under the MIT license found in the
5
+ * LICENSE file in the root directory of this source tree.
6
+ */
7
+
8
+import {CompilerError} from '..';
9
+import {
10
+ CallExpression,
11
+ getHookKind,
12
+ HIRFunction,
13
+ IdentifierId,
14
+ InstructionValue,
15
+ isArrayType,
16
+ isPlainObjectType,
17
+ isPrimitiveType,
18
+ isSetStateType,
19
+ isStartTransitionType,
20
+ LoadLocal,
21
+ StoreLocal,
22
+} from '../HIR';
23
+import {
24
+ eachInstructionValueOperand,
25
+ eachTerminalOperand,
26
+} from '../HIR/visitors';
27
+import {retainWhere} from '../Utils/utils';
28
+
29
+/**
30
+ * Optimizes the code for running specifically in an SSR environment. This optimization
31
+ * asssumes that setState will not be called during render during initial mount, which
32
+ * allows inlining useState/useReducer.
33
+ *
34
+ * Optimizations:
35
+ * - Inline useState/useReducer
36
+ * - Remove effects
37
+ * - Remove refs where known to be unused during render (eg directly passed to a dom node)
38
+ * - Remove event handlers
39
+ *
40
+ * Note that an earlier pass already inlines useMemo/useCallback
41
+ */
42
+export function optimizeForSSR(fn: HIRFunction): void {
43
+ const inlinedState = new Map<IdentifierId, InstructionValue>();
44
+ /**
45
+ * First pass identifies useState/useReducer which can be safely inlined. Any use
46
+ * of the hook return other than destructuring (with a specific pattern) prevents
47
+ * inlining.
48
+ *
49
+ * Supported cases:
50
+ * - `const [state, ] = useState( <primitive-array-or-object> )`
51
+ * - `const [state, ] = useReducer(..., <value>)`
52
+ * - `const [state, ] = useReducer[..., <value>, <init>]`
53
+ */
54
+ for (const block of fn.body.blocks.values()) {
55
+ for (const instr of block.instructions) {
56
+ const {value} = instr;
57
+ switch (value.kind) {
58
+ case 'Destructure': {
59
+ if (
60
+ inlinedState.has(value.value.identifier.id) &&
61
+ value.lvalue.pattern.kind === 'ArrayPattern' &&
62
+ value.lvalue.pattern.items.length >= 1 &&
63
+ value.lvalue.pattern.items[0].kind === 'Identifier'
64
+ ) {
65
+ // Allow destructuring of inlined states
66
+ continue;
67
+ }
68
+ break;
69
+ }
70
+ case 'MethodCall':
71
+ case 'CallExpression': {
72
+ const calleee =
73
+ value.kind === 'CallExpression' ? value.callee : value.property;
74
+ const hookKind = getHookKind(fn.env, calleee.identifier);
75
+ switch (hookKind) {
76
+ case 'useReducer': {
77
+ if (
78
+ value.args.length === 2 &&
79
+ value.args[1].kind === 'Identifier'
80
+ ) {
81
+ const arg = value.args[1];
82
+ const replace: LoadLocal = {
83
+ kind: 'LoadLocal',
84
+ place: arg,
85
+ loc: arg.loc,
86
+ };
87
+ inlinedState.set(instr.lvalue.identifier.id, replace);
88
+ } else if (
89
+ value.args.length === 3 &&
90
+ value.args[1].kind === 'Identifier' &&
91
+ value.args[2].kind === 'Identifier'
92
+ ) {
93
+ const arg = value.args[1];
94
+ const initializer = value.args[2];
95
+ const replace: CallExpression = {
96
+ kind: 'CallExpression',
97
+ callee: initializer,
98
+ args: [arg],
99
+ loc: value.loc,
100
+ };
101
+ inlinedState.set(instr.lvalue.identifier.id, replace);
102
+ }
103
+ break;
104
+ }
105
+ case 'useState': {
106
+ if (
107
+ value.args.length === 1 &&
108
+ value.args[0].kind === 'Identifier'
109
+ ) {
110
+ const arg = value.args[0];
111
+ if (
112
+ isPrimitiveType(arg.identifier) ||
113
+ isPlainObjectType(arg.identifier) ||
114
+ isArrayType(arg.identifier)
115
+ ) {
116
+ const replace: LoadLocal = {
117
+ kind: 'LoadLocal',
118
+ place: arg,
119
+ loc: arg.loc,
120
+ };
121
+ inlinedState.set(instr.lvalue.identifier.id, replace);
122
+ }
123
+ }
124
+ break;
125
+ }
126
+ }
127
+ }
128
+ }
129
+ // Any use of useState/useReducer return besides destructuring prevents inlining
130
+ if (inlinedState.size !== 0) {
131
+ for (const operand of eachInstructionValueOperand(value)) {
132
+ inlinedState.delete(operand.identifier.id);
133
+ }
134
+ }
135
+ }
136
+ if (inlinedState.size !== 0) {
137
+ for (const operand of eachTerminalOperand(block.terminal)) {
138
+ inlinedState.delete(operand.identifier.id);
139
+ }
140
+ }
141
+ }
142
+ for (const block of fn.body.blocks.values()) {
143
+ for (const instr of block.instructions) {
144
+ const {value} = instr;
145
+ switch (value.kind) {
146
+ case 'FunctionExpression': {
147
+ if (hasKnownNonRenderCall(value.loweredFunc.func)) {
148
+ instr.value = {
149
+ kind: 'Primitive',
150
+ value: undefined,
151
+ loc: value.loc,
152
+ };
153
+ }
154
+ break;
155
+ }
156
+ case 'JsxExpression': {
157
+ if (
158
+ value.tag.kind === 'BuiltinTag' &&
159
+ value.tag.name.indexOf('-') === -1
160
+ ) {
161
+ const tag = value.tag.name;
162
+ retainWhere(value.props, prop => {
163
+ return (
164
+ prop.kind === 'JsxSpreadAttribute' ||
165
+ (!isKnownEventHandler(tag, prop.name) && prop.name !== 'ref')
166
+ );
167
+ });
168
+ }
169
+ break;
170
+ }
171
+ case 'Destructure': {
172
+ if (inlinedState.has(value.value.identifier.id)) {
173
+ // Canonical check is part of determining if state can inline, this is for TS
174
+ CompilerError.invariant(
175
+ value.lvalue.pattern.kind === 'ArrayPattern' &&
176
+ value.lvalue.pattern.items.length >= 1 &&
177
+ value.lvalue.pattern.items[0].kind === 'Identifier',
178
+ {
179
+ reason:
180
+ 'Expected a valid destructuring pattern for inlined state',
181
+ description: null,
182
+ details: [
183
+ {
184
+ kind: 'error',
185
+ message: 'Expected a valid destructuring pattern',
186
+ loc: value.loc,
187
+ },
188
+ ],
189
+ },
190
+ );
191
+ const store: StoreLocal = {
192
+ kind: 'StoreLocal',
193
+ loc: value.loc,
194
+ type: null,
195
+ lvalue: {
196
+ kind: value.lvalue.kind,
197
+ place: value.lvalue.pattern.items[0],
198
+ },
199
+ value: value.value,
200
+ };
201
+ instr.value = store;
202
+ }
203
+ break;
204
+ }
205
+ case 'MethodCall':
206
+ case 'CallExpression': {
207
+ const calleee =
208
+ value.kind === 'CallExpression' ? value.callee : value.property;
209
+ const hookKind = getHookKind(fn.env, calleee.identifier);
210
+ switch (hookKind) {
211
+ case 'useEffectEvent': {
212
+ if (
213
+ value.args.length === 1 &&
214
+ value.args[0].kind === 'Identifier'
215
+ ) {
216
+ const load: LoadLocal = {
217
+ kind: 'LoadLocal',
218
+ place: value.args[0],
219
+ loc: value.loc,
220
+ };
221
+ instr.value = load;
222
+ }
223
+ break;
224
+ }
225
+ case 'useEffect':
226
+ case 'useLayoutEffect':
227
+ case 'useInsertionEffect': {
228
+ // Drop effects
229
+ instr.value = {
230
+ kind: 'Primitive',
231
+ value: undefined,
232
+ loc: value.loc,
233
+ };
234
+ break;
235
+ }
236
+ case 'useReducer':
237
+ case 'useState': {
238
+ const replace = inlinedState.get(instr.lvalue.identifier.id);
239
+ if (replace != null) {
240
+ instr.value = replace;
241
+ }
242
+ break;
243
+ }
244
+ }
245
+ }
246
+ }
247
+ }
248
+ }
249
+}
250
+
251
+function hasKnownNonRenderCall(fn: HIRFunction): boolean {
252
+ for (const block of fn.body.blocks.values()) {
253
+ for (const instr of block.instructions) {
254
+ if (
255
+ instr.value.kind === 'CallExpression' &&
256
+ (isSetStateType(instr.value.callee.identifier) ||
257
+ isStartTransitionType(instr.value.callee.identifier))
258
+ ) {
259
+ return true;
260
+ }
261
+ }
262
+ }
263
+ return false;
264
+}
265
+
266
+const EVENT_HANDLER_PATTERN = /^on[A-Z]/;
267
+function isKnownEventHandler(_tag: string, prop: string): boolean {
268
+ return EVENT_HANDLER_PATTERN.test(prop);
269
+}