Turbopack: don't revisit nodes

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<StoreImpl, Key, KeyExtractor>
+where
+    StoreImpl: GraphStore,
+    Key: Send + Eq + Hash,
+    KeyExtractor: Send + Fn(&StoreImpl::Node) -> &Key,
+{
+    store: StoreImpl,
+    visited: FxHashSet<Key>,
+    key_extractor: KeyExtractor,
+}
+
+impl<StoreImpl, Key, KeyExtractor> SkipDuplicatesWithKey<StoreImpl, Key, KeyExtractor>
+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<StoreImpl, Key, KeyExtractor> GraphStore
+    for SkipDuplicatesWithKey<StoreImpl, Key, KeyExtractor>
+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<Self::Handle>,
+        node: GraphNode<StoreImpl::Node>,
+    ) -> 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<StoreImpl, Key, KeyExtractor> SkipDuplicatesWithKey<StoreImpl, Key, KeyExtractor>
+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<Key>) {
+        (self.store, VisitedNodes(self.visited))
+    }
+}

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<Self::Node>,
     ) -> SkipDuplicates<Self>;
+
+    fn skip_duplicates_with_key<
+        Key: Send + Eq + std::hash::Hash + Clone,
+        KeyExtractor: Send + Fn(&Self::Node) -> &Key,
+    >(
+        self,
+        key_extractor: KeyExtractor,
+    ) -> SkipDuplicatesWithKey<Self, Key, KeyExtractor>;
 }
 
 impl<Store> GraphTraversal for Store
@@ -130,6 +138,16 @@ where
     ) -> SkipDuplicates<Self> {
         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<Self, Key, KeyExtractor> {
+        SkipDuplicatesWithKey::new(self, key_extractor)
+    }
 }
 
 pub enum GraphTraversalResult<Completed, Aborted> {

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

turbopack/crates/turbopack-core/src/reference/mod.rs

@@ -290,7 +290,7 @@ type ModulesVec = Vec<ResolvedVc<Box<dyn Module>>>;
 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