| 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 | //! Inlines immediately invoked function expressions (IIFEs) to allow more |
| 7 | //! fine-grained memoization of the values they produce. |
| 8 | //! |
| 9 | //! Example: |
| 10 | //! ```text |
| 11 | //! const x = (() => { |
| 12 | //! const x = []; |
| 13 | //! x.push(foo()); |
| 14 | //! return x; |
| 15 | //! })(); |
| 16 | //! |
| 17 | //! => |
| 18 | //! |
| 19 | //! bb0: |
| 20 | //! // placeholder for the result, all return statements will assign here |
| 21 | //! let t0; |
| 22 | //! // Label allows using a goto (break) to exit out of the body |
| 23 | //! Label block=bb1 fallthrough=bb2 |
| 24 | //! bb1: |
| 25 | //! // code within the function expression |
| 26 | //! const x0 = []; |
| 27 | //! x0.push(foo()); |
| 28 | //! // return is replaced by assignment to the result variable... |
| 29 | //! t0 = x0; |
| 30 | //! // ...and a goto to the code after the function expression invocation |
| 31 | //! Goto bb2 |
| 32 | //! bb2: |
| 33 | //! // code after the IIFE call |
| 34 | //! const x = t0; |
| 35 | //! ``` |
| 36 | //! |
| 37 | //! If the inlined function has only one return, we avoid the labeled block |
| 38 | //! and fully inline the code. The original return is replaced with an assignment |
| 39 | //! to the IIFE's call expression lvalue. |
| 40 | //! |
| 41 | //! Analogous to TS `Inference/InlineImmediatelyInvokedFunctionExpressions.ts`. |
| 42 | |
| 43 | use indexmap::IndexSet; |
| 44 | use rustc_hash::{FxHashMap, FxHashSet}; |
| 45 | |
| 46 | use react_compiler_hir::environment::Environment; |
| 47 | use react_compiler_hir::visitors; |
| 48 | use react_compiler_hir::{ |
| 49 | BasicBlock, BlockId, BlockKind, EvaluationOrder, FunctionId, GENERATED_SOURCE, GotoVariant, |
| 50 | HirFunction, IdentifierId, IdentifierName, Instruction, InstructionId, InstructionKind, |
| 51 | InstructionValue, LValue, Place, Terminal, |
| 52 | }; |
| 53 | use react_compiler_lowering::{ |
| 54 | create_temporary_place, get_reverse_postordered_blocks, mark_instruction_ids, mark_predecessors, |
| 55 | }; |
| 56 | |
| 57 | use crate::merge_consecutive_blocks::merge_consecutive_blocks; |
| 58 | |
| 59 | /// Inline immediately invoked function expressions into the enclosing function's |
| 60 | /// control flow graph. |
| 61 | pub fn inline_immediately_invoked_function_expressions( |
| 62 | func: &mut HirFunction, |
| 63 | env: &mut Environment, |
| 64 | ) { |
| 65 | // Track all function expressions that are assigned to a temporary |
| 66 | let mut functions: FxHashMap<IdentifierId, FunctionId> = FxHashMap::default(); |
| 67 | // Functions that are inlined (by identifier id of the callee) |
| 68 | let mut inlined_functions: FxHashSet<IdentifierId> = FxHashSet::default(); |
| 69 | |
| 70 | // Iterate the *existing* blocks from the outer component to find IIFEs |
| 71 | // and inline them. During iteration we will modify `func` (by inlining the CFG |
| 72 | // of IIFEs) so we explicitly copy references to just the original |
| 73 | // function's block IDs first. As blocks are split to make room for IIFE calls, |
| 74 | // the split portions of the blocks will be added to this queue. |
| 75 | let mut queue: Vec<BlockId> = func.body.blocks.keys().copied().collect(); |
| 76 | let mut queue_idx = 0; |
| 77 | |
| 78 | 'queue: while queue_idx < queue.len() { |
| 79 | let block_id = queue[queue_idx]; |
| 80 | queue_idx += 1; |
| 81 | |
| 82 | let block = match func.body.blocks.get(&block_id) { |
| 83 | Some(b) => b, |
| 84 | None => continue, |
| 85 | }; |
| 86 | |
| 87 | // We can't handle labels inside expressions yet, so we don't inline IIFEs |
| 88 | // if they are in an expression block. |
| 89 | if !is_statement_block_kind(block.kind) { |
| 90 | continue; |
| 91 | } |
| 92 | |
| 93 | let num_instructions = block.instructions.len(); |
| 94 | for ii in 0..num_instructions { |
| 95 | let instr_id = func.body.blocks[&block_id].instructions[ii]; |
| 96 | let instr = &func.instructions[instr_id.0 as usize]; |
| 97 | |
| 98 | match &instr.value { |
| 99 | InstructionValue::FunctionExpression { lowered_func, .. } => { |
| 100 | let identifier_id = instr.lvalue.identifier; |
| 101 | if env.identifiers[identifier_id.0 as usize].name.is_none() { |
| 102 | functions.insert(identifier_id, lowered_func.func); |
| 103 | } |
| 104 | continue; |
| 105 | } |
| 106 | InstructionValue::CallExpression { callee, args, .. } => { |
| 107 | if !args.is_empty() { |
| 108 | // We don't support inlining when there are arguments |
| 109 | continue; |
| 110 | } |
| 111 | |
| 112 | let callee_id = callee.identifier; |
| 113 | let inner_func_id = match functions.get(&callee_id) { |
| 114 | Some(id) => *id, |
| 115 | None => continue, // Not invoking a local function expression |
| 116 | }; |
| 117 | |
| 118 | let inner_func = &env.functions[inner_func_id.0 as usize]; |
| 119 | if !inner_func.params.is_empty() || inner_func.is_async || inner_func.generator |
| 120 | { |
| 121 | // Can't inline functions with params, or async/generator functions |
| 122 | continue; |
| 123 | } |
| 124 | |
| 125 | // We know this function is used for an IIFE and can prune it later |
| 126 | inlined_functions.insert(callee_id); |
| 127 | |
| 128 | // Capture the lvalue from the call instruction |
| 129 | let call_lvalue = func.instructions[instr_id.0 as usize].lvalue.clone(); |
| 130 | let block_terminal_id = func.body.blocks[&block_id].terminal.evaluation_order(); |
| 131 | let block_terminal_loc = func.body.blocks[&block_id].terminal.loc().cloned(); |
| 132 | let block_kind = func.body.blocks[&block_id].kind; |
| 133 | |
| 134 | // Create a new block which will contain code following the IIFE call |
| 135 | let continuation_block_id = env.next_block_id(); |
| 136 | let continuation_instructions: Vec<InstructionId> = |
| 137 | func.body.blocks[&block_id].instructions[ii + 1..].to_vec(); |
| 138 | let continuation_terminal = func.body.blocks[&block_id].terminal.clone(); |
| 139 | let continuation_block = BasicBlock { |
| 140 | id: continuation_block_id, |
| 141 | instructions: continuation_instructions, |
| 142 | kind: block_kind, |
| 143 | phis: Vec::new(), |
| 144 | preds: IndexSet::default(), |
| 145 | terminal: continuation_terminal, |
| 146 | }; |
| 147 | func.body |
| 148 | .blocks |
| 149 | .insert(continuation_block_id, continuation_block); |
| 150 | |
| 151 | // Trim the original block to contain instructions up to (but not including) |
| 152 | // the IIFE |
| 153 | func.body |
| 154 | .blocks |
| 155 | .get_mut(&block_id) |
| 156 | .unwrap() |
| 157 | .instructions |
| 158 | .truncate(ii); |
| 159 | |
| 160 | let has_single_return = |
| 161 | has_single_exit_return_terminal(&env.functions[inner_func_id.0 as usize]); |
| 162 | let inner_entry = env.functions[inner_func_id.0 as usize].body.entry; |
| 163 | |
| 164 | if has_single_return { |
| 165 | // Single-return path: simple goto replacement |
| 166 | func.body.blocks.get_mut(&block_id).unwrap().terminal = Terminal::Goto { |
| 167 | block: inner_entry, |
| 168 | id: block_terminal_id, |
| 169 | loc: block_terminal_loc, |
| 170 | variant: GotoVariant::Break, |
| 171 | }; |
| 172 | |
| 173 | // Take blocks and instructions from inner function |
| 174 | let inner_func = &mut env.functions[inner_func_id.0 as usize]; |
| 175 | let inner_blocks: Vec<(BlockId, BasicBlock)> = |
| 176 | inner_func.body.blocks.drain(..).collect(); |
| 177 | let inner_instructions: Vec<Instruction> = |
| 178 | inner_func.instructions.drain(..).collect(); |
| 179 | |
| 180 | // Append inner instructions first, then remap block instruction IDs |
| 181 | let instr_offset = func.instructions.len() as u32; |
| 182 | func.instructions.extend(inner_instructions); |
| 183 | |
| 184 | for (_, mut inner_block) in inner_blocks { |
| 185 | // Remap instruction IDs in the block |
| 186 | for iid in &mut inner_block.instructions { |
| 187 | *iid = InstructionId(iid.0 + instr_offset); |
| 188 | } |
| 189 | inner_block.preds.clear(); |
| 190 | |
| 191 | if let Terminal::Return { |
| 192 | value, |
| 193 | id: ret_id, |
| 194 | loc: ret_loc, |
| 195 | .. |
| 196 | } = &inner_block.terminal |
| 197 | { |
| 198 | // Replace return with LoadLocal + goto |
| 199 | let load_instr = Instruction { |
| 200 | id: EvaluationOrder(0), |
| 201 | loc: ret_loc.clone(), |
| 202 | lvalue: call_lvalue.clone(), |
| 203 | value: InstructionValue::LoadLocal { |
| 204 | place: value.clone(), |
| 205 | loc: ret_loc.clone(), |
| 206 | }, |
| 207 | effects: None, |
| 208 | }; |
| 209 | let load_instr_id = InstructionId(func.instructions.len() as u32); |
| 210 | func.instructions.push(load_instr); |
| 211 | inner_block.instructions.push(load_instr_id); |
| 212 | |
| 213 | let ret_id = *ret_id; |
| 214 | let ret_loc = ret_loc.clone(); |
| 215 | inner_block.terminal = Terminal::Goto { |
| 216 | block: continuation_block_id, |
| 217 | id: ret_id, |
| 218 | loc: ret_loc, |
| 219 | variant: GotoVariant::Break, |
| 220 | }; |
| 221 | } |
| 222 | |
| 223 | func.body.blocks.insert(inner_block.id, inner_block); |
| 224 | } |
| 225 | } else { |
| 226 | // Multi-return path: uses LabelTerminal |
| 227 | let result = call_lvalue.clone(); |
| 228 | |
| 229 | // Set block terminal to Label |
| 230 | func.body.blocks.get_mut(&block_id).unwrap().terminal = Terminal::Label { |
| 231 | block: inner_entry, |
| 232 | id: EvaluationOrder(0), |
| 233 | fallthrough: continuation_block_id, |
| 234 | loc: block_terminal_loc, |
| 235 | }; |
| 236 | |
| 237 | // Declare the IIFE temporary |
| 238 | declare_temporary(env, func, block_id, &result); |
| 239 | |
| 240 | // Promote the temporary with a name as we require this to persist |
| 241 | let identifier_id = result.identifier; |
| 242 | if env.identifiers[identifier_id.0 as usize].name.is_none() { |
| 243 | promote_temporary(env, identifier_id); |
| 244 | } |
| 245 | |
| 246 | // Take blocks and instructions from inner function |
| 247 | let inner_func = &mut env.functions[inner_func_id.0 as usize]; |
| 248 | let inner_blocks: Vec<(BlockId, BasicBlock)> = |
| 249 | inner_func.body.blocks.drain(..).collect(); |
| 250 | let inner_instructions: Vec<Instruction> = |
| 251 | inner_func.instructions.drain(..).collect(); |
| 252 | |
| 253 | // Append inner instructions first, then remap block instruction IDs |
| 254 | let instr_offset = func.instructions.len() as u32; |
| 255 | func.instructions.extend(inner_instructions); |
| 256 | |
| 257 | for (_, mut inner_block) in inner_blocks { |
| 258 | for iid in &mut inner_block.instructions { |
| 259 | *iid = InstructionId(iid.0 + instr_offset); |
| 260 | } |
| 261 | inner_block.preds.clear(); |
| 262 | |
| 263 | // Rewrite return terminals to StoreLocal + goto |
| 264 | if matches!(inner_block.terminal, Terminal::Return { .. }) { |
| 265 | rewrite_block( |
| 266 | env, |
| 267 | &mut func.instructions, |
| 268 | &mut inner_block, |
| 269 | continuation_block_id, |
| 270 | &result, |
| 271 | ); |
| 272 | } |
| 273 | |
| 274 | func.body.blocks.insert(inner_block.id, inner_block); |
| 275 | } |
| 276 | } |
| 277 | |
| 278 | // Ensure we visit the continuation block, since there may have been |
| 279 | // sequential IIFEs that need to be visited. |
| 280 | queue.push(continuation_block_id); |
| 281 | continue 'queue; |
| 282 | } |
| 283 | _ => { |
| 284 | // Any other use of a function expression means it isn't an IIFE |
| 285 | for id in visitors::each_instruction_value_operand_ids(&instr.value, env) { |
| 286 | functions.remove(&id); |
| 287 | } |
| 288 | } |
| 289 | } |
| 290 | } |
| 291 | } |
| 292 | |
| 293 | if !inlined_functions.is_empty() { |
| 294 | // Remove instructions that define lambdas which we inlined |
| 295 | for block in func.body.blocks.values_mut() { |
| 296 | block.instructions.retain(|instr_id| { |
| 297 | let instr = &func.instructions[instr_id.0 as usize]; |
| 298 | !inlined_functions.contains(&instr.lvalue.identifier) |
| 299 | }); |
| 300 | } |
| 301 | |
| 302 | // If terminals have changed then blocks may have become newly unreachable. |
| 303 | // Re-run minification of the graph (incl reordering instruction ids). |
| 304 | func.body.blocks = get_reverse_postordered_blocks(&func.body, &func.instructions); |
| 305 | mark_instruction_ids(&mut func.body, &mut func.instructions); |
| 306 | mark_predecessors(&mut func.body); |
| 307 | merge_consecutive_blocks(func, &mut env.functions); |
| 308 | } |
| 309 | } |
| 310 | |
| 311 | /// Returns true for "block" and "catch" block kinds which correspond to statements |
| 312 | /// in the source. |
| 313 | fn is_statement_block_kind(kind: BlockKind) -> bool { |
| 314 | matches!(kind, BlockKind::Block | BlockKind::Catch) |
| 315 | } |
| 316 | |
| 317 | /// Returns true if the function has a single exit terminal (throw/return) which is a return. |
| 318 | fn has_single_exit_return_terminal(func: &HirFunction) -> bool { |
| 319 | let mut has_return = false; |
| 320 | let mut exit_count = 0; |
| 321 | for block in func.body.blocks.values() { |
| 322 | match &block.terminal { |
| 323 | Terminal::Return { .. } => { |
| 324 | has_return = true; |
| 325 | exit_count += 1; |
| 326 | } |
| 327 | Terminal::Throw { .. } => { |
| 328 | exit_count += 1; |
| 329 | } |
| 330 | _ => {} |
| 331 | } |
| 332 | } |
| 333 | exit_count == 1 && has_return |
| 334 | } |
| 335 | |
| 336 | /// Rewrites the block so that all `return` terminals are replaced: |
| 337 | /// * Add a StoreLocal <return_value> = <terminal.value> |
| 338 | /// * Replace the terminal with a Goto to <return_target> |
| 339 | fn rewrite_block( |
| 340 | env: &mut Environment, |
| 341 | instructions: &mut Vec<Instruction>, |
| 342 | block: &mut BasicBlock, |
| 343 | return_target: BlockId, |
| 344 | return_value: &Place, |
| 345 | ) { |
| 346 | if let Terminal::Return { |
| 347 | value, |
| 348 | loc: ret_loc, |
| 349 | .. |
| 350 | } = &block.terminal |
| 351 | { |
| 352 | let store_lvalue = create_temporary_place(env, ret_loc.clone()); |
| 353 | let store_instr = Instruction { |
| 354 | id: EvaluationOrder(0), |
| 355 | loc: ret_loc.clone(), |
| 356 | lvalue: store_lvalue, |
| 357 | value: InstructionValue::StoreLocal { |
| 358 | lvalue: LValue { |
| 359 | kind: InstructionKind::Reassign, |
| 360 | place: return_value.clone(), |
| 361 | }, |
| 362 | value: value.clone(), |
| 363 | type_annotation: None, |
| 364 | loc: ret_loc.clone(), |
| 365 | }, |
| 366 | effects: None, |
| 367 | }; |
| 368 | let store_instr_id = InstructionId(instructions.len() as u32); |
| 369 | instructions.push(store_instr); |
| 370 | block.instructions.push(store_instr_id); |
| 371 | |
| 372 | let ret_loc = ret_loc.clone(); |
| 373 | block.terminal = Terminal::Goto { |
| 374 | block: return_target, |
| 375 | id: EvaluationOrder(0), |
| 376 | variant: GotoVariant::Break, |
| 377 | loc: ret_loc, |
| 378 | }; |
| 379 | } |
| 380 | } |
| 381 | |
| 382 | /// Emits a DeclareLocal instruction for the result temporary. |
| 383 | fn declare_temporary( |
| 384 | env: &mut Environment, |
| 385 | func: &mut HirFunction, |
| 386 | block_id: BlockId, |
| 387 | result: &Place, |
| 388 | ) { |
| 389 | let declare_lvalue = create_temporary_place(env, result.loc.clone()); |
| 390 | let declare_instr = Instruction { |
| 391 | id: EvaluationOrder(0), |
| 392 | loc: GENERATED_SOURCE, |
| 393 | lvalue: declare_lvalue, |
| 394 | value: InstructionValue::DeclareLocal { |
| 395 | lvalue: LValue { |
| 396 | place: result.clone(), |
| 397 | kind: InstructionKind::Let, |
| 398 | }, |
| 399 | type_annotation: None, |
| 400 | loc: result.loc.clone(), |
| 401 | }, |
| 402 | effects: None, |
| 403 | }; |
| 404 | let instr_id = InstructionId(func.instructions.len() as u32); |
| 405 | func.instructions.push(declare_instr); |
| 406 | func.body |
| 407 | .blocks |
| 408 | .get_mut(&block_id) |
| 409 | .unwrap() |
| 410 | .instructions |
| 411 | .push(instr_id); |
| 412 | } |
| 413 | |
| 414 | /// Promote a temporary identifier to a named identifier. |
| 415 | fn promote_temporary(env: &mut Environment, identifier_id: IdentifierId) { |
| 416 | let decl_id = env.identifiers[identifier_id.0 as usize].declaration_id; |
| 417 | env.identifiers[identifier_id.0 as usize].name = |
| 418 | Some(IdentifierName::Promoted(format!("#t{}", decl_id.0))); |
| 419 | } |