From 3a591d27cee19910852108e34ad6eb657b7fea7d Mon Sep 17 00:00:00 2001 From: Alex Kirszenberg Date: Mon, 1 May 2023 10:48:38 +0200 Subject: [PATCH] More fluent GraphTraversal API (vercel/turbo#4598) ### Description This PR refactors the API of `GraphTraversal` to be more fluent, and less magical. It should make it easier to use and understand. Had this in store from trying to debug a lifetime issue in another branch. ### Testing Instructions Surface-level change. Next.js PR: https://github.com/vercel/turbo/pulls/alexkirsz --- crates/turbo-tasks/src/graph/graph_store.rs | 62 +++++++++----- .../turbo-tasks/src/graph/graph_traversal.rs | 81 ++++++++++++------- .../src/graph/non_deterministic.rs | 15 +++- .../src/graph/reverse_topological.rs | 20 +++-- crates/turbopack-core/src/changed.rs | 21 +++-- .../src/chunk/available_assets.rs | 15 ++-- crates/turbopack-core/src/chunk/mod.rs | 7 +- crates/turbopack-dev/src/chunking_context.rs | 29 +++---- crates/turbopack-node/src/lib.rs | 15 ++-- 9 files changed, 160 insertions(+), 105 deletions(-) diff --git a/crates/turbo-tasks/src/graph/graph_store.rs b/crates/turbo-tasks/src/graph/graph_store.rs index 918dc643bae81..77e590ffe0749 100644 --- a/crates/turbo-tasks/src/graph/graph_store.rs +++ b/crates/turbo-tasks/src/graph/graph_store.rs @@ -2,7 +2,8 @@ use std::collections::HashSet; /// A graph store is a data structure that will be built up during a graph /// traversal. It is used to store the results of the traversal. -pub trait GraphStore: Default { +pub trait GraphStore { + type Node; type Handle: Clone; // TODO(alexkirsz) An `entry(from_handle) -> Entry` API would be more @@ -13,48 +14,71 @@ pub trait GraphStore: Default { fn insert( &mut self, from_handle: Option, - node: Node, - ) -> Option<(Self::Handle, &Node)>; + node: GraphNode, + ) -> Option<(Self::Handle, &Self::Node)>; +} + +/// Utility type to ensure that GraphStore::insert can only ever be called from +/// within this module, as a GraphNode can't be constructed outside of it. +#[derive(Clone, Copy, Eq, PartialEq, Debug, Hash, Ord, PartialOrd)] +pub struct GraphNode(pub(super) Node); + +impl GraphNode { + /// Consumes this `GraphNode` and returns the underlying node. + pub fn into_node(self) -> Node { + self.0 + } + + /// Returns a reference the underlying node. + pub fn node(&self) -> &Node { + &self.0 + } + + /// Returns a mutable reference the underlying node. + pub fn node_mut(&mut self) -> &mut Node { + &mut self.0 + } } /// A [`GraphStore`] wrapper that skips nodes that have already been /// visited. 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 SkipDuplicates +pub struct SkipDuplicates where - StoreImpl: GraphStore, + StoreImpl: GraphStore, { store: StoreImpl, - visited: HashSet, + visited: HashSet, } -impl Default for SkipDuplicates +impl SkipDuplicates where - StoreImpl: GraphStore, + StoreImpl: GraphStore, { - fn default() -> Self { + pub fn new(store: StoreImpl) -> Self { Self { - store: Default::default(), + store: store, visited: Default::default(), } } } -impl GraphStore for SkipDuplicates +impl GraphStore for SkipDuplicates where - StoreImpl: GraphStore, - Node: Eq + std::hash::Hash + Clone, + StoreImpl: GraphStore, + StoreImpl::Node: Eq + std::hash::Hash + Clone, { + type Node = StoreImpl::Node; type Handle = StoreImpl::Handle; fn insert( &mut self, from_handle: Option, - node: Node, - ) -> Option<(Self::Handle, &Node)> { - if !self.visited.contains(&node) { - self.visited.insert(node.clone()); + node: GraphNode, + ) -> Option<(Self::Handle, &StoreImpl::Node)> { + if !self.visited.contains(node.node()) { + self.visited.insert(node.node().clone()); self.store.insert(from_handle, node) } else { // Always insert the node into the store, even if we've already @@ -66,9 +90,9 @@ where } } -impl SkipDuplicates +impl SkipDuplicates where - StoreImpl: GraphStore, + StoreImpl: GraphStore, { /// Consumes the wrapper and returns the underlying store. pub fn into_inner(self) -> StoreImpl { diff --git a/crates/turbo-tasks/src/graph/graph_traversal.rs b/crates/turbo-tasks/src/graph/graph_traversal.rs index 9d8dafa18cac5..5a77b867a416c 100644 --- a/crates/turbo-tasks/src/graph/graph_traversal.rs +++ b/crates/turbo-tasks/src/graph/graph_traversal.rs @@ -3,39 +3,54 @@ use std::{future::Future, pin::Pin}; use anyhow::Result; use futures::{stream::FuturesUnordered, Stream}; -use super::{graph_store::GraphStore, with_future::With, Visit, VisitControlFlow}; +use super::{ + graph_store::{GraphNode, GraphStore}, + with_future::With, + SkipDuplicates, Visit, VisitControlFlow, +}; /// [`GraphTraversal`] is a utility type that can be used to traverse a graph of /// nodes, where each node can have a variable number of outgoing edges. The /// traversal is done in parallel, and the order of the nodes in the traversal /// result is determined by the [`GraphStore`] parameter. -pub struct GraphTraversal { - _store: std::marker::PhantomData, +pub trait GraphTraversal: GraphStore + Sized { + fn visit( + self, + root_edges: RootEdgesIt, + visit: VisitImpl, + ) -> GraphTraversalFuture + where + VisitImpl: Visit, + RootEdgesIt: IntoIterator; + + fn skip_duplicates(self) -> SkipDuplicates; } -impl GraphTraversal { +impl GraphTraversal for Store +where + Store: GraphStore, +{ /// Visits the graph starting from the given `roots`, and returns a future /// that will resolve to the traversal result. - pub fn visit( + fn visit( + mut self, root_edges: RootEdgesIt, mut visit: VisitImpl, - ) -> GraphTraversalFuture + ) -> GraphTraversalFuture where - Store: GraphStore, - VisitImpl: Visit, + VisitImpl: Visit, RootEdgesIt: IntoIterator, { - let mut store = Store::default(); let futures = FuturesUnordered::new(); for edge in root_edges { match visit.visit(edge) { VisitControlFlow::Continue(node) => { - if let Some((parent_handle, node_ref)) = store.insert(None, node) { + if let Some((parent_handle, node_ref)) = self.insert(None, GraphNode(node)) { futures.push(With::new(visit.edges(node_ref), parent_handle)); } } VisitControlFlow::Skip(node) => { - store.insert(None, node); + self.insert(None, GraphNode(node)); } VisitControlFlow::Abort(abort) => { return GraphTraversalFuture { @@ -46,42 +61,46 @@ impl GraphTraversal { } GraphTraversalFuture { state: GraphTraversalState::Running(GraphTraversalRunningState { - store, + store: self, futures, visit, }), } } + + fn skip_duplicates(self) -> SkipDuplicates { + SkipDuplicates::new(self) + } } /// A future that resolves to a [`GraphStore`] containing the result of a graph /// traversal. -pub struct GraphTraversalFuture +pub struct GraphTraversalFuture where - Store: GraphStore, - VisitImpl: Visit, + Store: GraphStore, + VisitImpl: Visit, { - state: GraphTraversalState, + state: GraphTraversalState, } #[derive(Default)] -enum GraphTraversalState +enum GraphTraversalState where - Store: GraphStore, - VisitImpl: Visit, + Store: GraphStore, + VisitImpl: Visit, { #[default] Completed, Aborted { abort: Abort, }, - Running(GraphTraversalRunningState), + Running(GraphTraversalRunningState), } -struct GraphTraversalRunningState +struct GraphTraversalRunningState where - Store: GraphStore, - VisitImpl: Visit, + Store: GraphStore, + VisitImpl: Visit, { store: Store, futures: FuturesUnordered>, @@ -102,11 +121,10 @@ impl GraphTraversalResult { } } -impl Future - for GraphTraversalFuture +impl Future for GraphTraversalFuture where - Store: GraphStore, - VisitImpl: Visit, + Store: GraphStore, + VisitImpl: Visit, { type Output = GraphTraversalResult, Abort>; @@ -132,8 +150,9 @@ where for edge in edges { match running.visit.visit(edge) { VisitControlFlow::Continue(node) => { - if let Some((node_handle, node_ref)) = - running.store.insert(Some(parent_handle.clone()), node) + if let Some((node_handle, node_ref)) = running + .store + .insert(Some(parent_handle.clone()), GraphNode(node)) { running.futures.push(With::new( running.visit.edges(node_ref), @@ -142,7 +161,9 @@ where } } VisitControlFlow::Skip(node) => { - running.store.insert(Some(parent_handle.clone()), node); + running + .store + .insert(Some(parent_handle.clone()), GraphNode(node)); } VisitControlFlow::Abort(abort) => { break 'outer ( diff --git a/crates/turbo-tasks/src/graph/non_deterministic.rs b/crates/turbo-tasks/src/graph/non_deterministic.rs index 1316c23b5b903..ef64f514176f4 100644 --- a/crates/turbo-tasks/src/graph/non_deterministic.rs +++ b/crates/turbo-tasks/src/graph/non_deterministic.rs @@ -1,4 +1,4 @@ -use super::graph_store::GraphStore; +use super::graph_store::{GraphNode, GraphStore}; /// A graph traversal that does not guarantee any particular order, and may not /// return the same order every time it is run. @@ -8,19 +8,26 @@ pub struct NonDeterministic { impl Default for NonDeterministic { fn default() -> Self { + Self::new() + } +} + +impl NonDeterministic { + pub fn new() -> Self { Self { output: Vec::new() } } } -impl GraphStore for NonDeterministic { +impl GraphStore for NonDeterministic { + type Node = T; type Handle = (); fn insert( &mut self, _from_handle: Option, - node: T, + node: GraphNode, ) -> Option<(Self::Handle, &T)> { - self.output.push(node); + self.output.push(node.into_node()); Some(((), self.output.last().unwrap())) } } diff --git a/crates/turbo-tasks/src/graph/reverse_topological.rs b/crates/turbo-tasks/src/graph/reverse_topological.rs index 4996641dd955b..a2fb6af73bb00 100644 --- a/crates/turbo-tasks/src/graph/reverse_topological.rs +++ b/crates/turbo-tasks/src/graph/reverse_topological.rs @@ -1,6 +1,6 @@ use std::collections::{HashMap, HashSet}; -use super::graph_store::GraphStore; +use super::graph_store::{GraphNode, GraphStore}; /// A graph traversal that returns nodes in reverse topological order. pub struct ReverseTopological @@ -16,6 +16,15 @@ where T: Eq + std::hash::Hash + Clone, { fn default() -> Self { + Self::new() + } +} + +impl ReverseTopological +where + T: Eq + std::hash::Hash + Clone, +{ + pub fn new() -> Self { Self { adjacency_map: HashMap::new(), roots: Vec::new(), @@ -23,13 +32,14 @@ where } } -impl GraphStore for ReverseTopological +impl GraphStore for ReverseTopological where T: Eq + std::hash::Hash + Clone, { + type Node = T; type Handle = T; - fn insert(&mut self, from_handle: Option, node: T) -> Option<(Self::Handle, &T)> { + fn insert(&mut self, from_handle: Option, node: GraphNode) -> Option<(Self::Handle, &T)> { let vec = if let Some(from_handle) = from_handle { self.adjacency_map .entry(from_handle) @@ -38,8 +48,8 @@ where &mut self.roots }; - vec.push(node.clone()); - Some((node, vec.last().unwrap())) + vec.push(node.node().clone()); + Some((node.into_node(), vec.last().unwrap())) } } diff --git a/crates/turbopack-core/src/changed.rs b/crates/turbopack-core/src/changed.rs index c8a6fdf84dd7a..e640f35de2631 100644 --- a/crates/turbopack-core/src/changed.rs +++ b/crates/turbopack-core/src/changed.rs @@ -1,6 +1,6 @@ use anyhow::Result; use turbo_tasks::{ - graph::{GraphTraversal, NonDeterministic, SkipDuplicates}, + graph::{GraphTraversal, NonDeterministic}, CompletionVc, CompletionsVc, }; @@ -20,16 +20,15 @@ async fn get_referenced_assets(parent: AssetVc) -> Result Result { - let completions = GraphTraversal::, _>>::visit( - [root], - get_referenced_assets, - ) - .await - .completed()? - .into_inner() - .into_iter() - .map(content_changed) - .collect(); + let completions = NonDeterministic::new() + .skip_duplicates() + .visit([root], get_referenced_assets) + .await + .completed()? + .into_inner() + .into_iter() + .map(content_changed) + .collect(); Ok(CompletionsVc::cell(completions).completed()) } diff --git a/crates/turbopack-core/src/chunk/available_assets.rs b/crates/turbopack-core/src/chunk/available_assets.rs index dc77db35f66d6..624ddbe17a997 100644 --- a/crates/turbopack-core/src/chunk/available_assets.rs +++ b/crates/turbopack-core/src/chunk/available_assets.rs @@ -2,7 +2,7 @@ use std::iter::once; use anyhow::Result; use turbo_tasks::{ - graph::{GraphTraversal, ReverseTopological, SkipDuplicates}, + graph::{GraphTraversal, ReverseTopological}, primitives::{BoolVc, U64Vc}, TryJoinIterExt, ValueToString, }; @@ -82,9 +82,9 @@ impl AvailableAssetsVc { #[turbo_tasks::function] async fn chunkable_assets_set(root: AssetVc) -> Result { - let assets = GraphTraversal::, _>>::visit( - once(root), - |&asset: &AssetVc| async move { + let assets = ReverseTopological::new() + .skip_duplicates() + .visit(once(root), |&asset: &AssetVc| async move { let mut results = Vec::new(); for reference in asset.references().await?.iter() { if let Some(chunkable) = ChunkableAssetReferenceVc::resolve_from(reference).await? { @@ -108,9 +108,8 @@ async fn chunkable_assets_set(root: AssetVc) -> Result { } } Ok(results) - }, - ) - .await - .completed()?; + }) + .await + .completed()?; Ok(AssetsSetVc::cell(assets.into_inner().into_iter().collect())) } diff --git a/crates/turbopack-core/src/chunk/mod.rs b/crates/turbopack-core/src/chunk/mod.rs index 16b60fda8e6ea..723cf0a7edcc7 100644 --- a/crates/turbopack-core/src/chunk/mod.rs +++ b/crates/turbopack-core/src/chunk/mod.rs @@ -562,10 +562,9 @@ where _phantom: PhantomData, }; - let GraphTraversalResult::Completed(traversal_result) = - GraphTraversal::>::visit(root_edges, visit).await else { - return Ok(None); - }; + let GraphTraversalResult::Completed(traversal_result) = ReverseTopological::new().visit(root_edges, visit).await else { + return Ok(None); + }; let graph_nodes: Vec<_> = traversal_result?.into_iter().collect(); diff --git a/crates/turbopack-dev/src/chunking_context.rs b/crates/turbopack-dev/src/chunking_context.rs index 2596be59c7986..acdc3e5e2c9dd 100644 --- a/crates/turbopack-dev/src/chunking_context.rs +++ b/crates/turbopack-dev/src/chunking_context.rs @@ -3,7 +3,7 @@ use std::fmt::Write; use anyhow::Result; use indexmap::IndexSet; use turbo_tasks::{ - graph::{GraphTraversal, ReverseTopological, SkipDuplicates}, + graph::{GraphTraversal, ReverseTopological}, primitives::{BoolVc, StringVc}, TryJoinIterExt, Value, ValueToString, }; @@ -441,24 +441,21 @@ async fn get_parallel_chunks(entries: I) -> Result, { - Ok( - GraphTraversal::, _>>::visit( - entries, - |chunk: ChunkVc| async move { - Ok(chunk - .parallel_chunks() - .await? - .iter() - .copied() - .collect::>() - .into_iter()) - }, - ) + Ok(ReverseTopological::new() + .skip_duplicates() + .visit(entries, |chunk: ChunkVc| async move { + Ok(chunk + .parallel_chunks() + .await? + .iter() + .copied() + .collect::>() + .into_iter()) + }) .await .completed()? .into_inner() - .into_iter(), - ) + .into_iter()) } async fn get_optimized_chunks(chunks: I) -> Result diff --git a/crates/turbopack-node/src/lib.rs b/crates/turbopack-node/src/lib.rs index df948f7cff9a3..eb3e424416085 100644 --- a/crates/turbopack-node/src/lib.rs +++ b/crates/turbopack-node/src/lib.rs @@ -10,7 +10,7 @@ pub use node_entry::{ NodeEntry, NodeEntryVc, NodeRenderingEntriesVc, NodeRenderingEntry, NodeRenderingEntryVc, }; use turbo_tasks::{ - graph::{GraphTraversal, ReverseTopological, SkipDuplicates}, + graph::{GraphTraversal, ReverseTopological}, CompletionVc, CompletionsVc, TryJoinIterExt, ValueToString, }; use turbo_tasks_env::{ProcessEnv, ProcessEnvVc}; @@ -170,13 +170,12 @@ async fn separate_assets( .await }; - let graph = GraphTraversal::, _>>::visit( - once(Type::Internal(intermediate_asset)), - get_asset_children, - ) - .await - .completed()? - .into_inner(); + let graph = ReverseTopological::new() + .skip_duplicates() + .visit(once(Type::Internal(intermediate_asset)), get_asset_children) + .await + .completed()? + .into_inner(); let mut internal_assets = IndexSet::new(); let mut external_asset_entrypoints = IndexSet::new();