| 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 | |
| 20 | use rustc_hash::FxHashMap; |
| 21 | |
| 22 | use react_compiler_hir::environment::Environment; |
| 23 | use react_compiler_hir::object_shape::HookKind; |
| 24 | use react_compiler_hir::visitors::{each_instruction_value_operand, each_terminal_operand}; |
| 25 | use 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`. |
| 34 | pub 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)] |
| 298 | enum 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`. |
| 317 | fn 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]`. |
| 334 | fn 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. |
| 346 | fn 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 | } |