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