diff --git a/crates/turbo-tasks/src/graph/control_flow.rs b/crates/turbo-tasks/src/graph/control_flow.rs new file mode 100644 index 0000000000000..7402d204000bd --- /dev/null +++ b/crates/turbo-tasks/src/graph/control_flow.rs @@ -0,0 +1,35 @@ +/// The control flow of visiting an edge during a graph traversal. +pub enum VisitControlFlow { + /// The traversal should continue on the outgoing edges of the given node. + Continue(Node), + /// The traversal should skip visiting the edges the given node. + Skip(Node), + /// The traversal should abort and return immediately. + Abort(Abort), +} + +impl VisitControlFlow { + /// Map the continue and skip values of this control flow. + pub fn map_node(self, mut map: Map) -> VisitControlFlow + where + Map: FnMut(Node) -> Mapped, + { + match self { + VisitControlFlow::Continue(node) => VisitControlFlow::Continue(map(node)), + VisitControlFlow::Skip(node) => VisitControlFlow::Skip(map(node)), + VisitControlFlow::Abort(abort) => VisitControlFlow::Abort(abort), + } + } + + /// Map the abort value of this control flow. + pub fn map_abort(self, mut map: Map) -> VisitControlFlow + where + Map: FnMut(Abort) -> Mapped, + { + match self { + VisitControlFlow::Continue(node) => VisitControlFlow::Continue(node), + VisitControlFlow::Skip(node) => VisitControlFlow::Skip(node), + VisitControlFlow::Abort(abort) => VisitControlFlow::Abort(map(abort)), + } + } +} diff --git a/crates/turbo-tasks/src/graph/get_children.rs b/crates/turbo-tasks/src/graph/get_children.rs deleted file mode 100644 index ee1a9c5bc1cbc..0000000000000 --- a/crates/turbo-tasks/src/graph/get_children.rs +++ /dev/null @@ -1,226 +0,0 @@ -use std::{collections::HashSet, future::Future}; - -use anyhow::Result; - -use super::GraphTraversalControlFlow; - -/// A trait that allows a graph traversal to get the children of a node. -pub trait Visit { - type Children: IntoIterator; - type MapChildren: IntoIterator; - type Future: Future>; - - fn get_children(&mut self, item: &T) -> GraphTraversalControlFlow; - fn map_children( - &mut self, - children: Self::Children, - ) -> GraphTraversalControlFlow; -} - -// The different `Impl*` here are necessary in order to avoid the `Conflicting -// implementations of trait` error when implementing `GetChildren` on different -// kinds of `FnMut`. -// See https://users.rust-lang.org/t/conflicting-implementation-when-implementing-traits-for-fn/53359/3 - -pub struct ImplRef; - -impl Visit for C -where - C: FnMut(&T) -> F, - F: Future>, - CI: IntoIterator, -{ - type Children = CI; - type MapChildren = Self::Children; - type Future = F; - - fn get_children(&mut self, item: &T) -> GraphTraversalControlFlow { - GraphTraversalControlFlow::Continue((self)(item)) - } - - fn map_children( - &mut self, - children: Self::Children, - ) -> GraphTraversalControlFlow { - GraphTraversalControlFlow::Continue(children) - } -} - -pub struct ImplRefOption; - -impl Visit for C -where - C: FnMut(&T) -> Option, - F: Future>, - CI: IntoIterator, -{ - type Children = CI; - type MapChildren = Self::Children; - type Future = F; - - fn get_children(&mut self, item: &T) -> GraphTraversalControlFlow { - match (self)(item) { - Some(future) => GraphTraversalControlFlow::Continue(future), - None => GraphTraversalControlFlow::Skip, - } - } - - fn map_children( - &mut self, - children: Self::Children, - ) -> GraphTraversalControlFlow { - GraphTraversalControlFlow::Continue(children) - } -} - -pub struct ImplRefControlFlow; - -impl Visit for C -where - T: Clone, - C: FnMut(&T) -> GraphTraversalControlFlow, - F: Future>, - CI: IntoIterator, -{ - type Children = CI; - type MapChildren = Self::Children; - type Future = F; - - fn get_children(&mut self, item: &T) -> GraphTraversalControlFlow { - (self)(item) - } - - fn map_children( - &mut self, - children: Self::Children, - ) -> GraphTraversalControlFlow { - GraphTraversalControlFlow::Continue(children) - } -} - -pub struct ImplValue; - -impl Visit for C -where - T: Clone, - C: FnMut(T) -> F, - F: Future>, - CI: IntoIterator, -{ - type Children = CI; - type MapChildren = Self::Children; - type Future = F; - - fn get_children(&mut self, item: &T) -> GraphTraversalControlFlow { - GraphTraversalControlFlow::Continue((self)(item.clone())) - } - - fn map_children( - &mut self, - children: Self::Children, - ) -> GraphTraversalControlFlow { - GraphTraversalControlFlow::Continue(children) - } -} - -pub struct ImplValueOption; - -impl Visit for C -where - T: Clone, - C: FnMut(T) -> Option, - F: Future>, - CI: IntoIterator, -{ - type Children = CI; - type MapChildren = Self::Children; - type Future = F; - - fn get_children(&mut self, item: &T) -> GraphTraversalControlFlow { - match (self)(item.clone()) { - Some(future) => GraphTraversalControlFlow::Continue(future), - None => GraphTraversalControlFlow::Skip, - } - } - - fn map_children( - &mut self, - children: Self::Children, - ) -> GraphTraversalControlFlow { - GraphTraversalControlFlow::Continue(children) - } -} - -pub struct ImplValueControlFlow; - -impl Visit for C -where - T: Clone, - C: FnMut(T) -> GraphTraversalControlFlow, - F: Future>, - CI: IntoIterator, -{ - type Children = CI; - type MapChildren = Self::Children; - type Future = F; - - fn get_children(&mut self, item: &T) -> GraphTraversalControlFlow { - (self)(item.clone()) - } - - fn map_children( - &mut self, - children: Self::Children, - ) -> GraphTraversalControlFlow { - GraphTraversalControlFlow::Continue(children) - } -} - -/// A [`GetChildren`] implementation that skips nodes that have already been -/// visited. This is necessary to avoid repeated work when traversing non-tree -/// graphs (i.e. where a child can have more than one parent). -#[derive(Debug)] -pub struct SkipDuplicates { - visit: C, - visited: HashSet, - _a: std::marker::PhantomData, - _impl: std::marker::PhantomData, -} - -impl SkipDuplicates { - /// Create a new [`SkipDuplicates`] that wraps the given [`GetChildren`]. - pub fn new(visit: C) -> Self { - Self { - visit, - visited: HashSet::new(), - _a: std::marker::PhantomData, - _impl: std::marker::PhantomData, - } - } -} - -impl Visit for SkipDuplicates -where - T: Eq + std::hash::Hash + Clone, - C: Visit, -{ - type Children = C::Children; - type MapChildren = C::MapChildren; - type Future = C::Future; - - fn get_children(&mut self, item: &T) -> GraphTraversalControlFlow { - if !self.visited.contains(item) { - self.visited.insert(item.clone()); - self.visit.get_children(item) - } else { - GraphTraversalControlFlow::Skip - } - } - - fn map_children( - &mut self, - children: Self::Children, - ) -> GraphTraversalControlFlow { - self.visit.map_children(children) - } -} diff --git a/crates/turbo-tasks/src/graph/graph_store.rs b/crates/turbo-tasks/src/graph/graph_store.rs index 7fda8a5b3f6d6..918dc643bae81 100644 --- a/crates/turbo-tasks/src/graph/graph_store.rs +++ b/crates/turbo-tasks/src/graph/graph_store.rs @@ -1,10 +1,77 @@ +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: Default { type Handle: Clone; - // TODO(alexkirsz) An `entry(parent_handle) -> Entry` API would be more + // TODO(alexkirsz) An `entry(from_handle) -> Entry` API would be more // efficient, as right now we're getting the same key multiple times. /// Inserts a node into the graph store, and returns a handle to it. - fn insert(&mut self, parent_handle: Option, node: T) -> (Self::Handle, &T); + /// + /// If this method returns `None`, the node edges will not be visited. + fn insert( + &mut self, + from_handle: Option, + node: Node, + ) -> Option<(Self::Handle, &Node)>; +} + +/// 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 +where + StoreImpl: GraphStore, +{ + store: StoreImpl, + visited: HashSet, +} + +impl Default for SkipDuplicates +where + StoreImpl: GraphStore, +{ + fn default() -> Self { + Self { + store: Default::default(), + visited: Default::default(), + } + } +} + +impl GraphStore for SkipDuplicates +where + StoreImpl: GraphStore, + Node: Eq + std::hash::Hash + Clone, +{ + 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()); + 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 SkipDuplicates +where + StoreImpl: GraphStore, +{ + /// Consumes the wrapper and returns the underlying store. + pub fn into_inner(self) -> StoreImpl { + self.store + } } diff --git a/crates/turbo-tasks/src/graph/graph_traversal.rs b/crates/turbo-tasks/src/graph/graph_traversal.rs index 96ad03fb13d6a..685d0bb3ba16f 100644 --- a/crates/turbo-tasks/src/graph/graph_traversal.rs +++ b/crates/turbo-tasks/src/graph/graph_traversal.rs @@ -3,39 +3,45 @@ use std::{future::Future, pin::Pin}; use anyhow::Result; use futures::{stream::FuturesUnordered, Stream}; -use super::{graph_store::GraphStore, with_future::With, Visit}; +use super::{graph_store::GraphStore, with_future::With, 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 children. 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, +/// 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, } -impl GraphTraversal { +impl GraphTraversal { /// Visits the graph starting from the given `roots`, and returns a future /// that will resolve to the traversal result. - pub fn visit(roots: I, mut visit: C) -> GraphTraversalFuture + pub fn visit( + root_edges: RootEdgesIt, + mut visit: VisitImpl, + ) -> GraphTraversalFuture where - S: GraphStore, - I: IntoIterator, - C: Visit, + Store: GraphStore, + VisitImpl: Visit, + RootEdgesIt: IntoIterator, { - let mut store = S::default(); + let mut store = Store::default(); let futures = FuturesUnordered::new(); - for item in roots { - let (parent_handle, item) = store.insert(None, item); - match visit.get_children(item) { - GraphTraversalControlFlow::Continue(future) => { - futures.push(With::new(future, parent_handle)); + for edge in root_edges { + match visit.visit(edge) { + VisitControlFlow::Continue(node) => { + if let Some((parent_handle, node_ref)) = store.insert(None, node) { + futures.push(With::new(visit.edges(node_ref), parent_handle)); + } + } + VisitControlFlow::Skip(node) => { + store.insert(None, node); } - GraphTraversalControlFlow::Abort(abort) => { + VisitControlFlow::Abort(abort) => { return GraphTraversalFuture { state: GraphTraversalState::Aborted { abort }, }; } - GraphTraversalControlFlow::Skip => {} } } GraphTraversalFuture { @@ -50,64 +56,66 @@ impl GraphTraversal { /// A future that resolves to a [`GraphStore`] containing the result of a graph /// traversal. -pub struct GraphTraversalFuture +pub struct GraphTraversalFuture where - S: GraphStore, - C: Visit, + Store: GraphStore, + VisitImpl: Visit, { - state: GraphTraversalState, + state: GraphTraversalState, } -enum GraphTraversalState +enum GraphTraversalState where - S: GraphStore, - C: Visit, + Store: GraphStore, + VisitImpl: Visit, { Completed, - Aborted { abort: A }, - Running(GraphTraversalRunningState), + Aborted { abort: Abort }, + Running(GraphTraversalRunningState), } -struct GraphTraversalRunningState +struct GraphTraversalRunningState where - S: GraphStore, - C: Visit, + Store: GraphStore, + VisitImpl: Visit, { - store: S, - futures: FuturesUnordered>, - visit: C, + store: Store, + futures: FuturesUnordered>, + visit: VisitImpl, } -impl Default for GraphTraversalState +impl Default + for GraphTraversalState where - S: GraphStore, - C: Visit, + Store: GraphStore, + VisitImpl: Visit, { fn default() -> Self { GraphTraversalState::Completed } } -pub enum GraphTraversalResult { - Completed(R), - Aborted(A), +pub enum GraphTraversalResult { + Completed(Completed), + Aborted(Aborted), } -impl GraphTraversalResult { - pub fn completed(self) -> R { +impl GraphTraversalResult { + pub fn completed(self) -> Completed { match self { - GraphTraversalResult::Completed(result) => result, - GraphTraversalResult::Aborted(_) => unreachable!("the type parameter `A` is `!`"), + GraphTraversalResult::Completed(completed) => completed, + GraphTraversalResult::Aborted(_) => unreachable!("the type parameter `Aborted` is `!`"), } } } -impl Future for GraphTraversalFuture +impl Future + for GraphTraversalFuture where - S: GraphStore, - C: Visit, + Store: GraphStore, + VisitImpl: Visit, { - type Output = GraphTraversalResult, A>; + type Output = GraphTraversalResult, Abort>; fn poll( self: std::pin::Pin<&mut Self>, @@ -127,35 +135,30 @@ where GraphTraversalState::Running(mut running) => 'outer: loop { let futures_pin = unsafe { Pin::new_unchecked(&mut running.futures) }; match futures_pin.poll_next(cx) { - std::task::Poll::Ready(Some((parent_handle, Ok(children)))) => { - match running.visit.map_children(children) { - GraphTraversalControlFlow::Continue(children) => { - for item in children { - let (child_handle, item) = - running.store.insert(Some(parent_handle.clone()), item); - - match running.visit.get_children(item) { - GraphTraversalControlFlow::Continue(future) => { - running.futures.push(With::new(future, child_handle)); - } - GraphTraversalControlFlow::Skip => {} - GraphTraversalControlFlow::Abort(abort) => { - break 'outer ( - GraphTraversalState::Completed, - std::task::Poll::Ready( - GraphTraversalResult::Aborted(abort), - ), - ); - } + std::task::Poll::Ready(Some((parent_handle, Ok(edges)))) => { + 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) + { + running.futures.push(With::new( + running.visit.edges(node_ref), + node_handle, + )); } } - } - GraphTraversalControlFlow::Skip => {} - GraphTraversalControlFlow::Abort(abort) => { - break 'outer ( - GraphTraversalState::Completed, - std::task::Poll::Ready(GraphTraversalResult::Aborted(abort)), - ); + VisitControlFlow::Skip(node) => { + running.store.insert(Some(parent_handle.clone()), node); + } + VisitControlFlow::Abort(abort) => { + break 'outer ( + GraphTraversalState::Completed, + std::task::Poll::Ready(GraphTraversalResult::Aborted( + abort, + )), + ); + } } } } @@ -186,43 +189,3 @@ where result } } - -/// The control flow of a graph traversal. -pub enum GraphTraversalControlFlow { - /// The traversal should continue with the given future. - Continue(C), - /// The traversal should abort and return immediately. - Abort(A), - /// The traversal should skip the current node. - Skip, -} - -impl GraphTraversalControlFlow { - /// Map the continue value of this control flow. - pub fn map_continue(self, mut map: Map) -> GraphTraversalControlFlow - where - Map: FnMut(C) -> Mapped, - { - match self { - GraphTraversalControlFlow::Continue(future) => { - GraphTraversalControlFlow::Continue(map(future)) - } - GraphTraversalControlFlow::Abort(abort) => GraphTraversalControlFlow::Abort(abort), - GraphTraversalControlFlow::Skip => GraphTraversalControlFlow::Skip, - } - } - - /// Map the abort value of this control flow. - pub fn map_abort(self, mut map: Map) -> GraphTraversalControlFlow - where - Map: FnMut(A) -> Mapped, - { - match self { - GraphTraversalControlFlow::Continue(future) => { - GraphTraversalControlFlow::Continue(future) - } - GraphTraversalControlFlow::Abort(abort) => GraphTraversalControlFlow::Abort(map(abort)), - GraphTraversalControlFlow::Skip => GraphTraversalControlFlow::Skip, - } - } -} diff --git a/crates/turbo-tasks/src/graph/mod.rs b/crates/turbo-tasks/src/graph/mod.rs index f1bd232efa32a..1b59ae689a982 100644 --- a/crates/turbo-tasks/src/graph/mod.rs +++ b/crates/turbo-tasks/src/graph/mod.rs @@ -1,11 +1,14 @@ -mod get_children; +mod control_flow; mod graph_store; mod graph_traversal; mod non_deterministic; mod reverse_topological; +mod visit; mod with_future; -pub use get_children::{SkipDuplicates, Visit}; -pub use graph_traversal::{GraphTraversal, GraphTraversalControlFlow, GraphTraversalResult}; +pub use control_flow::VisitControlFlow; +pub use graph_store::{GraphStore, SkipDuplicates}; +pub use graph_traversal::{GraphTraversal, GraphTraversalResult}; pub use non_deterministic::NonDeterministic; pub use reverse_topological::ReverseTopological; +pub use visit::Visit; diff --git a/crates/turbo-tasks/src/graph/non_deterministic.rs b/crates/turbo-tasks/src/graph/non_deterministic.rs index c90f59bd991a1..1316c23b5b903 100644 --- a/crates/turbo-tasks/src/graph/non_deterministic.rs +++ b/crates/turbo-tasks/src/graph/non_deterministic.rs @@ -15,9 +15,13 @@ impl Default for NonDeterministic { impl GraphStore for NonDeterministic { type Handle = (); - fn insert(&mut self, _parent_handle: Option, node: T) -> (Self::Handle, &T) { + fn insert( + &mut self, + _from_handle: Option, + node: T, + ) -> Option<(Self::Handle, &T)> { self.output.push(node); - ((), self.output.last().unwrap()) + 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 c6e26c9e03a6b..4996641dd955b 100644 --- a/crates/turbo-tasks/src/graph/reverse_topological.rs +++ b/crates/turbo-tasks/src/graph/reverse_topological.rs @@ -29,17 +29,17 @@ where { type Handle = T; - fn insert(&mut self, parent: Option, node: T) -> (Self::Handle, &T) { - let vec = if let Some(parent) = parent { + fn insert(&mut self, from_handle: Option, node: T) -> Option<(Self::Handle, &T)> { + let vec = if let Some(from_handle) = from_handle { self.adjacency_map - .entry(parent) + .entry(from_handle) .or_insert_with(|| Vec::with_capacity(1)) } else { &mut self.roots }; vec.push(node.clone()); - (node, vec.last().unwrap()) + Some((node, vec.last().unwrap())) } } @@ -99,15 +99,15 @@ where self.visited.insert(current.clone()); - let Some(children) = self.adjacency_map.get(¤t) else { + let Some(neighbors) = self.adjacency_map.get(¤t) else { break current; }; self.stack.push((ReverseTopologicalPass::Post, current)); self.stack.extend( - children + neighbors .iter() - .map(|child| (ReverseTopologicalPass::Pre, child.clone())), + .map(|neighbor| (ReverseTopologicalPass::Pre, neighbor.clone())), ); } } diff --git a/crates/turbo-tasks/src/graph/visit.rs b/crates/turbo-tasks/src/graph/visit.rs new file mode 100644 index 0000000000000..2acf7dcdf64d0 --- /dev/null +++ b/crates/turbo-tasks/src/graph/visit.rs @@ -0,0 +1,72 @@ +use std::future::Future; + +use anyhow::Result; + +use super::VisitControlFlow; + +/// A trait that allows a graph traversal to visit the edges of a node +/// transitively. +pub trait Visit { + type Edge; + type EdgesIntoIter: IntoIterator; + type EdgesFuture: Future>; + + /// Visits an edge to get to the neighbor node. Should return a + /// [`VisitControlFlow`] that indicates whether to: + /// * continue visiting the neighbor node edges; + /// * skip visiting the neighbor node's edges; + /// * abort the traversal entirely. + fn visit(&mut self, edge: Self::Edge) -> VisitControlFlow; + + /// Returns a future that resolves to the outgoing edges of the given + /// `node`. + fn edges(&mut self, node: &Node) -> Self::EdgesFuture; +} + +// The different `Impl*` here are necessary in order to avoid the `Conflicting +// implementations of trait` error when implementing `Visit` on different +// kinds of `FnMut`. +// See https://users.rust-lang.org/t/conflicting-implementation-when-implementing-traits-for-fn/53359/3 + +pub struct ImplRef; + +impl Visit for VisitFn +where + VisitFn: FnMut(&Node) -> NeighFut, + NeighFut: Future>, + NeighIt: IntoIterator, +{ + type Edge = Node; + type EdgesIntoIter = NeighIt; + type EdgesFuture = NeighFut; + + fn visit(&mut self, edge: Self::Edge) -> VisitControlFlow { + VisitControlFlow::Continue(edge) + } + + fn edges(&mut self, node: &Node) -> Self::EdgesFuture { + (self)(node) + } +} + +pub struct ImplValue; + +impl Visit for VisitFn +where + Node: Clone, + VisitFn: FnMut(Node) -> NeighFut, + NeighFut: Future>, + NeighIt: IntoIterator, +{ + type Edge = Node; + type EdgesIntoIter = NeighIt; + type EdgesFuture = NeighFut; + + fn visit(&mut self, edge: Self::Edge) -> VisitControlFlow { + VisitControlFlow::Continue(edge) + } + + fn edges(&mut self, node: &Node) -> Self::EdgesFuture { + (self)(node.clone()) + } +} diff --git a/crates/turbopack-core/src/changed.rs b/crates/turbopack-core/src/changed.rs index 61291c7a049d9..c8a6fdf84dd7a 100644 --- a/crates/turbopack-core/src/changed.rs +++ b/crates/turbopack-core/src/changed.rs @@ -20,12 +20,13 @@ async fn get_referenced_assets(parent: AssetVc) -> Result Result { - let completions = GraphTraversal::>::visit( + let completions = GraphTraversal::, _>>::visit( [root], - SkipDuplicates::new(get_referenced_assets), + get_referenced_assets, ) .await .completed()? + .into_inner() .into_iter() .map(content_changed) .collect(); diff --git a/crates/turbopack-core/src/chunk/mod.rs b/crates/turbopack-core/src/chunk/mod.rs index 0216531a9ebc0..e17267cbea021 100644 --- a/crates/turbopack-core/src/chunk/mod.rs +++ b/crates/turbopack-core/src/chunk/mod.rs @@ -14,8 +14,8 @@ use serde::{Deserialize, Serialize}; use turbo_tasks::{ debug::ValueDebugFormat, graph::{ - GraphTraversal, GraphTraversalControlFlow, GraphTraversalResult, ReverseTopological, - SkipDuplicates, Visit, + GraphTraversal, GraphTraversalResult, ReverseTopological, SkipDuplicates, Visit, + VisitControlFlow, }, primitives::{BoolVc, StringVc}, trace::TraceRawVcs, @@ -131,12 +131,13 @@ impl ChunkGroupVc { /// All chunks should be loaded in parallel. #[turbo_tasks::function] pub async fn chunks(self) -> Result { - let chunks: Vec<_> = GraphTraversal::>::visit( + let chunks: Vec<_> = GraphTraversal::, _>>::visit( [self.await?.entry], - SkipDuplicates::new(get_chunk_children), + get_chunk_children, ) .await .completed()? + .into_inner() .into_iter() .collect(); @@ -535,24 +536,6 @@ where Ok(graph_nodes) } -async fn chunk_item_to_graph_nodes( - context: ChunkContentContext, - chunk_item: I, -) -> Result> -where - I: FromChunkableAsset + Eq + std::hash::Hash + Clone, -{ - Ok(chunk_item - .references() - .await? - .into_iter() - .map(|reference| reference_to_graph_nodes::(context, *reference)) - .try_join() - .await? - .into_iter() - .flatten()) -} - /// The maximum number of chunk items that can be in a chunk before we split it /// into multiple chunks. const MAX_CHUNK_ITEMS_COUNT: usize = 5000; @@ -564,66 +547,71 @@ struct ChunkContentVisit { _phantom: PhantomData, } -type ChunkItemToGraphNodesChildren = +type ChunkItemToGraphNodesEdges = impl Iterator, ChunkContentGraphNode)>; -type ChunkItemToGraphNodesFuture = - impl Future>>; +type ChunkItemToGraphNodesFuture = + impl Future>>; impl Visit, ()> for ChunkContentVisit where I: FromChunkableAsset + Eq + std::hash::Hash + Clone, { - type Children = ChunkItemToGraphNodesChildren; - type MapChildren = Vec>; - type Future = ChunkItemToGraphNodesFuture; + type Edge = (Option<(AssetVc, ChunkingType)>, ChunkContentGraphNode); + type EdgesIntoIter = ChunkItemToGraphNodesEdges; + type EdgesFuture = ChunkItemToGraphNodesFuture; - fn get_children( + fn visit( &mut self, - node: &ChunkContentGraphNode, - ) -> GraphTraversalControlFlow { - // TODO(alexkirsz) Right now the output type of the traversal is the same as the - // input type. If we can have them be separate types, we wouldn't need - // to call `get_children` on graph nodes that we know not to have any children - // (i.e. anything that's not a chunk item). - let ChunkContentGraphNode::ChunkItem(chunk_item) = node else { - return GraphTraversalControlFlow::Skip; + (option_key, node): (Option<(AssetVc, ChunkingType)>, ChunkContentGraphNode), + ) -> VisitControlFlow, ()> { + let Some((asset, chunking_type)) = option_key else { + return VisitControlFlow::Continue(node); }; - GraphTraversalControlFlow::Continue(chunk_item_to_graph_nodes::( - self.context, - chunk_item.clone(), - )) - } + if !self.processed_assets.insert((chunking_type, asset)) { + return VisitControlFlow::Skip(node); + } - fn map_children( - &mut self, - children: Self::Children, - ) -> GraphTraversalControlFlow { - let mut map_children = vec![]; - - for (option_key, child) in children { - if let Some((asset, chunking_type)) = option_key { - if self.processed_assets.insert((chunking_type, asset)) { - if let ChunkContentGraphNode::ChunkItem(_) = &child { - self.chunk_items_count += 1; - } - map_children.push(child); - } - } else { - map_children.push(child); + if let ChunkContentGraphNode::ChunkItem(_) = &node { + self.chunk_items_count += 1; + + // Make sure the chunk doesn't become too large. + // This will hurt performance in many aspects. + if !self.context.split && self.chunk_items_count >= MAX_CHUNK_ITEMS_COUNT { + // Chunk is too large, cancel this algorithm and restart with splitting from the + // start. + return VisitControlFlow::Abort(()); } } - // Make sure the chunk doesn't become too large. - // This will hurt performance in many aspects. - if !self.context.split && self.chunk_items_count >= MAX_CHUNK_ITEMS_COUNT { - // Chunk is too large, cancel this algorithm and restart with splitting from the - // start. - return GraphTraversalControlFlow::Abort(()); - } + VisitControlFlow::Continue(node) + } - GraphTraversalControlFlow::Continue(map_children) + fn edges(&mut self, node: &ChunkContentGraphNode) -> Self::EdgesFuture { + let chunk_item = if let ChunkContentGraphNode::ChunkItem(chunk_item) = node { + Some(chunk_item.clone()) + } else { + None + }; + + let context = self.context; + + async move { + let Some(chunk_item) = chunk_item else { + return Ok(vec![].into_iter().flatten()); + }; + + Ok(chunk_item + .references() + .await? + .into_iter() + .map(|reference| reference_to_graph_nodes::(context, *reference)) + .try_join() + .await? + .into_iter() + .flatten()) + } } } @@ -636,30 +624,26 @@ async fn chunk_content_internal_parallel( where I: FromChunkableAsset + Eq + std::hash::Hash + Clone, { - let mut processed_assets: HashSet<(ChunkingType, AssetVc)> = HashSet::default(); - let additional_entries = if let Some(additional_entries) = additional_entries { additional_entries.await?.clone_value().into_iter() } else { vec![].into_iter() }; - let entries = [entry].into_iter().chain(additional_entries); - for entry in entries.clone() { - processed_assets.insert((ChunkingType::Placed, entry)); - } - - let root_graph_nodes = entries + let root_edges = [entry] + .into_iter() + .chain(additional_entries) .map(|entry| async move { - Ok(ChunkContentGraphNode::ChunkItem( - I::from_asset(chunking_context, entry).await?.unwrap(), + Ok(( + Some((entry, ChunkingType::Placed)), + ChunkContentGraphNode::ChunkItem( + I::from_asset(chunking_context, entry).await?.unwrap(), + ), )) }) .try_join() .await?; - processed_assets.insert((ChunkingType::Placed, entry)); - let context = ChunkContentContext { chunking_context, entry, @@ -668,13 +652,13 @@ where let visit = ChunkContentVisit { context, - chunk_items_count: root_graph_nodes.len(), - processed_assets, + chunk_items_count: 0, + processed_assets: Default::default(), _phantom: PhantomData, }; let GraphTraversalResult::Completed(traversal_result) = - GraphTraversal::>::visit(root_graph_nodes, visit).await else { + GraphTraversal::>::visit(root_edges, visit).await else { return Ok(None); };