main
rs 670 lines 22.8 KB
Raw
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 //! Port of OutlineJsx from TypeScript.
7 //!
8 //! Outlines JSX expressions in callbacks into separate component functions.
9 //! This pass is conditional on `env.config.enable_jsx_outlining` (defaults to false).
10
11 use indexmap::{IndexMap, IndexSet};
12 use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
13
14 use react_compiler_hir::environment::Environment;
15 use react_compiler_hir::{
16 BasicBlock, BlockId, BlockKind, EvaluationOrder, FunctionId, HIR, HirFunction, IdentifierId,
17 IdentifierName, Instruction, InstructionId, InstructionKind, InstructionValue, JsxAttribute,
18 JsxTag, LValuePattern, NonLocalBinding, ObjectPattern, ObjectProperty, ObjectPropertyKey,
19 ObjectPropertyOrSpread, ObjectPropertyType, ParamPattern, Pattern, Place, ReactFunctionType,
20 ReturnVariant, Terminal,
21 };
22
23 /// Outline JSX expressions in inner functions into separate outlined components.
24 ///
25 /// Ported from TS `outlineJSX` in `Optimization/OutlineJsx.ts`.
26 pub fn outline_jsx(func: &mut HirFunction, env: &mut Environment) {
27 let mut outlined_fns: Vec<HirFunction> = Vec::new();
28 outline_jsx_impl(func, env, &mut outlined_fns);
29
30 for outlined_fn in outlined_fns {
31 env.outline_function(outlined_fn, Some(ReactFunctionType::Component));
32 }
33 }
34
35 /// Data about a JSX instruction for outlining
36 struct JsxInstrInfo {
37 instr_idx: usize, // index into func.instructions
38 #[allow(dead_code)]
39 instr_id: InstructionId, // the InstructionId
40 lvalue_id: IdentifierId,
41 eval_order: EvaluationOrder,
42 }
43
44 struct OutlinedJsxAttribute {
45 original_name: String,
46 new_name: String,
47 place: Place,
48 }
49
50 struct OutlinedResult {
51 instrs: Vec<Instruction>,
52 func: HirFunction,
53 }
54
55 fn outline_jsx_impl(
56 func: &mut HirFunction,
57 env: &mut Environment,
58 outlined_fns: &mut Vec<HirFunction>,
59 ) {
60 // Collect LoadGlobal instructions (tag -> instr)
61 let mut globals: FxHashMap<IdentifierId, usize> = FxHashMap::default(); // id -> instr_idx
62
63 // Process each block
64 let block_ids: Vec<BlockId> = func.body.blocks.keys().copied().collect();
65 for block_id in &block_ids {
66 let block = &func.body.blocks[block_id];
67 let instr_ids = block.instructions.clone();
68
69 let mut rewrite_instr: FxHashMap<EvaluationOrder, Vec<Instruction>> = FxHashMap::default();
70 let mut jsx_group: Vec<JsxInstrInfo> = Vec::new();
71 let mut children_ids: FxHashSet<IdentifierId> = FxHashSet::default();
72
73 // First pass: collect all instruction info without borrowing func mutably
74 enum InstrAction {
75 LoadGlobal {
76 lvalue_id: IdentifierId,
77 instr_idx: usize,
78 },
79 FunctionExpr {
80 func_id: FunctionId,
81 },
82 JsxExpr {
83 lvalue_id: IdentifierId,
84 instr_idx: usize,
85 eval_order: EvaluationOrder,
86 child_ids: Vec<IdentifierId>,
87 },
88 Other,
89 }
90
91 let mut actions: Vec<InstrAction> = Vec::new();
92 for i in (0..instr_ids.len()).rev() {
93 let iid = instr_ids[i];
94 let instr = &func.instructions[iid.0 as usize];
95 let lvalue_id = instr.lvalue.identifier;
96
97 match &instr.value {
98 InstructionValue::LoadGlobal { .. } => {
99 actions.push(InstrAction::LoadGlobal {
100 lvalue_id,
101 instr_idx: iid.0 as usize,
102 });
103 }
104 InstructionValue::FunctionExpression { lowered_func, .. } => {
105 actions.push(InstrAction::FunctionExpr {
106 func_id: lowered_func.func,
107 });
108 }
109 InstructionValue::JsxExpression { children, .. } => {
110 let child_ids = children
111 .as_ref()
112 .map(|kids| kids.iter().map(|c| c.identifier).collect())
113 .unwrap_or_default();
114 actions.push(InstrAction::JsxExpr {
115 lvalue_id,
116 instr_idx: iid.0 as usize,
117 eval_order: instr.id,
118 child_ids,
119 });
120 }
121 _ => {
122 actions.push(InstrAction::Other);
123 }
124 }
125 }
126
127 // Second pass: process actions
128 for action in actions {
129 match action {
130 InstrAction::LoadGlobal {
131 lvalue_id,
132 instr_idx,
133 } => {
134 globals.insert(lvalue_id, instr_idx);
135 }
136 InstrAction::FunctionExpr { func_id } => {
137 let mut inner_func = std::mem::replace(
138 &mut env.functions[func_id.0 as usize],
139 react_compiler_ssa::enter_ssa::placeholder_function(),
140 );
141 outline_jsx_impl(&mut inner_func, env, outlined_fns);
142 env.functions[func_id.0 as usize] = inner_func;
143 }
144 InstrAction::JsxExpr {
145 lvalue_id,
146 instr_idx,
147 eval_order,
148 child_ids,
149 } => {
150 if !children_ids.contains(&lvalue_id) {
151 process_and_outline_jsx(
152 func,
153 env,
154 &mut jsx_group,
155 &globals,
156 &mut rewrite_instr,
157 outlined_fns,
158 );
159 jsx_group.clear();
160 children_ids.clear();
161 }
162 jsx_group.push(JsxInstrInfo {
163 instr_idx,
164 instr_id: InstructionId(instr_idx as u32),
165 lvalue_id,
166 eval_order,
167 });
168 for child_id in child_ids {
169 children_ids.insert(child_id);
170 }
171 }
172 InstrAction::Other => {}
173 }
174 }
175 // Process remaining JSX group after the loop
176 process_and_outline_jsx(
177 func,
178 env,
179 &mut jsx_group,
180 &globals,
181 &mut rewrite_instr,
182 outlined_fns,
183 );
184 if !rewrite_instr.is_empty() {
185 let block = func.body.blocks.get_mut(block_id).unwrap();
186 let old_instr_ids = block.instructions.clone();
187 let mut new_instr_ids = Vec::new();
188 for &iid in &old_instr_ids {
189 let eval_order = func.instructions[iid.0 as usize].id;
190 if let Some(replacement_instrs) = rewrite_instr.get(&eval_order) {
191 // Add replacement instructions to the instruction table and reference them
192 for new_instr in replacement_instrs {
193 let new_idx = func.instructions.len();
194 func.instructions.push(new_instr.clone());
195 new_instr_ids.push(InstructionId(new_idx as u32));
196 }
197 } else {
198 new_instr_ids.push(iid);
199 }
200 }
201 let block = func.body.blocks.get_mut(block_id).unwrap();
202 block.instructions = new_instr_ids;
203
204 // Run dead code elimination after rewriting
205 super::dead_code_elimination(func, env);
206 }
207 }
208 }
209
210 fn process_and_outline_jsx(
211 func: &mut HirFunction,
212 env: &mut Environment,
213 jsx_group: &mut Vec<JsxInstrInfo>,
214 globals: &FxHashMap<IdentifierId, usize>,
215 rewrite_instr: &mut FxHashMap<EvaluationOrder, Vec<Instruction>>,
216 outlined_fns: &mut Vec<HirFunction>,
217 ) {
218 if jsx_group.len() <= 1 {
219 return;
220 }
221 // Sort by eval order ascending (TS: sort by a.id - b.id)
222 jsx_group.sort_by_key(|j| j.eval_order);
223
224 let result = process_jsx_group(func, env, jsx_group, globals);
225 if let Some(result) = result {
226 outlined_fns.push(result.func);
227 // Map from the LAST JSX instruction's eval order to the replacement instructions
228 // In the TS code, `state.jsx.at(0)` is the first element pushed during reverse iteration,
229 // which is the last JSX in forward block order (highest eval order).
230 // After sorting by eval_order ascending, that's jsx_group.last().
231 let last_eval_order = jsx_group.last().unwrap().eval_order;
232 rewrite_instr.insert(last_eval_order, result.instrs);
233 }
234 }
235
236 fn process_jsx_group(
237 func: &HirFunction,
238 env: &mut Environment,
239 jsx_group: &[JsxInstrInfo],
240 globals: &FxHashMap<IdentifierId, usize>,
241 ) -> Option<OutlinedResult> {
242 // Only outline in callbacks, not top-level components
243 if func.fn_type == ReactFunctionType::Component {
244 return None;
245 }
246
247 let props = collect_props(func, env, jsx_group)?;
248
249 let outlined_tag = env.generate_globally_unique_identifier_name(None);
250 let new_instrs = emit_outlined_jsx(func, env, jsx_group, &props, &outlined_tag)?;
251 let outlined_fn = emit_outlined_fn(func, env, jsx_group, &props, globals)?;
252
253 // Set the outlined function's id
254 let mut outlined_fn = outlined_fn;
255 outlined_fn.id = Some(outlined_tag);
256
257 Some(OutlinedResult {
258 instrs: new_instrs,
259 func: outlined_fn,
260 })
261 }
262
263 fn collect_props(
264 func: &HirFunction,
265 env: &mut Environment,
266 jsx_group: &[JsxInstrInfo],
267 ) -> Option<Vec<OutlinedJsxAttribute>> {
268 let mut id_counter = 1u32;
269 let mut seen: FxHashSet<String> = FxHashSet::default();
270 let mut attributes = Vec::new();
271 let jsx_ids: FxHashSet<IdentifierId> = jsx_group.iter().map(|j| j.lvalue_id).collect();
272
273 let mut generate_name = |old_name: &str, _env: &mut Environment| -> String {
274 let mut new_name = old_name.to_string();
275 while seen.contains(&new_name) {
276 new_name = format!("{}{}", old_name, id_counter);
277 id_counter += 1;
278 }
279 seen.insert(new_name.clone());
280 // TS: env.programContext.addNewReference(newName)
281 // We don't have programContext in Rust, but this is needed for unique name tracking
282 new_name
283 };
284
285 for info in jsx_group {
286 let instr = &func.instructions[info.instr_idx];
287 if let InstructionValue::JsxExpression {
288 props, children, ..
289 } = &instr.value
290 {
291 for attr in props {
292 match attr {
293 JsxAttribute::SpreadAttribute { .. } => return None,
294 JsxAttribute::Attribute { name, place } => {
295 let new_name = generate_name(name, env);
296 attributes.push(OutlinedJsxAttribute {
297 original_name: name.clone(),
298 new_name,
299 place: place.clone(),
300 });
301 }
302 }
303 }
304
305 if let Some(kids) = children {
306 for child in kids {
307 if jsx_ids.contains(&child.identifier) {
308 continue;
309 }
310 // Promote the child's identifier to a named temporary
311 let child_id = child.identifier;
312 let decl_id = env.identifiers[child_id.0 as usize].declaration_id;
313 if env.identifiers[child_id.0 as usize].name.is_none() {
314 env.identifiers[child_id.0 as usize].name =
315 Some(IdentifierName::Promoted(format!("#t{}", decl_id.0)));
316 }
317
318 let child_name = match &env.identifiers[child_id.0 as usize].name {
319 Some(IdentifierName::Named(n)) => n.clone(),
320 Some(IdentifierName::Promoted(n)) => n.clone(),
321 None => format!("#t{}", decl_id.0),
322 };
323 let new_name = generate_name("t", env);
324 attributes.push(OutlinedJsxAttribute {
325 original_name: child_name,
326 new_name,
327 place: child.clone(),
328 });
329 }
330 }
331 }
332 }
333
334 Some(attributes)
335 }
336
337 fn emit_outlined_jsx(
338 func: &HirFunction,
339 env: &mut Environment,
340 jsx_group: &[JsxInstrInfo],
341 outlined_props: &[OutlinedJsxAttribute],
342 outlined_tag: &str,
343 ) -> Option<Vec<Instruction>> {
344 let props: Vec<JsxAttribute> = outlined_props
345 .iter()
346 .map(|p| JsxAttribute::Attribute {
347 name: p.new_name.clone(),
348 place: p.place.clone(),
349 })
350 .collect();
351
352 // Create LoadGlobal for the outlined component
353 let load_id = env.next_identifier_id();
354 // Promote it as a JSX tag temporary
355 let decl_id = env.identifiers[load_id.0 as usize].declaration_id;
356 env.identifiers[load_id.0 as usize].name =
357 Some(IdentifierName::Promoted(format!("#T{}", decl_id.0)));
358
359 let load_place = Place {
360 identifier: load_id,
361 effect: react_compiler_hir::Effect::Unknown,
362 reactive: false,
363 loc: None,
364 };
365
366 let load_jsx = Instruction {
367 id: EvaluationOrder(0),
368 lvalue: load_place.clone(),
369 value: InstructionValue::LoadGlobal {
370 binding: NonLocalBinding::ModuleLocal {
371 name: outlined_tag.to_string(),
372 },
373 loc: None,
374 },
375 loc: None,
376 effects: None,
377 };
378
379 // Create the replacement JsxExpression using the last JSX instruction's lvalue
380 let last_info = jsx_group.last().unwrap();
381 let last_instr = &func.instructions[last_info.instr_idx];
382 let jsx_expr = Instruction {
383 id: EvaluationOrder(0),
384 lvalue: last_instr.lvalue.clone(),
385 value: InstructionValue::JsxExpression {
386 tag: JsxTag::Place(load_place),
387 props,
388 children: None,
389 loc: None,
390 opening_loc: None,
391 closing_loc: None,
392 },
393 loc: None,
394 effects: None,
395 };
396
397 Some(vec![load_jsx, jsx_expr])
398 }
399
400 fn emit_outlined_fn(
401 func: &HirFunction,
402 env: &mut Environment,
403 jsx_group: &[JsxInstrInfo],
404 old_props: &[OutlinedJsxAttribute],
405 globals: &FxHashMap<IdentifierId, usize>,
406 ) -> Option<HirFunction> {
407 let old_to_new_props = create_old_to_new_props_mapping(env, old_props);
408
409 // Create props parameter
410 let props_obj_id = env.next_identifier_id();
411 let decl_id = env.identifiers[props_obj_id.0 as usize].declaration_id;
412 env.identifiers[props_obj_id.0 as usize].name =
413 Some(IdentifierName::Promoted(format!("#t{}", decl_id.0)));
414 let props_obj = Place {
415 identifier: props_obj_id,
416 effect: react_compiler_hir::Effect::Unknown,
417 reactive: false,
418 loc: None,
419 };
420
421 // Create destructure instruction
422 let destructure_instr = emit_destructure_props(env, &props_obj, &old_to_new_props);
423
424 // Emit load globals for JSX tags
425 let load_global_instrs = emit_load_globals(func, jsx_group, globals)?;
426
427 // Emit updated JSX instructions
428 let updated_jsx_instrs = emit_updated_jsx(func, jsx_group, &old_to_new_props);
429
430 // Build instructions list
431 let mut instructions = Vec::new();
432 instructions.push(destructure_instr);
433 instructions.extend(load_global_instrs);
434 instructions.extend(updated_jsx_instrs);
435
436 // Build instruction table and instruction IDs
437 let mut instr_table = Vec::new();
438 let mut instr_ids = Vec::new();
439 for instr in instructions {
440 let idx = instr_table.len();
441 instr_table.push(instr);
442 instr_ids.push(InstructionId(idx as u32));
443 }
444
445 // Return terminal uses the last instruction's lvalue
446 let last_lvalue = instr_table.last().unwrap().lvalue.clone();
447
448 // Create return place
449 let returns_id = env.next_identifier_id();
450 let returns_place = Place {
451 identifier: returns_id,
452 effect: react_compiler_hir::Effect::Unknown,
453 reactive: false,
454 loc: None,
455 };
456
457 let block = BasicBlock {
458 kind: BlockKind::Block,
459 id: BlockId(0),
460 instructions: instr_ids,
461 preds: IndexSet::default(),
462 terminal: Terminal::Return {
463 value: last_lvalue,
464 return_variant: ReturnVariant::Explicit,
465 id: EvaluationOrder(0),
466 loc: None,
467 effects: None,
468 },
469 phis: Vec::new(),
470 };
471
472 let mut blocks = IndexMap::default();
473 blocks.insert(BlockId(0), block);
474
475 let outlined_fn = HirFunction {
476 id: None,
477 name_hint: None,
478 fn_type: ReactFunctionType::Other,
479 params: vec![ParamPattern::Place(props_obj)],
480 return_type_annotation: None,
481 returns: returns_place,
482 context: Vec::new(),
483 body: HIR {
484 entry: BlockId(0),
485 blocks,
486 },
487 instructions: instr_table,
488 generator: false,
489 is_async: false,
490 directives: Vec::new(),
491 aliasing_effects: Some(vec![]),
492 loc: None,
493 };
494
495 Some(outlined_fn)
496 }
497
498 fn emit_load_globals(
499 func: &HirFunction,
500 jsx_group: &[JsxInstrInfo],
501 globals: &FxHashMap<IdentifierId, usize>,
502 ) -> Option<Vec<Instruction>> {
503 let mut instructions = Vec::new();
504 for info in jsx_group {
505 let instr = &func.instructions[info.instr_idx];
506 if let InstructionValue::JsxExpression { tag, .. } = &instr.value {
507 if let JsxTag::Place(tag_place) = tag {
508 let global_instr_idx = globals.get(&tag_place.identifier)?;
509 instructions.push(func.instructions[*global_instr_idx].clone());
510 }
511 }
512 }
513 Some(instructions)
514 }
515
516 fn emit_updated_jsx(
517 func: &HirFunction,
518 jsx_group: &[JsxInstrInfo],
519 old_to_new_props: &IndexMap<IdentifierId, OutlinedJsxAttribute, FxBuildHasher>,
520 ) -> Vec<Instruction> {
521 let jsx_ids: FxHashSet<IdentifierId> = jsx_group.iter().map(|j| j.lvalue_id).collect();
522 let mut new_instrs = Vec::new();
523
524 for info in jsx_group {
525 let instr = &func.instructions[info.instr_idx];
526 if let InstructionValue::JsxExpression {
527 tag,
528 props,
529 children,
530 loc,
531 opening_loc,
532 closing_loc,
533 } = &instr.value
534 {
535 let mut new_props = Vec::new();
536 for prop in props {
537 // TS: invariant(prop.kind === 'JsxAttribute', ...)
538 // Spread attributes would have caused collectProps to return null earlier
539 let (name, place) = match prop {
540 JsxAttribute::Attribute { name, place } => (name, place),
541 JsxAttribute::SpreadAttribute { .. } => {
542 unreachable!("Expected only JsxAttribute, not spread")
543 }
544 };
545 if name == "key" {
546 continue;
547 }
548 // TS: invariant(newProp !== undefined, ...)
549 let new_prop = old_to_new_props
550 .get(&place.identifier)
551 .expect("Expected a new property for identifier");
552 new_props.push(JsxAttribute::Attribute {
553 name: new_prop.original_name.clone(),
554 place: new_prop.place.clone(),
555 });
556 }
557
558 let new_children = children.as_ref().map(|kids| {
559 kids.iter()
560 .map(|child| {
561 if jsx_ids.contains(&child.identifier) {
562 child.clone()
563 } else {
564 // TS: invariant(newChild !== undefined, ...)
565 let new_prop = old_to_new_props
566 .get(&child.identifier)
567 .expect("Expected a new prop for child identifier");
568 new_prop.place.clone()
569 }
570 })
571 .collect()
572 });
573
574 new_instrs.push(Instruction {
575 id: instr.id,
576 lvalue: instr.lvalue.clone(),
577 value: InstructionValue::JsxExpression {
578 tag: tag.clone(),
579 props: new_props,
580 children: new_children,
581 loc: *loc,
582 opening_loc: *opening_loc,
583 closing_loc: *closing_loc,
584 },
585 loc: instr.loc,
586 effects: instr.effects.clone(),
587 });
588 }
589 }
590
591 new_instrs
592 }
593
594 fn create_old_to_new_props_mapping(
595 env: &mut Environment,
596 old_props: &[OutlinedJsxAttribute],
597 ) -> IndexMap<IdentifierId, OutlinedJsxAttribute, FxBuildHasher> {
598 let mut old_to_new = IndexMap::default();
599
600 for old_prop in old_props {
601 if old_prop.original_name == "key" {
602 continue;
603 }
604
605 let new_id = env.next_identifier_id();
606 env.identifiers[new_id.0 as usize].name =
607 Some(IdentifierName::Named(old_prop.new_name.clone()));
608
609 let new_place = Place {
610 identifier: new_id,
611 effect: react_compiler_hir::Effect::Unknown,
612 reactive: false,
613 loc: None,
614 };
615
616 old_to_new.insert(
617 old_prop.place.identifier,
618 OutlinedJsxAttribute {
619 original_name: old_prop.original_name.clone(),
620 new_name: old_prop.new_name.clone(),
621 place: new_place,
622 },
623 );
624 }
625
626 old_to_new
627 }
628
629 fn emit_destructure_props(
630 env: &mut Environment,
631 props_obj: &Place,
632 old_to_new_props: &IndexMap<IdentifierId, OutlinedJsxAttribute, FxBuildHasher>,
633 ) -> Instruction {
634 let mut properties = Vec::new();
635 for prop in old_to_new_props.values() {
636 properties.push(ObjectPropertyOrSpread::Property(ObjectProperty {
637 key: ObjectPropertyKey::String {
638 name: prop.new_name.clone(),
639 },
640 property_type: ObjectPropertyType::Property,
641 place: prop.place.clone(),
642 }));
643 }
644
645 let lvalue_id = env.next_identifier_id();
646 let lvalue = Place {
647 identifier: lvalue_id,
648 effect: react_compiler_hir::Effect::Unknown,
649 reactive: false,
650 loc: None,
651 };
652
653 Instruction {
654 id: EvaluationOrder(0),
655 lvalue,
656 value: InstructionValue::Destructure {
657 lvalue: LValuePattern {
658 pattern: Pattern::Object(ObjectPattern {
659 properties,
660 loc: None,
661 }),
662 kind: InstructionKind::Let,
663 },
664 value: props_obj.clone(),
665 loc: None,
666 },
667 loc: None,
668 effects: None,
669 }
670 }