react_compiler_optimization/optimize_for_ssr.rs
1// Copyright (c) Meta Platforms, Inc. and affiliates.
2//
3// This source code is licensed under the MIT license found in the
4// LICENSE file in the root directory of this source tree.
5
6//! Optimizes the code for running in an SSR environment.
7//!
8//! Assumes that setState will not be called during render during initial mount,
9//! which allows inlining useState/useReducer.
10//!
11//! Optimizations:
12//! - Inline useState/useReducer
13//! - Remove effects (useEffect, useLayoutEffect, useInsertionEffect)
14//! - Remove event handlers (functions that call setState or startTransition)
15//! - Remove known event handler props and ref props from builtin JSX tags
16//! - Inline useEffectEvent to its argument
17//!
18//! Ported from TypeScript `src/Optimization/OptimizeForSSR.ts`.
19
20use std::collections::HashMap;
21
22use react_compiler_hir::environment::Environment;
23use react_compiler_hir::object_shape::HookKind;
24use react_compiler_hir::visitors::{each_instruction_value_operand, each_terminal_operand};
25use react_compiler_hir::{
26 ArrayPatternElement, HirFunction, IdentifierId, InstructionValue, PlaceOrSpread,
27 PrimitiveValue, is_set_state_type, is_start_transition_type,
28};
29
30/// Optimizes a function for SSR by inlining state hooks, removing effects,
31/// removing event handlers, and stripping known event handler / ref JSX props.
32///
33/// Corresponds to TS `optimizeForSSR(fn: HIRFunction): void`.
34pub fn optimize_for_ssr(func: &mut HirFunction, env: &Environment) {
35 // Phase 1: Identify useState/useReducer calls that can be safely inlined.
36 //
37 // For useState(initialValue) where initialValue is primitive/object/array,
38 // store a LoadLocal of the initial value.
39 //
40 // For useReducer(reducer, initialArg) store a LoadLocal of initialArg.
41 // For useReducer(reducer, initialArg, init) store a CallExpression of init(initialArg).
42 //
43 // Any use of the hook return other than the expected destructuring pattern
44 // prevents inlining (we delete from inlined_state if we see the identifier used
45 // as an operand elsewhere).
46 let mut inlined_state: HashMap<IdentifierId, InlinedStateReplacement> = HashMap::new();
47
48 for (_block_id, block) in &func.body.blocks {
49 for &instr_id in &block.instructions {
50 let instr = &func.instructions[instr_id.0 as usize];
51 match &instr.value {
52 InstructionValue::Destructure { value, lvalue, .. } => {
53 if inlined_state.contains_key(&env.identifiers[value.identifier.0 as usize].id) {
54 if let react_compiler_hir::Pattern::Array(arr) = &lvalue.pattern {
55 if !arr.items.is_empty() {
56 if let ArrayPatternElement::Place(_) = &arr.items[0] {
57 // Allow destructuring of inlined states
58 continue;
59 }
60 }
61 }
62 }
63 }
64 InstructionValue::MethodCall {
65 property, args, ..
66 }
67 | InstructionValue::CallExpression {
68 callee: property,
69 args,
70 ..
71 } => {
72 // Determine callee based on instruction kind
73 let callee_id = property.identifier;
74 let hook_kind = get_hook_kind(env, callee_id);
75 match hook_kind {
76 Some(HookKind::UseReducer) => {
77 if args.len() == 2 {
78 if let (
79 PlaceOrSpread::Place(_),
80 PlaceOrSpread::Place(arg),
81 ) = (&args[0], &args[1])
82 {
83 let lvalue_id = env.identifiers
84 [instr.lvalue.identifier.0 as usize]
85 .id;
86 inlined_state.insert(
87 lvalue_id,
88 InlinedStateReplacement::LoadLocal {
89 place: arg.clone(),
90 loc: arg.loc,
91 },
92 );
93 }
94 } else if args.len() == 3 {
95 if let (
96 PlaceOrSpread::Place(_),
97 PlaceOrSpread::Place(arg),
98 PlaceOrSpread::Place(initializer),
99 ) = (&args[0], &args[1], &args[2])
100 {
101 let lvalue_id = env.identifiers
102 [instr.lvalue.identifier.0 as usize]
103 .id;
104 let call_loc = instr.value.loc().copied();
105 inlined_state.insert(
106 lvalue_id,
107 InlinedStateReplacement::CallExpression {
108 callee: initializer.clone(),
109 arg: arg.clone(),
110 loc: call_loc,
111 },
112 );
113 }
114 }
115 }
116 Some(HookKind::UseState) => {
117 if args.len() == 1 {
118 if let PlaceOrSpread::Place(arg) = &args[0] {
119 let arg_type = &env.types
120 [env.identifiers[arg.identifier.0 as usize].type_.0
121 as usize];
122 if react_compiler_hir::is_primitive_type(arg_type)
123 || react_compiler_hir::is_plain_object_type(arg_type)
124 || react_compiler_hir::is_array_type(arg_type)
125 {
126 let lvalue_id = env.identifiers
127 [instr.lvalue.identifier.0 as usize]
128 .id;
129 inlined_state.insert(
130 lvalue_id,
131 InlinedStateReplacement::LoadLocal {
132 place: arg.clone(),
133 loc: arg.loc,
134 },
135 );
136 }
137 }
138 }
139 }
140 _ => {}
141 }
142 }
143 _ => {}
144 }
145
146 // Any use of useState/useReducer return besides destructuring prevents inlining
147 if !inlined_state.is_empty() {
148 let operands =
149 each_instruction_value_operand(&instr.value, env);
150 for operand in &operands {
151 let id = env.identifiers[operand.identifier.0 as usize].id;
152 inlined_state.remove(&id);
153 }
154 }
155 }
156 if !inlined_state.is_empty() {
157 let operands = each_terminal_operand(&block.terminal);
158 for operand in &operands {
159 let id = env.identifiers[operand.identifier.0 as usize].id;
160 inlined_state.remove(&id);
161 }
162 }
163 }
164
165 // Phase 2: Apply transformations
166 //
167 // - Replace FunctionExpression with Primitive(undefined) if it calls setState/startTransition
168 // - Remove known event handler props and ref props from builtin JSX tags
169 // - Replace Destructure of inlined state with StoreLocal
170 // - Replace useEffectEvent(fn) with LoadLocal(fn)
171 // - Replace useEffect/useLayoutEffect/useInsertionEffect with Primitive(undefined)
172 // - Replace useState/useReducer with their inlined replacement
173 for (_block_id, block) in &mut func.body.blocks {
174 for &instr_id in &block.instructions {
175 let instr = &mut func.instructions[instr_id.0 as usize];
176 match &instr.value {
177 InstructionValue::FunctionExpression {
178 lowered_func, loc, ..
179 } => {
180 let inner_func = &env.functions[lowered_func.func.0 as usize];
181 if has_known_non_render_call(inner_func, env) {
182 let loc = *loc;
183 instr.value = InstructionValue::Primitive {
184 value: PrimitiveValue::Undefined,
185 loc,
186 };
187 }
188 }
189 InstructionValue::JsxExpression { tag, .. } => {
190 if let react_compiler_hir::JsxTag::Builtin(builtin) = tag {
191 // Only optimize non-custom-element builtin tags
192 if !builtin.name.contains('-') {
193 let tag_name = builtin.name.clone();
194 // Retain only props that are not known event handlers and not "ref"
195 if let InstructionValue::JsxExpression { props, .. } =
196 &mut instr.value
197 {
198 props.retain(|prop| match prop {
199 react_compiler_hir::JsxAttribute::SpreadAttribute { .. } => {
200 true
201 }
202 react_compiler_hir::JsxAttribute::Attribute {
203 name, ..
204 } => {
205 !is_known_event_handler(&tag_name, name)
206 && name != "ref"
207 }
208 });
209 }
210 }
211 }
212 }
213 InstructionValue::Destructure { value, lvalue, loc } => {
214 let value_id = env.identifiers[value.identifier.0 as usize].id;
215 if inlined_state.contains_key(&value_id) {
216 // Invariant: destructuring pattern must be ArrayPattern with at least one Identifier item
217 if let react_compiler_hir::Pattern::Array(arr) = &lvalue.pattern {
218 if !arr.items.is_empty() {
219 if let ArrayPatternElement::Place(first_place) = &arr.items[0] {
220 let loc = *loc;
221 let kind = lvalue.kind;
222 let store = InstructionValue::StoreLocal {
223 lvalue: react_compiler_hir::LValue {
224 place: first_place.clone(),
225 kind,
226 },
227 value: value.clone(),
228 type_annotation: None,
229 loc,
230 };
231 instr.value = store;
232 }
233 }
234 }
235 }
236 }
237 InstructionValue::MethodCall {
238 property, args, loc, ..
239 }
240 | InstructionValue::CallExpression {
241 callee: property,
242 args,
243 loc,
244 ..
245 } => {
246 let callee_id = property.identifier;
247 let hook_kind = get_hook_kind(env, callee_id);
248 match hook_kind {
249 Some(HookKind::UseEffectEvent) => {
250 if args.len() == 1 {
251 if let PlaceOrSpread::Place(arg) = &args[0] {
252 let loc = *loc;
253 instr.value = InstructionValue::LoadLocal {
254 place: arg.clone(),
255 loc,
256 };
257 }
258 }
259 }
260 Some(
261 HookKind::UseEffect
262 | HookKind::UseLayoutEffect
263 | HookKind::UseInsertionEffect,
264 ) => {
265 let loc = *loc;
266 instr.value = InstructionValue::Primitive {
267 value: PrimitiveValue::Undefined,
268 loc,
269 };
270 }
271 Some(HookKind::UseReducer | HookKind::UseState) => {
272 let lvalue_id =
273 env.identifiers[instr.lvalue.identifier.0 as usize].id;
274 if let Some(replacement) = inlined_state.get(&lvalue_id) {
275 instr.value = match replacement {
276 InlinedStateReplacement::LoadLocal { place, loc } => {
277 InstructionValue::LoadLocal {
278 place: place.clone(),
279 loc: *loc,
280 }
281 }
282 InlinedStateReplacement::CallExpression {
283 callee,
284 arg,
285 loc,
286 } => InstructionValue::CallExpression {
287 callee: callee.clone(),
288 args: vec![PlaceOrSpread::Place(arg.clone())],
289 loc: *loc,
290 },
291 };
292 }
293 }
294 _ => {}
295 }
296 }
297 _ => {}
298 }
299 }
300 }
301}
302
303/// Replacement values for inlined useState/useReducer calls.
304#[derive(Debug, Clone)]
305enum InlinedStateReplacement {
306 /// Replace with `LoadLocal { place }` — used for useState and useReducer(reducer, initialArg)
307 LoadLocal {
308 place: react_compiler_hir::Place,
309 loc: Option<react_compiler_hir::SourceLocation>,
310 },
311 /// Replace with `CallExpression { callee, args: [arg] }` — used for useReducer(reducer, initialArg, init)
312 CallExpression {
313 callee: react_compiler_hir::Place,
314 arg: react_compiler_hir::Place,
315 loc: Option<react_compiler_hir::SourceLocation>,
316 },
317}
318
319/// Returns true if the function body contains a call to setState or startTransition.
320/// This identifies functions that are event handlers and can be replaced with undefined
321/// during SSR.
322///
323/// Corresponds to TS `hasKnownNonRenderCall(fn: HIRFunction): boolean`.
324fn has_known_non_render_call(func: &HirFunction, env: &Environment) -> bool {
325 for (_block_id, block) in &func.body.blocks {
326 for &instr_id in &block.instructions {
327 let instr = &func.instructions[instr_id.0 as usize];
328 if let InstructionValue::CallExpression { callee, .. } = &instr.value {
329 let callee_type =
330 &env.types[env.identifiers[callee.identifier.0 as usize].type_.0 as usize];
331 if is_set_state_type(callee_type) || is_start_transition_type(callee_type) {
332 return true;
333 }
334 }
335 }
336 }
337 false
338}
339
340/// Returns true if the prop name matches the known event handler pattern `on[A-Z]`.
341fn is_known_event_handler(_tag: &str, prop: &str) -> bool {
342 if prop.len() < 3 {
343 return false;
344 }
345 if !prop.starts_with("on") {
346 return false;
347 }
348 let third_char = prop.as_bytes()[2];
349 third_char.is_ascii_uppercase()
350}
351
352/// Get the hook kind for an identifier, if its type represents a hook.
353fn get_hook_kind(env: &Environment, identifier_id: IdentifierId) -> Option<HookKind> {
354 env.get_hook_kind_for_id(identifier_id)
355 .ok()
356 .flatten()
357 .cloned()
358}