diff --git a/turbopack/crates/turbo-tasks/src/graph/graph_store.rs b/turbopack/crates/turbo-tasks/src/graph/graph_store.rs index f6ec1b50c7fda..cf1385b2e04b2 100644 --- a/turbopack/crates/turbo-tasks/src/graph/graph_store.rs +++ b/turbopack/crates/turbo-tasks/src/graph/graph_store.rs @@ -114,3 +114,83 @@ where (self.store, VisitedNodes(self.visited)) } } + +/// A [`GraphStore`] wrapper that skips nodes that have already been +/// visited, based on a key extracted from the node. +/// +/// This is necessary to avoid repeated work when traversing non-tree +/// graphs (i.e. where a node can have more than one incoming edge). +#[derive(Debug)] +pub struct SkipDuplicatesWithKey +where + StoreImpl: GraphStore, + Key: Send + Eq + Hash, + KeyExtractor: Send + Fn(&StoreImpl::Node) -> &Key, +{ + store: StoreImpl, + visited: FxHashSet, + key_extractor: KeyExtractor, +} + +impl SkipDuplicatesWithKey +where + StoreImpl: GraphStore, + Key: Send + Eq + std::hash::Hash + Clone, + KeyExtractor: Send + Fn(&StoreImpl::Node) -> &Key, +{ + pub fn new(store: StoreImpl, key_extractor: KeyExtractor) -> Self { + Self { + store, + visited: FxHashSet::default(), + key_extractor, + } + } +} + +impl GraphStore + for SkipDuplicatesWithKey +where + StoreImpl: GraphStore, + StoreImpl::Node: Eq + std::hash::Hash + Clone, + Key: Send + Eq + std::hash::Hash + Clone, + KeyExtractor: Send + Fn(&StoreImpl::Node) -> &Key, +{ + type Node = StoreImpl::Node; + type Handle = StoreImpl::Handle; + + fn insert( + &mut self, + from_handle: Option, + node: GraphNode, + ) -> Option<(Self::Handle, &StoreImpl::Node)> { + let key = (self.key_extractor)(node.node()); + + if !self.visited.contains(key) { + self.visited.insert(key.clone()); + self.store.insert(from_handle, node) + } else { + // Always insert the node into the store, even if we've already + // visited it. This is necessary to ensure that the store sees all + // edges. + self.store.insert(from_handle, node); + None + } + } +} + +impl SkipDuplicatesWithKey +where + StoreImpl: GraphStore, + Key: Send + Eq + std::hash::Hash + Clone, + KeyExtractor: Send + Fn(&StoreImpl::Node) -> &Key, +{ + /// Consumes the wrapper and returns the underlying store. + pub fn into_inner(self) -> StoreImpl { + self.store + } + + /// Consumes the wrapper and returns the underlying store along with the visited nodes. + pub fn into_inner_with_visited(self) -> (StoreImpl, VisitedNodes) { + (self.store, VisitedNodes(self.visited)) + } +} diff --git a/turbopack/crates/turbo-tasks/src/graph/graph_traversal.rs b/turbopack/crates/turbo-tasks/src/graph/graph_traversal.rs index 4ed18d480a861..0afad0eb2c687 100644 --- a/turbopack/crates/turbo-tasks/src/graph/graph_traversal.rs +++ b/turbopack/crates/turbo-tasks/src/graph/graph_traversal.rs @@ -6,7 +6,7 @@ use rustc_hash::FxHashSet; use super::{ SkipDuplicates, Visit, VisitControlFlow, - graph_store::{GraphNode, GraphStore}, + graph_store::{GraphNode, GraphStore, SkipDuplicatesWithKey}, with_future::With, }; @@ -35,6 +35,14 @@ pub trait GraphTraversal: GraphStore + Sized { self, visited: VisitedNodes, ) -> SkipDuplicates; + + fn skip_duplicates_with_key< + Key: Send + Eq + std::hash::Hash + Clone, + KeyExtractor: Send + Fn(&Self::Node) -> &Key, + >( + self, + key_extractor: KeyExtractor, + ) -> SkipDuplicatesWithKey; } impl GraphTraversal for Store @@ -130,6 +138,16 @@ where ) -> SkipDuplicates { SkipDuplicates::new_with_visited_nodes(self, visited.0) } + + fn skip_duplicates_with_key< + Key: Send + Eq + std::hash::Hash + Clone, + KeyExtractor: Send + Fn(&Self::Node) -> &Key, + >( + self, + key_extractor: KeyExtractor, + ) -> SkipDuplicatesWithKey { + SkipDuplicatesWithKey::new(self, key_extractor) + } } pub enum GraphTraversalResult { diff --git a/turbopack/crates/turbopack-core/src/module_graph/mod.rs b/turbopack/crates/turbopack-core/src/module_graph/mod.rs index 623baefe465b8..c6392e541d6e1 100644 --- a/turbopack/crates/turbopack-core/src/module_graph/mod.rs +++ b/turbopack/crates/turbopack-core/src/module_graph/mod.rs @@ -233,7 +233,7 @@ impl SingleModuleGraph { .await?; let (children_nodes_iter, visited_nodes) = AdjacencyMap::new() - .skip_duplicates() + .skip_duplicates_with_key(|node: &(SingleModuleGraphBuilderNode, ExportUsage)| &node.0) .visit( root_edges, SingleModuleGraphBuilder { @@ -1669,6 +1669,8 @@ impl Visit<(SingleModuleGraphBuilderNode, ExportUsage)> for SingleModuleGraphBui fn edges( &mut self, + // The `skip_duplicates_with_key()` above ensures only a single `edges()` call per module + // (and not per `(module, export)` pair), so the export must not be read here! (node, _): &(SingleModuleGraphBuilderNode, ExportUsage), ) -> Self::EdgesFuture { // Destructure beforehand to not have to clone the whole node when entering the async block diff --git a/turbopack/crates/turbopack-core/src/reference/mod.rs b/turbopack/crates/turbopack-core/src/reference/mod.rs index 0d75b1171dafd..b0f0df2deafaa 100644 --- a/turbopack/crates/turbopack-core/src/reference/mod.rs +++ b/turbopack/crates/turbopack-core/src/reference/mod.rs @@ -290,7 +290,7 @@ type ModulesVec = Vec>>; pub struct ModulesWithRefData(Vec<(ChunkingType, ExportUsage, ModulesVec)>); /// Aggregates all primary [Module]s referenced by an [Module] via [ChunkableModuleReference]s. -/// This does not include transitively references [Module]s, only includes +/// This does not include transitively referenced [Module]s, only includes /// primary [Module]s referenced. /// /// [Module]: crate::module::Module