merge_nodes_with_same_starting_point

Author: kdy1

turbopack/crates/turbopack-ecmascript/src/tree_shake/optimizations.rs

@@ -125,6 +125,11 @@ impl GraphOptimizer<'_> {
         did_work
     }
 
+    /// This function merges nodes that can only be reached from a single starting point.
+    /// Example:
+    /// If we have a graph with edges: A->B, B->C, A->C, B->E, D->E
+    /// Then B and C can only be reached from A, so they will be merged into A.
+    /// The resulting graph would have edges like: (A,B,C)->E, D->E
     pub(super) fn merge_nodes_with_same_starting_point<N>(
         &self,
         g: &mut Graph<Vec<N>, Dependency>,
@@ -134,64 +139,115 @@ impl GraphOptimizer<'_> {
         Self: Index<N, Output = ItemId>,
     {
         let mut did_work = false;
-        let mut nodes_to_merge = Vec::new();
+        let mut reachability: FxHashMap<_, FxHashSet<_>> = FxHashMap::default();
 
-        for node in g.node_indices() {
-            let items = g.node_weight(node).expect("Node should exist");
-            for item in items {
-                let item_id = &self[*item];
-                if matches!(item_id, ItemId::Group(..)) {
-                    let incoming_edges: Vec<_> =
-                        g.edges_directed(node, Direction::Incoming).collect();
-                    if incoming_edges.len() == 1 && incoming_edges[0].source() == node {
-                        nodes_to_merge.push(node);
-                        break;
-                    }
+        // Step 1: Build a reverse reachability map (which starting nodes can reach each node)
+        // We consider a "starting node" as one with no incoming edges
+        let starting_nodes: Vec<_> = g
+            .node_indices()
+            .filter(|&node| g.edges_directed(node, Direction::Incoming).count() == 0)
+            .collect();
+
+        // For each starting node, find all nodes reachable from it
+        for &start in &starting_nodes {
+            let mut visited = FxHashSet::default();
+            let mut queue = vec![start];
+
+            while let Some(node) = queue.pop() {
+                if !visited.insert(node) {
+                    continue;
+                }
+
+                // For each outgoing edge, add the target to queue
+                for edge in g.edges_directed(node, Direction::Outgoing) {
+                    let target = edge.target();
+                    queue.push(target);
+
+                    // Add this starting node to the set of starting nodes that can reach target
+                    reachability.entry(target).or_default().insert(start);
                 }
             }
         }
 
-        for node in nodes_to_merge {
-            let mut dependencies = g
-                .edges_directed(node, Direction::Outgoing)
-                .map(|e| (e.target(), *e.weight()))
-                .collect::<Vec<_>>();
+        // Step 2: Find nodes that are reachable from exactly one starting node
+        // and group them by that starting node
+        let mut merge_groups: FxHashMap<_, Vec<_>> = FxHashMap::default();
 
-            // Handle transitive dependencies
-            let mut visited = FxHashSet::default();
-            let mut stack = dependencies.clone();
-            while let Some((dependency, _weight)) = stack.pop() {
-                if visited.insert(dependency) {
-                    let transitive_deps = g
-                        .edges_directed(dependency, Direction::Outgoing)
-                        .map(|e| (e.target(), *e.weight()))
-                        .collect::<Vec<_>>();
-                    stack.extend(transitive_deps.clone());
-                    dependencies.extend(transitive_deps);
+        for node in g.node_indices() {
+            // Skip starting nodes
+            if starting_nodes.contains(&node) {
+                continue;
+            }
+
+            // Skip nodes that should not be merged
+            if self.should_not_merge_iter(g.node_weight(node).expect("Node should exist")) {
+                continue;
+            }
+
+            // If this node is reachable from exactly one starting node, add it to that group
+            if let Some(reachable_from) = reachability.get(&node) {
+                if reachable_from.len() == 1 {
+                    let start = *reachable_from.iter().next().unwrap();
+
+                    // Don't merge if the starting node should not be merged
+                    if self.should_not_merge_iter(g.node_weight(start).expect("Node should exist"))
+                    {
+                        continue;
+                    }
+
+                    merge_groups.entry(start).or_default().push(node);
                 }
             }
+        }
 
-            for (dependency, weight) in dependencies {
-                let edge = g
-                    .find_edge(node, dependency)
-                    .and_then(|e| g.edge_weight_mut(e));
-                match edge {
-                    Some(v) => {
-                        if matches!(v, Dependency::Weak) {
-                            *v = weight;
+        // Step 3: Merge nodes into their starting points
+        for (start, nodes_to_merge) in merge_groups {
+            if nodes_to_merge.is_empty() {
+                continue;
+            }
+
+            let mut nodes_to_remove = Vec::new();
+
+            for node in nodes_to_merge {
+                // Move outgoing edges from node to start
+                let outgoing_edges: Vec<_> = g
+                    .edges_directed(node, Direction::Outgoing)
+                    .map(|e| (e.target(), *e.weight()))
+                    .collect();
+
+                for (target, weight) in outgoing_edges {
+                    // If there's already an edge from start to target, only update if necessary
+                    let existing_edge = g.find_edge(start, target);
+                    match existing_edge {
+                        Some(e) => {
+                            let edge_weight = g.edge_weight_mut(e).unwrap();
+                            // Only upgrade from weak to strong dependency
+                            if matches!(edge_weight, Dependency::Weak)
+                                && !matches!(weight, Dependency::Weak)
+                            {
+                                *edge_weight = weight;
+                            }
+                        }
+                        None => {
+                            // Add a new edge
+                            g.add_edge(start, target, weight);
                         }
-                    }
-                    None => {
-                        g.add_edge(node, dependency, weight);
                     }
                 }
-            }
 
-            let items = g.node_weight(node).expect("Node should exist").clone();
-            g.node_weight_mut(node).unwrap().extend(items);
+                // Move items from this node to the starting node
+                let items = g.node_weight(node).expect("Node should exist").clone();
+                g.node_weight_mut(start).unwrap().extend(items);
 
-            g.remove_node(node).expect("Node should exist");
-            did_work = true;
+                nodes_to_remove.push(node);
+            }
+
+            // Remove merged nodes (in reverse order to preserve indices)
+            nodes_to_remove.sort();
+            for node in nodes_to_remove.into_iter().rev() {
+                g.remove_node(node);
+                did_work = true;
+            }
         }
 
         did_work