@samitouri / QOS-React-2 / commits / b939280bb3

React Compiler: avoid deep ast clone and quadratic-time codegen (#37206)

tl;dr this reduces peak memory allocation by 5-15%, and reduces codegen time by 30-70% depending on payload. On heavy components with deep ASTs the impact is more exaggerated. ## Summary Codegen of temp vars records the expressions they replaced, so that they can be unwound. In the TS version this uses a `Map` and stores pointers to AST nodes - relatively cheap. For borrow-checking reasons, the Rust version clones the AST. This results in recurring deep clones, making codegen accidentally quadratic and using significant amounts of memory. This introduces a convenience data structure for emitting temp vars in an unwindable manner, without heavy AST allocation. It also avoids a separate AST deep clone when propagating null values. ## How did you test this change? All fixtures pass with byte-identical outputs. Ran this against real codebases and pathological benchmark cases, confirming byte-identical output as well.

Andrew Imm committed Aug 24, 2026 at 10:38 UTC b939280bb3e2d79f62fd62aa264cbf2a3d5b90a2

3 files changed +171 -46

compiler/crates/react_compiler_inference/src/propagate_scope_dependencies_hir.rs

+19 -15

@@ -1472,32 +1472,36 @@ fn recursively_propagate_non_null(

1472 }

1473

1474 // Compute intersection of 'done' neighbors only (filter out 'active' = cycle nodes)

1475 - let done_neighbor_sets: Vec<BTreeSet<usize>> = neighbors

1476 - .iter()

1477 - .filter(|n| traversal_state.get(n) == Some(&TraversalState::Done))

1478 - .filter_map(|n| working.get(n).cloned())

1479 - .collect();

1475 + let neighbor_intersection = {

1476 + let done_neighbor_sets: Vec<&BTreeSet<usize>> = neighbors

1477 + .iter()

1478 + .filter(|n| traversal_state.get(n) == Some(&TraversalState::Done))

1479 + .filter_map(|n| working.get(n))

1480 + .collect();

1481

1481 - let neighbor_intersection = if done_neighbor_sets.is_empty() {

1482 - BTreeSet::new()

1483 - } else {

1484 - let mut iter = done_neighbor_sets.into_iter();

1485 - let first = iter.next().unwrap();

1486 - iter.fold(first, |acc, s| acc.intersection(&s).copied().collect())

1482 + match done_neighbor_sets.split_first() {

1483 + None => BTreeSet::new(),

1484 + Some((first, rest)) => rest.iter().fold((*first).clone(), |acc, s| {

1485 + acc.intersection(s).copied().collect()

1486 + }),

1487 + }

1488 };

1489

1489 - let prev_objects = working.get(&node_id).cloned().unwrap_or_default();

1490 + // Temporarily remove the previous set out of the map so it can be safely

1491 + // borrowed and compared without a heavy deep clone.

1492 + let prev_objects = working.remove(&node_id).unwrap_or_default();

1493 let mut merged: BTreeSet<usize> = prev_objects

1494 .union(&neighbor_intersection)

1495 .copied()

1496 .collect();

1497 reduce_maybe_optional_chains(&mut merged, registry);

1498

1496 - working.insert(node_id, merged.clone());

1497 - traversal_state.insert(node_id, TraversalState::Done);

1498 -

1499 // Compare with previous value — can't just check size due to reduce_maybe_optional_chains

1500 changed |= prev_objects != merged;

1501 +

1502 + working.insert(node_id, merged);

1503 + traversal_state.insert(node_id, TraversalState::Done);

1504 +

1505 changed

1506 }

1507

compiler/crates/react_compiler_reactive_scopes/src/codegen_reactive_function.rs

+145 -25

@@ -555,14 +555,121 @@ pub fn codegen_function(

555 // Context

556 // =============================================================================

557

558 -type Temporaries = FxHashMap<DeclarationId, Option<ExpressionOrJsxText>>;

559 -

558 #[derive(Clone)]

559 enum ExpressionOrJsxText {

560 Expression(Expression),

561 JsxText(JSXText),

562 }

563

564 +/// The entry a write to [`Temporaries`] displaced, kept so the write can be

565 +/// undone.

566 +///

567 +/// The expression is boxed because `ExpressionOrJsxText` is ~900 bytes (it

568 +/// inlines an `Expression`). Unboxed, the undo log would be a `Vec` of

569 +/// ~900-byte slots that are almost always `Absent`, costing more peak heap than

570 +/// the copy it replaces on shallow functions. Boxed, an entry is 16 bytes and

571 +/// only allocates when a write actually displaces a buffered expression.

572 +enum Displaced {

573 + /// The key was not present before the write.

574 + Absent,

575 + /// The key was present as a declared temporary with no buffered value.

576 + Empty,

577 + /// The key was present with this buffered value.

578 + Value(Box<ExpressionOrJsxText>),

579 +}

580 +

581 +/// A position in a [`Temporaries`] undo log, produced by [`Temporaries::mark`].

582 +#[derive(Clone, Copy)]

583 +struct TempMark(usize);

584 +

585 +/// Expressions buffered for temporaries that have not been emitted yet, plus an

586 +/// undo log allowing a nested block or scope to be codegen'd and its additions

587 +/// discarded.

588 +///

589 +/// The TypeScript implementation snapshots this with `new Map(cx.temp)`, which

590 +/// is a *shallow* copy: it duplicates references, not the AST nodes behind them.

591 +/// The equivalent Rust `.clone()` deep-copies every buffered `Expression` tree,

592 +/// which made codegen quadratic in component size and dominated both allocation

593 +/// volume and peak heap.

594 +///

595 +/// `TS CodegenReactiveFunction.codegenBlock` asserts that pre-existing entries

596 +/// are never mutated ("Expected temporary value to be unchanged"), so a

597 +/// snapshot's only job is to discard entries added by the nested block. Because

598 +/// entries are only ever inserted (never removed, nor mutated in place),

599 +/// rewinding an insert log restores the map exactly, with no copying.

600 +///

601 +/// All writes go through [`Temporaries::set`] so the log cannot drift out of

602 +/// sync with the map.

603 +#[derive(Default)]

604 +struct Temporaries {

605 + values: FxHashMap<DeclarationId, Option<ExpressionOrJsxText>>,

606 + journal: Vec<(DeclarationId, Displaced)>,

607 +}

608 +

609 +impl Temporaries {

610 + fn get(&self, declaration_id: DeclarationId) -> Option<&Option<ExpressionOrJsxText>> {

611 + self.values.get(&declaration_id)

612 + }

613 +

614 + fn contains_key(&self, declaration_id: DeclarationId) -> bool {

615 + self.values.contains_key(&declaration_id)

616 + }

617 +

618 + /// Buffers `value` for `declaration_id`, journaling the displaced entry.

619 + /// `HashMap::insert` returns that entry by move, so journaling costs no

620 + /// clones.

621 + fn set(&mut self, declaration_id: DeclarationId, value: Option<ExpressionOrJsxText>) {

622 + let displaced = match self.values.insert(declaration_id, value) {

623 + None => Displaced::Absent,

624 + Some(None) => Displaced::Empty,

625 + Some(Some(previous)) => Displaced::Value(Box::new(previous)),

626 + };

627 + self.journal.push((declaration_id, displaced));

628 + }

629 +

630 + /// Marks the current state, for a later [`Temporaries::rewind`].

631 + fn mark(&self) -> TempMark {

632 + TempMark(self.journal.len())

633 + }

634 +

635 + /// Restores the state captured by `mark`, discarding every write since.

636 + fn rewind(&mut self, mark: TempMark) {

637 + while self.journal.len() > mark.0 {

638 + let (declaration_id, displaced) = self.journal.pop().unwrap();

639 + match displaced {

640 + Displaced::Absent => {

641 + self.values.remove(&declaration_id);

642 + }

643 + Displaced::Empty => {

644 + self.values.insert(declaration_id, None);

645 + }

646 + Displaced::Value(previous) => {

647 + self.values.insert(declaration_id, Some(*previous));

648 + }

649 + }

650 + }

651 + }

652 +

653 + /// Hands the buffered expressions to a nested function's context, which may

654 + /// read them but must not leak its own additions back out.

655 + ///

656 + /// The borrower gets a fresh log, so [`Temporaries::reclaim`] can undo

657 + /// exactly the borrower's writes rather than the lender's whole history.

658 + fn lend(&mut self) -> Temporaries {

659 + Temporaries {

660 + values: std::mem::take(&mut self.values),

661 + journal: Vec::new(),

662 + }

663 + }

664 +

665 + /// Takes back expressions handed out by [`Temporaries::lend`], discarding

666 + /// every write the borrower made.

667 + fn reclaim(&mut self, mut lent: Temporaries) {

668 + lent.rewind(TempMark(0));

669 + self.values = lent.values;

670 + }

671 +}

672 +

673 struct Context<'env> {

674 env: &'env mut Environment,

675 #[allow(dead_code)]

@@ -594,7 +701,7 @@ impl<'env> Context<'env> {

701 fn_name,

702 next_cache_index: 0,

703 declarations: FxHashSet::default(),

597 - temp: FxHashMap::default(),

704 + temp: Temporaries::default(),

705 object_methods: FxHashMap::default(),

706 unique_identifiers,

707 fbt_operands,

@@ -653,8 +760,8 @@ fn codegen_reactive_function(

760 ParamPattern::Place(p) => p,

761 ParamPattern::Spread(sp) => &sp.place,

762 };

656 - let ident = &cx.env.identifiers[place.identifier.0 as usize];

657 - cx.temp.insert(ident.declaration_id, None);

763 + let declaration_id = cx.env.identifiers[place.identifier.0 as usize].declaration_id;

764 + cx.temp.set(declaration_id, None);

765 cx.declare(place.identifier);

766 }

767

@@ -732,9 +839,9 @@ fn convert_parameter(

839 // =============================================================================

840

841 fn codegen_block(cx: &mut Context, block: &ReactiveBlock) -> Result<BlockStatement, CompilerError> {

735 - let temp_snapshot: Temporaries = cx.temp.clone();

842 + let mark = cx.temp.mark();

843 let result = codegen_block_no_reset(cx, block)?;

737 - cx.temp = temp_snapshot;

844 + cx.temp.rewind(mark);

845 Ok(result)

846 }

847

@@ -758,9 +865,9 @@ fn codegen_block_no_reset(

865 scope,

866 instructions,

867 }) => {

761 - let temp_snapshot = cx.temp.clone();

868 + let mark = cx.temp.mark();

869 codegen_reactive_scope(cx, &mut statements, *scope, instructions)?;

763 - cx.temp = temp_snapshot;

870 + cx.temp.rewind(mark);

871 }

872 ReactiveStatement::Terminal(term_stmt) => {

873 let stmt = codegen_terminal(cx, &term_stmt.terminal)?;

@@ -1258,8 +1365,9 @@ fn codegen_terminal(

1365 } => {

1366 let catch_param = match handler_binding.as_ref() {

1367 Some(binding) => {

1261 - let ident = &cx.env.identifiers[binding.identifier.0 as usize];

1262 - cx.temp.insert(ident.declaration_id, None);

1368 + let declaration_id =

1369 + cx.env.identifiers[binding.identifier.0 as usize].declaration_id;

1370 + cx.temp.set(declaration_id, None);

1371 Some(PatternLike::Identifier(convert_identifier(

1372 binding.identifier,

1373 cx.env,

@@ -1724,8 +1832,10 @@ fn codegen_store_or_declare(

1832 // Register temporaries for unnamed pattern operands

1833 for place in react_compiler_hir::visitors::each_pattern_operand(&lvalue.pattern) {

1834 let ident = &cx.env.identifiers[place.identifier.0 as usize];

1727 - if kind != InstructionKind::Reassign && ident.name.is_none() {

1728 - cx.temp.insert(ident.declaration_id, None);

1835 + let declaration_id = ident.declaration_id;

1836 + let is_unnamed = ident.name.is_none();

1837 + if kind != InstructionKind::Reassign && is_unnamed {

1838 + cx.temp.set(declaration_id, None);

1839 }

1840 }

1841 let rhs = codegen_place_to_expression(cx, val)?;

@@ -1838,11 +1948,10 @@ fn emit_store(

1948 ReactiveValue::Instruction(InstructionValue::StoreContext { .. })

1949 );

1950 if !is_store_context {

1841 - let ident = &cx.env.identifiers[lvalue_place.identifier.0 as usize];

1842 - cx.temp.insert(

1843 - ident.declaration_id,

1844 - Some(ExpressionOrJsxText::Expression(expr)),

1845 - );

1951 + let declaration_id =

1952 + cx.env.identifiers[lvalue_place.identifier.0 as usize].declaration_id;

1953 + cx.temp

1954 + .set(declaration_id, Some(ExpressionOrJsxText::Expression(expr)));

1955 return Ok(None);

1956 } else {

1957 let stmt =

@@ -1886,9 +1995,10 @@ fn codegen_instruction(

1995 }));

1996 };

1997 let ident = &cx.env.identifiers[lvalue.identifier.0 as usize];

1998 + let declaration_id = ident.declaration_id;

1999 if ident.name.is_none() {

2000 // temporary

1891 - cx.temp.insert(ident.declaration_id, Some(value));

2001 + cx.temp.set(declaration_id, Some(value));

2002 return Ok(Statement::EmptyStatement(EmptyStatement {

2003 base: BaseNode::typed("EmptyStatement"),

2004 }));

@@ -2663,9 +2773,16 @@ fn codegen_function_expression(

2773 cx.unique_identifiers.clone(),

2774 cx.fbt_operands.clone(),

2775 );

2666 - inner_cx.temp = cx.temp.clone();

2776 + // The inner function reads the enclosing temporaries but must not leak its

2777 + // own back out. Lend the map to `inner_cx` and rewind its writes on the way

2778 + // out, rather than deep-cloning every buffered expression tree. The map is

2779 + // restored on the error path too, so `cx` is never left empty.

2780 + inner_cx.temp = cx.temp.lend();

2781

2668 - let fn_result = codegen_reactive_function(&mut inner_cx, &reactive_fn_mut)?;

2782 + let fn_result = codegen_reactive_function(&mut inner_cx, &reactive_fn_mut);

2783 +

2784 + cx.temp.reclaim(std::mem::take(&mut inner_cx.temp));

2785 + let fn_result = fn_result?;

2786

2787 let value = match expr_type {

2788 FunctionExpressionType::ArrowFunctionExpression => {

@@ -2798,9 +2915,12 @@ fn codegen_object_expression(

2915 cx.unique_identifiers.clone(),

2916 cx.fbt_operands.clone(),

2917 );

2801 - inner_cx.temp = cx.temp.clone();

2918 + inner_cx.temp = cx.temp.lend();

2919 +

2920 + let fn_result = codegen_reactive_function(&mut inner_cx, &reactive_fn_mut);

2921

2803 - let fn_result = codegen_reactive_function(&mut inner_cx, &reactive_fn_mut)?;

2922 + cx.temp.reclaim(std::mem::take(&mut inner_cx.temp));

2923 + let fn_result = fn_result?;

2924

2925 ast_properties.push(ast_expr::ObjectExpressionProperty::ObjectMethod(

2926 ast_expr::ObjectMethod {

@@ -3309,14 +3429,14 @@ fn codegen_place_to_expression(

3429

3430 fn codegen_place(cx: &mut Context, place: &Place) -> Result<ExpressionOrJsxText, CompilerError> {

3431 let ident = &cx.env.identifiers[place.identifier.0 as usize];

3312 - if let Some(tmp) = cx.temp.get(&ident.declaration_id) {

3432 + if let Some(tmp) = cx.temp.get(ident.declaration_id) {

3433 if let Some(val) = tmp {

3434 return Ok(val.clone());

3435 }

3436 // tmp is None — means declared but no temp value, fall through

3437 }

3438 // Check if it's an unnamed identifier without a temp

3319 - if ident.name.is_none() && !cx.temp.contains_key(&ident.declaration_id) {

3439 + if ident.name.is_none() && !cx.temp.contains_key(ident.declaration_id) {

3440 return Err(invariant_err(

3441 &format!(

3442 "[Codegen] No value found for temporary, identifier id={}",

compiler/crates/react_compiler_validation/src/validate_preserved_manual_memoization.rs

+7 -6

@@ -142,18 +142,19 @@ fn visit_scope(scope_block: &ReactiveScopeBlock, state: &mut VisitorState) {

142 if let Some(ref memo_state) = state.manual_memo_state {

143 if let Some(ref deps_from_source) = memo_state.deps_from_source {

144 let scope = &state.env.scopes[scope_block.scope.0 as usize];

145 + // `dependencies` still has to be cloned because `env` is passed

146 + // mutably below. `temporaries`, `decls` and `deps_from_source` do

147 + // not: they live in fields disjoint from `env`, so they can simply

148 + // be borrowed.

149 let deps = scope.dependencies.clone();

150 let memo_loc = memo_state.loc;

147 - let decls = memo_state.decls.clone();

148 - let deps_from_source = deps_from_source.clone();

149 - let temporaries = state.temporaries.clone();

151 for dep in &deps {

152 validate_inferred_dep(

153 dep.identifier,

154 &dep.path,

154 - &temporaries,

155 - &decls,

156 - &deps_from_source,

155 + &state.temporaries,

156 + &memo_state.decls,

157 + deps_from_source,

158 state.env,

159 memo_loc,

160 );