diff --git a/src/react_compiler/reactive_scopes/merge_reactive_scopes_that_invalidate_together.rs b/src/react_compiler/reactive_scopes/merge_reactive_scopes_that_invalidate_together.rs index 87ec7dd426e9..e38bb735f9b4 100644 --- a/src/react_compiler/reactive_scopes/merge_reactive_scopes_that_invalidate_together.rs +++ b/src/react_compiler/reactive_scopes/merge_reactive_scopes_that_invalidate_together.rs @@ -8,14 +8,18 @@ //! //! Corresponds to `src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts`. +use std::cmp::Reverse; +use std::collections::BinaryHeap; + use crate::collections::FxHashSet as HashSet; -use crate::collections::IdMap; +use crate::collections::{FxHashMap, IdMap}; use crate::diagnostics::CompilerError; use crate::hir::{ - DeclarationId, DependencyPathEntry, EvaluationOrder, HirVec, InstructionKind, InstructionValue, - Place, ReactiveBlock, ReactiveFunction, ReactiveScopeBlock, ReactiveScopeDependency, - ReactiveStatement, ReactiveValue, ScopeId, Type, + DeclarationId, DependencyPathEntry, EvaluationOrder, HirVec, IdentifierId, InstructionKind, + InstructionValue, Place, ReactiveBlock, ReactiveFunction, ReactiveScopeBlock, + ReactiveScopeDeclaration, ReactiveScopeDependency, ReactiveStatement, ReactiveValue, ScopeId, + Type, environment::Environment, object_shape::{BUILT_IN_ARRAY_ID, BUILT_IN_FUNCTION_ID, BUILT_IN_JSX_ID, BUILT_IN_OBJECT_ID}, }; @@ -140,6 +144,7 @@ impl<'a> MergeTransform<'a> { from: usize, to: usize, lvalues: HashSet, + declarations: CandidateDeclarations, } let mut current: Option = None; @@ -252,6 +257,7 @@ impl<'a> MergeTransform<'a> { let current_scope_id = c.scope_id; if can_merge_scopes( current_scope_id, + &c.declarations, next_scope_id, self.env, &self.temporaries, @@ -266,27 +272,16 @@ impl<'a> MergeTransform<'a> { self.env.scopes[next_scope_id.0 as usize].range.end; let current_range_end = self.env.scopes[current_scope_id.0 as usize].range.end; - self.env.scopes[current_scope_id.0 as usize].range.end = + let range_end = EvaluationOrder(current_range_end.0.max(next_range_end.0)); + self.env.scopes[current_scope_id.0 as usize].range.end = range_end; // Merge declarations from next into current - let next_decls = self.env.scopes[next_scope_id.0 as usize] - .declarations - .clone(); - for (key, value) in next_decls { - let current_decls = - &mut self.env.scopes[current_scope_id.0 as usize].declarations; - if let Some(existing) = - current_decls.iter_mut().find(|(k, _)| *k == key) - { - existing.1 = value; - } else { - current_decls.push((key, value)); - } - } + c.declarations + .extend(next_scope_id, &self.last_usage, self.env); // Prune declarations that are no longer used after the merged scope - update_scope_declarations(current_scope_id, &self.last_usage, self.env); + c.declarations.update_scope_declarations(range_end); c.to = i + 1; c.lvalues.clear(); @@ -310,6 +305,11 @@ impl<'a> MergeTransform<'a> { from: i, to: i + 1, lvalues: HashSet::default(), + declarations: CandidateDeclarations::new( + next_scope_id, + &self.last_usage, + self.env, + ), }); } } @@ -321,6 +321,11 @@ impl<'a> MergeTransform<'a> { from: i, to: i + 1, lvalues: HashSet::default(), + declarations: CandidateDeclarations::new( + next_scope_id, + &self.last_usage, + self.env, + ), }); } } @@ -343,7 +348,7 @@ impl<'a> MergeTransform<'a> { let mut index = 0; let all_stmts: Vec = std::mem::take(block); - for entry in &merged { + for entry in merged { // Push everything before the merge range while index < entry.from { next_instructions.push(all_stmts[index].clone()); @@ -360,6 +365,8 @@ impl<'a> MergeTransform<'a> { )); } }; + self.env.scopes[merged_scope.scope.0 as usize].declarations = + entry.declarations.into_declarations(); index += 1; while index < entry.to { let stmt = &all_stmts[index]; @@ -395,23 +402,105 @@ impl<'a> MergeTransform<'a> { // Helper functions // ============================================================================= -/// Updates scope declarations to remove any that are not used after the scope. -fn update_scope_declarations( - scope_id: ScopeId, - last_usage: &IdMap, - env: &mut Environment, -) { - let range_end = env.scopes[scope_id.0 as usize].range.end; - env.scopes[scope_id.0 as usize] - .declarations - .retain(|(_id, decl)| { - let decl_declaration_id = env.identifiers[decl.identifier.0 as usize].declaration_id; - match last_usage.get(decl_declaration_id) { - Some(last_used_at) => *last_used_at >= range_end, - // If not tracked, keep the declaration (conservative) - None => true, +/// Not in upstream: the declarations of a merge candidate, as the `Map` that TS keeps on the scope. +struct CandidateDeclarations { + /// Each declaration with its position in insertion order, in which RenameVariables names them. + entries: FxHashMap, + next_position: u32, + /// The entries that a later range end removes, earliest last usage first. + by_last_usage: BinaryHeap>, + /// The number of entries per declaration id. + declaration_ids: FxHashMap, +} + +impl CandidateDeclarations { + fn new( + scope_id: ScopeId, + last_usage: &IdMap, + env: &Environment, + ) -> Self { + let mut declarations = Self { + entries: FxHashMap::default(), + next_position: 0, + by_last_usage: BinaryHeap::new(), + declaration_ids: FxHashMap::default(), + }; + declarations.extend(scope_id, last_usage, env); + declarations + } + + /// TS: `declarations.set(key, value)` for each declaration of the scope. + fn extend( + &mut self, + scope_id: ScopeId, + last_usage: &IdMap, + env: &Environment, + ) { + for (key, value) in &env.scopes[scope_id.0 as usize].declarations { + // A declaration is keyed by its identifier, so a new value has the same last usage. + debug_assert_eq!(*key, value.identifier); + match self.entries.entry(*key) { + std::collections::hash_map::Entry::Occupied(mut existing) => { + existing.get_mut().1 = value.clone(); + } + std::collections::hash_map::Entry::Vacant(vacant) => { + vacant.insert((self.next_position, value.clone())); + self.next_position += 1; + let declaration_id = + env.identifiers[value.identifier.0 as usize].declaration_id; + *self.declaration_ids.entry(declaration_id).or_insert(0) += 1; + // If not tracked, keep the declaration (conservative) + if let Some(&last_used_at) = last_usage.get(declaration_id) { + self.by_last_usage + .push(Reverse((last_used_at, *key, declaration_id))); + } + } } - }); + } + } + + /// Updates scope declarations to remove any that are not used after the scope. + fn update_scope_declarations(&mut self, range_end: EvaluationOrder) { + while let Some(&Reverse((last_used_at, key, declaration_id))) = self.by_last_usage.peek() { + if last_used_at >= range_end { + break; + } + self.by_last_usage.pop(); + self.entries.remove(&key); + if let Some(count) = self.declaration_ids.get_mut(&declaration_id) { + *count -= 1; + } + } + } + + fn len(&self) -> usize { + self.entries.len() + } + + fn iter(&self) -> impl Iterator { + self.entries.values().map(|(_position, decl)| decl) + } + + fn has_declaration_id(&self, declaration_id: DeclarationId) -> bool { + self.declaration_ids + .get(&declaration_id) + .is_some_and(|count| *count > 0) + } + + fn into_declarations(self) -> HirVec<(IdentifierId, ReactiveScopeDeclaration)> { + let mut entries: Vec<_> = self.entries.into_iter().collect(); + bun_collections::index_sort::sort_slice_unstable_by( + &mut entries, + |(_, (a, _)), (_, (b, _))| a.cmp(b), + ); + let mut declarations = crate::hir_vec![]; + declarations.extend( + entries + .into_iter() + .map(|(key, (_position, decl))| (key, decl)), + ); + declarations + } } /// Returns whether all lvalues are last used at or before the given scope. @@ -435,6 +524,7 @@ fn are_lvalues_last_used_by_scope( /// Check if two scopes can be merged. fn can_merge_scopes( current_id: ScopeId, + current_declarations: &CandidateDeclarations, next_id: ScopeId, env: &Environment, temporaries: &IdMap, @@ -453,25 +543,27 @@ fn can_merge_scopes( } // Merge scopes where outputs of current are inputs of next - // Build synthetic dependencies from current's declarations - let current_decl_deps: Vec = current - .declarations - .iter() - .map(|(_key, decl)| ReactiveScopeDependency { - identifier: decl.identifier, - reactive: true, - path: crate::hir_vec![], - loc: None, - }) - .collect(); + // Not in upstream: the synthetic dependencies can only be equal when the lengths match. + if current_declarations.len() == next.dependencies.len() { + // Build synthetic dependencies from current's declarations + let current_decl_deps: Vec = current_declarations + .iter() + .map(|decl| ReactiveScopeDependency { + identifier: decl.identifier, + reactive: true, + path: crate::hir_vec![], + loc: None, + }) + .collect(); - if are_equal_dependencies(¤t_decl_deps, &next.dependencies, env) { - return true; + if are_equal_dependencies(¤t_decl_deps, &next.dependencies, env) { + return true; + } } // Check if all next deps have empty paths, always-invalidating types, // and correspond to current declarations (possibly through temporaries) - if !next.dependencies.is_empty() + !next.dependencies.is_empty() && next.dependencies.iter().all(|dep| { if !dep.path.is_empty() { return false; @@ -481,16 +573,11 @@ fn can_merge_scopes( return false; } let dep_decl = env.identifiers[dep.identifier.0 as usize].declaration_id; - current.declarations.iter().any(|(_key, decl)| { - let decl_decl_id = env.identifiers[decl.identifier.0 as usize].declaration_id; - decl_decl_id == dep_decl || temporaries.get(dep_decl).copied() == Some(decl_decl_id) - }) + current_declarations.has_declaration_id(dep_decl) + || temporaries + .get(dep_decl) + .is_some_and(|decl_id| current_declarations.has_declaration_id(*decl_id)) }) - { - return true; - } - - false } /// Check if a type is always invalidating (guaranteed to change when inputs change). diff --git a/test/bundler/transpiler/react-compiler.test.ts b/test/bundler/transpiler/react-compiler.test.ts index 536449e02a89..dbaddf250053 100644 --- a/test/bundler/transpiler/react-compiler.test.ts +++ b/test/bundler/transpiler/react-compiler.test.ts @@ -4382,6 +4382,109 @@ test("react-compiler keeps one copy of each dependency of a phi", async () => { expect(ladder.peakMB - empty.peakMB).toBeLessThan(100); }); +// MergeReactiveScopesThatInvalidateTogether collects the declarations of every +// scope that it merges into one. They were a list: for each scope it merged, +// the pass searched the list, ran `retain` over it, and scanned it once per +// dependency of that scope. Each pair below is two scopes, and all of them +// merge into one. Every `` stays declared until the `
` reads it, so N +// pairs took N^2 steps. +// +// The time of a whole build depends on every other pass too, so the test reads +// the time of this pass from BUN_REACT_COMPILER_TIMING. Only debug and ASAN +// builds have that report, and the test needs a debug build: optimized code +// needs several times as many pairs for the same ratio. +test.skipIf(!isDebug)( + "react-compiler MergeReactiveScopesThatInvalidateTogether time does not grow with the square of the scopes it merges into one", + async () => { + const pairs = 1600; + const perComponent = 10; + const pair = ""; + const source = (components: number, pairsEach: number) => + Array.from( + { length: components }, + (_, c) => + `export function App${c}(props) {\n return
${Buffer.alloc(pair.length * pairsEach, pair)}
;\n}\n`, + ).join(""); + using dir = tempDir("react-compiler-merge-scopes", { + "control.jsx": source(pairs / perComponent, perComponent), + "large.jsx": source(1, pairs), + }); + + // The milliseconds that the pass took in all the components of `entry`. + const passTime = async (entry: string, components: number) => { + await using proc = Bun.spawn({ + cmd: [bunExe(), "build", "--react-compiler", "--target=browser", "--external=*", entry], + env: { ...bunEnv, BUN_REACT_COMPILER_TIMING: "1" }, + cwd: String(dir), + stdout: "pipe", + stderr: "pipe", + }); + const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]); + const [, ms, calls] = + stderr.match(/([\d.]+)ms +[\d.]+% +(\d+)× +MergeReactiveScopesThatInvalidateTogether$/m) ?? []; + // Every component compiled, and each one is a single scope: one slot for + // `props.x` and one for the `
`. + expect({ calls: calls ?? stderr, scopes: stdout.match(/\b_c\(\d+\)/g), exitCode }).toEqual({ + calls: String(components), + scopes: Array(components).fill("_c(2)"), + exitCode: 0, + }); + return Number(ms); + }; + + // Both builds merge the same scopes, so the pass takes about as long in one + // as in the other. Without the fix the large component takes 6 times as + // long. The times are wall time. The two builds of a round run at the same + // time, under the same load, and the test takes the best of up to three + // rounds. + let control = Infinity; + let large = Infinity; + for (let round = 0; round < 3; round++) { + const [controlTime, largeTime] = await Promise.all([ + passTime("control.jsx", pairs / perComponent), + passTime("large.jsx", 1), + ]); + control = Math.min(control, controlTime); + large = Math.min(large, largeTime); + if (large / control < 3) break; + } + expect(large / control).toBeLessThan(3); + }, + // A round takes 5 seconds, and the test takes three rounds when it fails. + 60_000, +); + +// The scopes that merge into one keep their declarations in the order in which +// they merged, and RenameVariables names the declarations in that order. Each +// odd value is a nested array: its inner array is a declaration of the merged +// scope until the outer array merges and is its last use. +test("react-compiler names the declarations of a merged scope in source order", async () => { + const count = 64; + const values = Array.from({ length: count }, (_, i) => (i % 2 ? `{[["v${i}"]]}` : `{["v${i}"]}`)); + using dir = tempDir("react-compiler-merge-order", { + "entry.jsx": ` + export function App(props) { + return
${values.join("")}
; + } + `, + }); + + const result = await Bun.build({ + entrypoints: [join(String(dir), "entry.jsx")], + target: "browser", + external: ["*"], + reactCompiler: true, + }); + const output = await result.outputs[0].text(); + + // `t7 = [["v7"]]` + const names = Array.from(output.matchAll(/\bt(\d+) = \[\s*\[?\s*"v(\d+)"/g), ([, name, value]) => [name, value]); + expect(names).toEqual(Array.from({ length: count }, (_, i) => [String(i), String(i)])); + // One scope declares every value. The other one declares the `
` and + // depends on `props.t`. + expect(output).toContain(`_c(${count + 2})`); +}); + // validate_locals_not_reassigned_after_render (src/react_compiler/validation) // records the locals a component's closures capture while walking the // component body, and reports a nested function that assigns to one of them,