perf: optimize group_chunks, speed up group_chunks by 500%

crates/mako/Cargo.toml

@@ -91,6 +91,7 @@ indicatif             = "0.17.8"
 md5                   = "0.7.0"
 mdxjs                 = "0.1.14"
 mime_guess            = "2.0.4"
+nohash-hasher         = "0.2.0"
 notify                = { version = "6.1.1", default-features = false, features = ["macos_kqueue"] }
 notify-debouncer-full = { version = "0.3.1", default-features = false }
 parking_lot           = { version = "0.12", features = ["nightly"] }

crates/mako/src/generate/group_chunk.rs

@@ -388,16 +388,11 @@ impl Compiler {
 
         let chunk_id = entry_module_id.generate(&self.context);
         let mut chunk = Chunk::new(chunk_id.into(), chunk_type.clone());
-        let mut visited_modules: Vec<&ModuleId> = vec![entry_module_id];
+        let mut normal_deps = HashSet::<&ModuleId>::from_iter([entry_module_id]);
 
         let module_graph = self.context.module_graph.read().unwrap();
 
         visit_modules(vec![entry_module_id.clone()], None, |head| {
-            let parent_index = visited_modules
-                .iter()
-                .position(|m| m.id == head.id)
-                .unwrap_or(0);
-            let mut normal_deps = vec![];
             let mut next_module_ids = vec![];
 
             for (dep_module_id, dep) in module_graph.get_dependencies(head) {
@@ -418,21 +413,22 @@ impl Compiler {
                     {
                         next_module_ids.push(dep_module_id.clone());
                         // collect normal deps for current head
-                        normal_deps.push(dep_module_id);
+                        normal_deps.insert(dep_module_id);
                     }
                     _ => {}
                 }
             }
 
-            // insert normal deps before head, so that we can keep the dfs order
-            visited_modules.splice(parent_index..parent_index, normal_deps);
-
             next_module_ids
         });
 
+        let mut deps = module_graph.get_dependencies_by_link_back_dfs(entry_module_id);
+
         // add modules to chunk as dfs order
-        for module_id in visited_modules {
-            chunk.add_module(module_id.clone());
+        while let Some(module_id) = deps.pop() {
+            if let Some(module_id) = normal_deps.take(module_id) {
+                chunk.add_module(module_id.clone());
+            }
         }
 
         (chunk, dynamic_entries, worker_entries)

crates/mako/src/generate/optimize_chunk.rs

@@ -125,7 +125,7 @@ impl Compiler {
         let async_chunk_root_modules = chunks
             .iter()
             .filter_map(|chunk| match chunk.chunk_type {
-                ChunkType::Async => chunk.modules.iter().last(),
+                ChunkType::Async => chunk.modules.back(),
                 _ => None,
             })
             .collect::<Vec<_>>();

crates/mako/src/module_graph/link_back_dfs_visit.rs

@@ -0,0 +1,75 @@
+use std::fmt::Debug;
+
+use fixedbitset::FixedBitSet;
+use hashlink::{linked_hash_map, LinkedHashMap};
+use nohash_hasher::BuildNoHashHasher;
+use petgraph::csr::IndexType;
+use petgraph::graph::{EdgeIndex, NodeIndex};
+use petgraph::stable_graph::StableDiGraph;
+use petgraph::visit::{VisitMap, Visitable};
+use petgraph::Direction;
+
+/// This visitor is for css ordering, see test cases of css-merge-in-js and css-merge-in-css
+#[derive(Clone, Debug)]
+pub struct LinkBackDfs {
+    pub stack: LinkedHashMap<usize, (), BuildNoHashHasher<usize>>,
+    pub discovered: FixedBitSet,
+    pub finished: FixedBitSet,
+}
+
+impl LinkBackDfs {
+    pub fn new<N, E>(graph: &StableDiGraph<N, E>, start: NodeIndex) -> Self {
+        let mut dfs = Self::empty(graph);
+        dfs.move_to(start);
+        dfs
+    }
+
+    pub fn empty<N, E>(graph: &StableDiGraph<N, E>) -> Self {
+        LinkBackDfs {
+            stack: LinkedHashMap::with_hasher(BuildNoHashHasher::default()),
+            discovered: graph.visit_map(),
+            finished: graph.visit_map(),
+        }
+    }
+
+    pub fn move_to(&mut self, start: NodeIndex) {
+        self.stack.clear();
+        self.stack.insert(start.index(), ());
+    }
+
+    pub fn next<N, E, F>(
+        &mut self,
+        graph: &StableDiGraph<N, E>,
+        execlude_edge: F,
+    ) -> Option<NodeIndex>
+    where
+        F: Fn(EdgeIndex) -> bool,
+    {
+        while let Some((&nx, _)) = self.stack.back() {
+            if self.discovered.visit(nx) {
+                let mut nxs = Vec::<NodeIndex>::new();
+
+                let mut walker = graph
+                    .neighbors_directed(NodeIndex::new(nx), Direction::Outgoing)
+                    .detach();
+                while let Some((cex, cnx)) = walker.next(graph) {
+                    if !execlude_edge(cex) {
+                        nxs.push(cnx);
+                    }
+                }
+
+                while let Some(cnx) = nxs.pop() {
+                    if let linked_hash_map::Entry::Occupied(entry) = self.stack.entry(nx.index()) {
+                        entry.cursor_mut().insert_before(cnx.index(), ());
+                    }
+                }
+            } else {
+                self.stack.pop_back();
+                if self.finished.visit(nx) {
+                    return Some(NodeIndex::new(nx));
+                }
+            }
+        }
+        None
+    }
+}

crates/mako/src/module_graph.rs → crates/mako/src/module_graph/mod.rs

@@ -2,15 +2,20 @@ use std::collections::{HashMap, HashSet};
 use std::fmt;
 
 use fixedbitset::FixedBitSet;
+use hashlink::LinkedHashSet;
+use nohash_hasher::BuildNoHashHasher;
 use petgraph::graph::{DefaultIx, NodeIndex};
-use petgraph::prelude::{Dfs, EdgeRef};
+use petgraph::prelude::Dfs;
 use petgraph::stable_graph::{StableDiGraph, WalkNeighbors};
-use petgraph::visit::IntoEdgeReferences;
+use petgraph::visit::{EdgeRef, IntoEdgeReferences};
 use petgraph::Direction;
 use tracing::debug;
 
 use crate::module::{Dependencies, Dependency, Module, ModuleId, ResolveType};
 
+mod link_back_dfs_visit;
+use link_back_dfs_visit::LinkBackDfs;
+
 #[derive(Debug)]
 pub struct ModuleGraph {
     id_index_map: HashMap<ModuleId, NodeIndex<DefaultIx>>,
@@ -204,6 +209,28 @@ impl ModuleGraph {
         deps
     }
 
+    pub fn get_dependencies_by_link_back_dfs(&self, module_id: &ModuleId) -> Vec<&ModuleId> {
+        let start = self.id_index_map.get(module_id).unwrap();
+        let mut dfs_post_order = LinkBackDfs::new(&self.graph, *start);
+        let mut deps = LinkedHashSet::<usize, BuildNoHashHasher<usize>>::with_hasher(
+            BuildNoHashHasher::default(),
+        );
+        while let Some(nth) = dfs_post_order.next(&self.graph, |ex| {
+            let dependencies = self.graph.edge_weight(ex).unwrap();
+            dependencies.iter().any(|d| {
+                matches!(
+                    d.resolve_type,
+                    ResolveType::DynamicImport | ResolveType::Worker
+                )
+            })
+        }) {
+            deps.insert(nth.index());
+        }
+        deps.iter()
+            .map(|e| &self.graph[NodeIndex::new(*e)].id)
+            .collect()
+    }
+
     pub fn get_dependents(&self, module_id: &ModuleId) -> Vec<(&ModuleId, &Dependency)> {
         let mut edges = self.get_edges(module_id, Direction::Incoming);
         let mut deps: Vec<(&ModuleId, &Dependency)> = vec![];

crates/mako/src/stats.rs

@@ -126,11 +126,9 @@ impl Compiler {
                     ChunkType::Sync => chunk_graph
                         .dependents_chunk(&chunk.id)
                         .iter()
-                        .filter_map(|chunk_id| {
-                            chunk_graph.chunk(chunk_id).unwrap().modules.iter().last()
-                        })
+                        .filter_map(|chunk_id| chunk_graph.chunk(chunk_id).unwrap().modules.back())
                         .collect::<Vec<_>>(),
-                    _ => vec![chunk.modules.iter().last().unwrap()],
+                    _ => vec![chunk.modules.back().unwrap()],
                 };
                 let mut origins_set = IndexMap::new();
                 for origin_chunk_module in origin_chunk_modules {