// Copyright 2020 The Jujutsu Authors//// Licensed under the Apache License, Version 2.0 (the "License");// you may not use this file except in compliance with the License.// You may obtain a copy of the License at//// https://www.apache.org/licenses/LICENSE-2.0//// Unless required by applicable law or agreed to in writing, software// distributed under the License is distributed on an "AS IS" BASIS,// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.// See the License for the specific language governing permissions and// limitations under the License. //! General-purpose DAG algorithms. use std::collections::BinaryHeap;use std::collections::HashMap;use std::collections::HashSet;use std::convert::Infallible;use std::hash::Hash;use std::iter;use std::mem; use itertools::Itertools as _;use smallvec::SmallVec;use smallvec::smallvec_inline; /// Traverses nodes from `start` in depth-first order.pub fn dfs<T, ID, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI,) -> impl Iterator<Item = T>where ID: Hash + Eq, II: IntoIterator<Item = T>, NI: IntoIterator<Item = T>,{ let neighbors_fn = move |node: &T| to_infallibe_iter(neighbors_fn(node)); dfs_ok(to_infallibe_iter(start), id_fn, neighbors_fn).map(|Ok(node)| node)} /// Traverses nodes from `start` in depth-first order.////// An `Err` is emitted as a node with no neighbors. Caller may decide to/// short-circuit on it.pub fn dfs_ok<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI,) -> impl Iterator<Item = Result<T, E>>where ID: Hash + Eq, II: IntoIterator<Item = Result<T, E>>, NI: IntoIterator<Item = Result<T, E>>,{ let mut work: Vec<Result<T, E>> = start.into_iter().collect(); let mut visited: HashSet<ID> = HashSet::new(); iter::from_fn(move || { loop { let c = match work.pop() { Some(Ok(c)) => c, r @ (Some(Err(_)) | None) => return r, }; let id = id_fn(&c); if visited.contains(&id) { continue; } for p in neighbors_fn(&c) { work.push(p); } visited.insert(id); return Some(Ok(c)); } })} /// Builds a list of nodes reachable from the `start` where neighbors come/// before the node itself.////// If the graph has cycle, `cycle_fn()` is called with one of the nodes/// involved in the cycle.pub fn topo_order_forward<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnOnce(T) -> E,) -> Result<Vec<T>, E>where ID: Hash + Eq + Clone, II: IntoIterator<Item = T>, NI: IntoIterator<Item = T>,{ let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node)); topo_order_forward_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_fn)} /// Builds a list of `Ok` nodes reachable from the `start` where neighbors come/// before the node itself.////// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and/// returns the error. If the graph has cycle, `cycle_fn()` is called with one/// of the nodes involved in the cycle.pub fn topo_order_forward_ok<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnOnce(T) -> E,) -> Result<Vec<T>, E>where ID: Hash + Eq + Clone, II: IntoIterator<Item = Result<T, E>>, NI: IntoIterator<Item = Result<T, E>>,{ let mut stack: Vec<(T, bool)> = start.into_iter().map(|r| Ok((r?, false))).try_collect()?; let mut visiting = HashSet::new(); let mut emitted = HashSet::new(); let mut result = vec![]; while let Some((node, neighbors_visited)) = stack.pop() { let id = id_fn(&node); if emitted.contains(&id) { continue; } if !neighbors_visited { if !visiting.insert(id.clone()) { return Err(cycle_fn(node)); } let neighbors_iter = neighbors_fn(&node).into_iter(); stack.reserve(neighbors_iter.size_hint().0 + 1); stack.push((node, true)); for neighbor in neighbors_iter { stack.push((neighbor?, false)); } } else { visiting.remove(&id); emitted.insert(id); result.push(node); } } Ok(result)} /// Builds a list of nodes reachable from the `start` where neighbors come after/// the node itself.////// If the graph has cycle, `cycle_fn()` is called with one of the nodes/// involved in the cycle.pub fn topo_order_reverse<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnOnce(T) -> E,) -> Result<Vec<T>, E>where ID: Hash + Eq + Clone, II: IntoIterator<Item = T>, NI: IntoIterator<Item = T>,{ let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node)); topo_order_reverse_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_fn)} /// Builds a list of `Ok` nodes reachable from the `start` where neighbors come/// after the node itself.////// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and/// returns the error. If the graph has cycle, `cycle_fn()` is called with one/// of the nodes involved in the cycle.pub fn topo_order_reverse_ok<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnOnce(T) -> E,) -> Result<Vec<T>, E>where ID: Hash + Eq + Clone, II: IntoIterator<Item = Result<T, E>>, NI: IntoIterator<Item = Result<T, E>>,{ let mut result = topo_order_forward_ok(start, id_fn, neighbors_fn, cycle_fn)?; result.reverse(); Ok(result)} /// Like `topo_order_reverse()`, but can iterate linear DAG lazily.////// The DAG is supposed to be (mostly) topologically ordered by `T: Ord`./// For example, topological order of chronological data should respect/// timestamp (except a few outliers caused by clock skew.)////// Use `topo_order_reverse()` if the DAG is heavily branched. This can/// only process linear part lazily.////// If the graph has cycle, `cycle_fn()` is called with one of the nodes/// involved in the cycle.pub fn topo_order_reverse_lazy<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnMut(T) -> E,) -> impl Iterator<Item = Result<T, E>>where T: Ord, ID: Hash + Eq + Clone, II: IntoIterator<Item = T>, NI: IntoIterator<Item = T>,{ let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node)); topo_order_reverse_lazy_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_fn)} /// Like `topo_order_reverse_ok()`, but can iterate linear DAG lazily.////// The returned iterator short-circuits at an `Err`. Pending non-linear nodes/// before the `Err` will be discarded. If the graph has cycle, `cycle_fn()` is/// called with one of the nodes involved in the cycle.pub fn topo_order_reverse_lazy_ok<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, mut cycle_fn: impl FnMut(T) -> E,) -> impl Iterator<Item = Result<T, E>>where T: Ord, ID: Hash + Eq + Clone, II: IntoIterator<Item = Result<T, E>>, NI: IntoIterator<Item = Result<T, E>>,{ let mut inner = TopoOrderReverseLazyInner::empty(); inner.extend(start); iter::from_fn(move || inner.next(&id_fn, &mut neighbors_fn, &mut cycle_fn))} #[derive(Clone, Debug)]struct TopoOrderReverseLazyInner<T, ID, E> { start: Vec<T>, result: Vec<Result<T, E>>, emitted: HashSet<ID>,} impl<T: Ord, ID: Hash + Eq + Clone, E> TopoOrderReverseLazyInner<T, ID, E> { fn empty() -> Self { Self { start: Vec::new(), result: Vec::new(), emitted: HashSet::new(), } } fn extend(&mut self, iter: impl IntoIterator<Item = Result<T, E>>) { let iter = iter.into_iter(); self.start.reserve(iter.size_hint().0); for res in iter { if let Ok(node) = res { self.start.push(node); } else { // Emit the error and terminate self.start.clear(); self.result.insert(0, res); return; } } } fn next<NI: IntoIterator<Item = Result<T, E>>>( &mut self, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, mut cycle_fn: impl FnMut(T) -> E, ) -> Option<Result<T, E>> { if let Some(res) = self.result.pop() { return Some(res); } // Fast path for linear DAG if self.start.len() <= 1 { let node = self.start.pop()?; self.extend(neighbors_fn(&node)); if self.emitted.insert(id_fn(&node)) { return Some(Ok(node)); } else { return Some(Err(cycle_fn(node))); } } // Extract graph nodes based on T's order, and sort them by using ids // (because we wouldn't want to clone T itself) let start_ids = self.start.iter().map(&id_fn).collect_vec(); match look_ahead_sub_graph(mem::take(&mut self.start), &id_fn, &mut neighbors_fn) { Ok((mut node_map, neighbor_ids_map, remainder)) => { self.start = remainder; let sorted_ids = match topo_order_forward_ok( start_ids.iter().map(Ok), |id| *id, |id| neighbor_ids_map[id].iter().map(Ok), |id| cycle_fn(node_map.remove(id).unwrap()), ) { Ok(ids) => ids, Err(err) => return Some(Err(err)), }; self.result.reserve(sorted_ids.len()); for id in sorted_ids { let (id, node) = node_map.remove_entry(id).unwrap(); if self.emitted.insert(id) { self.result.push(Ok(node)); } else { self.result.push(Err(cycle_fn(node))); } } self.result.pop() } Err(err) => Some(Err(err)), } }} /// Splits DAG at the first single fork point, and builds a list of nodes/// reachable from the `start` where neighbors come after the node itself.////// This is a building block for lazy DAG iterators similar to/// [`topo_order_reverse_lazy()`]. The `start` list will be updated to include/// the next nodes to visit.////// If the split chunk of the graph has cycle, `cycle_fn()` is called with one/// of the nodes involved in the cycle.pub fn topo_order_reverse_chunked<T, ID, E, NI>( start: &mut Vec<T>, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, mut cycle_fn: impl FnMut(T) -> E,) -> Result<SmallVec<[T; 1]>, E>where T: Ord, ID: Hash + Eq + Clone, NI: IntoIterator<Item = Result<T, E>>,{ // Fast path for linear DAG if start.len() <= 1 { let Some(node) = start.pop() else { return Ok(SmallVec::new()); }; let neighbors_iter = neighbors_fn(&node).into_iter(); start.reserve(neighbors_iter.size_hint().0); for neighbor in neighbors_iter { start.push(neighbor?); } return Ok(smallvec_inline![node]); } // Extract graph nodes based on T's order, and sort them by using ids // (because we wouldn't want to clone T itself) let start_ids = start.iter().map(&id_fn).collect_vec(); let (mut node_map, neighbor_ids_map, remainder) = look_ahead_sub_graph(mem::take(start), &id_fn, &mut neighbors_fn)?; *start = remainder; let sorted_ids = topo_order_forward_ok( start_ids.iter().map(Ok), |id| *id, |id| neighbor_ids_map[id].iter().map(Ok), |id| cycle_fn(node_map.remove(id).unwrap()), )?; let sorted_nodes = sorted_ids .iter() .rev() .map(|&id| node_map.remove(id).unwrap()) .collect(); Ok(sorted_nodes)} /// Splits DAG at single fork point, and extracts branchy part as sub graph.////// ```text/// o | C/// | o B/// |/ <---- split here (A->B or A->C would create cycle)/// o A/// ```////// If a branch reached to root (empty neighbors), the graph can't be split/// anymore because the other branch may be connected to a descendant of/// the rooted branch.////// ```text/// o | C/// | o B/// | <---- can't split here (there may be edge A->B)/// o A/// ```////// We assume the graph is (mostly) topologically ordered by `T: Ord`.#[expect(clippy::type_complexity)]fn look_ahead_sub_graph<T, ID, E, NI>( start: Vec<T>, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI,) -> Result<(HashMap<ID, T>, HashMap<ID, Vec<ID>>, Vec<T>), E>where T: Ord, ID: Hash + Eq + Clone, NI: IntoIterator<Item = Result<T, E>>,{ let mut queue: BinaryHeap<T> = start.into(); // Build separate node/neighbors maps since lifetime is different at caller let mut node_map: HashMap<ID, T> = HashMap::new(); let mut neighbor_ids_map: HashMap<ID, Vec<ID>> = HashMap::new(); let mut has_reached_root = false; while queue.len() > 1 || node_map.is_empty() || has_reached_root { let Some(node) = queue.pop() else { break; }; let node_id = id_fn(&node); if node_map.contains_key(&node_id) { continue; } let mut neighbor_ids = Vec::new(); let mut neighbors_iter = neighbors_fn(&node).into_iter().peekable(); has_reached_root |= neighbors_iter.peek().is_none(); for neighbor in neighbors_iter { let neighbor = neighbor?; neighbor_ids.push(id_fn(&neighbor)); queue.push(neighbor); } node_map.insert(node_id.clone(), node); neighbor_ids_map.insert(node_id, neighbor_ids); } assert!(queue.len() <= 1, "order of remainder shouldn't matter"); let remainder = queue.into_vec(); // Omit unvisited neighbors if let Some(unvisited_id) = remainder.first().map(&id_fn) { for neighbor_ids in neighbor_ids_map.values_mut() { neighbor_ids.retain(|id| *id != unvisited_id); } } Ok((node_map, neighbor_ids_map, remainder))} /// Builds a list of nodes reachable from the `start` where neighbors come after/// the node itself.////// Unlike `topo_order_reverse()`, nodes are sorted in reverse `T: Ord` order so/// long as they can respect the topological requirement.////// If the graph has cycle, `cycle_fn()` is called with one of the nodes/// involved in the cycle.pub fn topo_order_reverse_ord<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnOnce(T) -> E,) -> Result<Vec<T>, E>where T: Ord, ID: Hash + Eq + Clone, II: IntoIterator<Item = T>, NI: IntoIterator<Item = T>,{ let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node)); topo_order_reverse_ord_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_fn)} /// Builds a list of `Ok` nodes reachable from the `start` where neighbors come/// after the node itself.////// Unlike `topo_order_reverse_ok()`, nodes are sorted in reverse `T: Ord` order/// so long as they can respect the topological requirement.////// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and/// returns the error. If the graph has cycle, `cycle_fn()` is called with one/// of the nodes involved in the cycle.pub fn topo_order_reverse_ord_ok<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnOnce(T) -> E,) -> Result<Vec<T>, E>where T: Ord, ID: Hash + Eq + Clone, II: IntoIterator<Item = Result<T, E>>, NI: IntoIterator<Item = Result<T, E>>,{ struct InnerNode<T> { node: Option<T>, indegree: usize, } // DFS to accumulate incoming edges let mut stack: Vec<T> = start.into_iter().try_collect()?; let mut head_node_map: HashMap<ID, T> = HashMap::new(); let mut inner_node_map: HashMap<ID, InnerNode<T>> = HashMap::new(); let mut neighbor_ids_map: HashMap<ID, Vec<ID>> = HashMap::new(); while let Some(node) = stack.pop() { let node_id = id_fn(&node); if neighbor_ids_map.contains_key(&node_id) { continue; // Already visitedLines 1–500 of 1498. Load more
// Copyright 2020 The Jujutsu Authors // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // https://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. //! General-purpose DAG algorithms. use std::collections::BinaryHeap; use std::collections::HashMap; use std::collections::HashSet; use std::convert::Infallible; use std::hash::Hash; use std::iter; use std::mem; use itertools::Itertools as _; use smallvec::SmallVec; use smallvec::smallvec_inline; /// Traverses nodes from `start` in depth-first order. pub fn dfs<T, ID, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, ) -> impl Iterator<Item = T> where ID: Hash + Eq, II: IntoIterator<Item = T>, NI: IntoIterator<Item = T>, { let neighbors_fn = move |node: &T| to_infallibe_iter(neighbors_fn(node)); dfs_ok(to_infallibe_iter(start), id_fn, neighbors_fn).map(|Ok(node)| node) } /// Traverses nodes from `start` in depth-first order. /// /// An `Err` is emitted as a node with no neighbors. Caller may decide to /// short-circuit on it. pub fn dfs_ok<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, ) -> impl Iterator<Item = Result<T, E>> where ID: Hash + Eq, II: IntoIterator<Item = Result<T, E>>, NI: IntoIterator<Item = Result<T, E>>, { let mut work: Vec<Result<T, E>> = start.into_iter().collect(); let mut visited: HashSet<ID> = HashSet::new(); iter::from_fn(move || { loop { let c = match work.pop() { Some(Ok(c)) => c, r @ (Some(Err(_)) | None) => return r, }; let id = id_fn(&c); if visited.contains(&id) { continue; } for p in neighbors_fn(&c) { work.push(p); } visited.insert(id); return Some(Ok(c)); } }) } /// Builds a list of nodes reachable from the `start` where neighbors come /// before the node itself. /// /// If the graph has cycle, `cycle_fn()` is called with one of the nodes /// involved in the cycle. pub fn topo_order_forward<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnOnce(T) -> E, ) -> Result<Vec<T>, E> where ID: Hash + Eq + Clone, II: IntoIterator<Item = T>, NI: IntoIterator<Item = T>, { let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node)); topo_order_forward_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_fn) } /// Builds a list of `Ok` nodes reachable from the `start` where neighbors come /// before the node itself. /// /// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and /// returns the error. If the graph has cycle, `cycle_fn()` is called with one /// of the nodes involved in the cycle. pub fn topo_order_forward_ok<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnOnce(T) -> E, ) -> Result<Vec<T>, E> where ID: Hash + Eq + Clone, II: IntoIterator<Item = Result<T, E>>, NI: IntoIterator<Item = Result<T, E>>, { let mut stack: Vec<(T, bool)> = start.into_iter().map(|r| Ok((r?, false))).try_collect()?; let mut visiting = HashSet::new(); let mut emitted = HashSet::new(); let mut result = vec![]; while let Some((node, neighbors_visited)) = stack.pop() { let id = id_fn(&node); if emitted.contains(&id) { continue; } if !neighbors_visited { if !visiting.insert(id.clone()) { return Err(cycle_fn(node)); } let neighbors_iter = neighbors_fn(&node).into_iter(); stack.reserve(neighbors_iter.size_hint().0 + 1); stack.push((node, true)); for neighbor in neighbors_iter { stack.push((neighbor?, false)); } } else { visiting.remove(&id); emitted.insert(id); result.push(node); } } Ok(result) } /// Builds a list of nodes reachable from the `start` where neighbors come after /// the node itself. /// /// If the graph has cycle, `cycle_fn()` is called with one of the nodes /// involved in the cycle. pub fn topo_order_reverse<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnOnce(T) -> E, ) -> Result<Vec<T>, E> where ID: Hash + Eq + Clone, II: IntoIterator<Item = T>, NI: IntoIterator<Item = T>, { let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node)); topo_order_reverse_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_fn) } /// Builds a list of `Ok` nodes reachable from the `start` where neighbors come /// after the node itself. /// /// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and /// returns the error. If the graph has cycle, `cycle_fn()` is called with one /// of the nodes involved in the cycle. pub fn topo_order_reverse_ok<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnOnce(T) -> E, ) -> Result<Vec<T>, E> where ID: Hash + Eq + Clone, II: IntoIterator<Item = Result<T, E>>, NI: IntoIterator<Item = Result<T, E>>, { let mut result = topo_order_forward_ok(start, id_fn, neighbors_fn, cycle_fn)?; result.reverse(); Ok(result) } /// Like `topo_order_reverse()`, but can iterate linear DAG lazily. /// /// The DAG is supposed to be (mostly) topologically ordered by `T: Ord`. /// For example, topological order of chronological data should respect /// timestamp (except a few outliers caused by clock skew.) /// /// Use `topo_order_reverse()` if the DAG is heavily branched. This can /// only process linear part lazily. /// /// If the graph has cycle, `cycle_fn()` is called with one of the nodes /// involved in the cycle. pub fn topo_order_reverse_lazy<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnMut(T) -> E, ) -> impl Iterator<Item = Result<T, E>> where T: Ord, ID: Hash + Eq + Clone, II: IntoIterator<Item = T>, NI: IntoIterator<Item = T>, { let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node)); topo_order_reverse_lazy_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_fn) } /// Like `topo_order_reverse_ok()`, but can iterate linear DAG lazily. /// /// The returned iterator short-circuits at an `Err`. Pending non-linear nodes /// before the `Err` will be discarded. If the graph has cycle, `cycle_fn()` is /// called with one of the nodes involved in the cycle. pub fn topo_order_reverse_lazy_ok<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, mut cycle_fn: impl FnMut(T) -> E, ) -> impl Iterator<Item = Result<T, E>> where T: Ord, ID: Hash + Eq + Clone, II: IntoIterator<Item = Result<T, E>>, NI: IntoIterator<Item = Result<T, E>>, { let mut inner = TopoOrderReverseLazyInner::empty(); inner.extend(start); iter::from_fn(move || inner.next(&id_fn, &mut neighbors_fn, &mut cycle_fn)) } #[derive(Clone, Debug)] struct TopoOrderReverseLazyInner<T, ID, E> { start: Vec<T>, result: Vec<Result<T, E>>, emitted: HashSet<ID>, } impl<T: Ord, ID: Hash + Eq + Clone, E> TopoOrderReverseLazyInner<T, ID, E> { fn empty() -> Self { Self { start: Vec::new(), result: Vec::new(), emitted: HashSet::new(), } } fn extend(&mut self, iter: impl IntoIterator<Item = Result<T, E>>) { let iter = iter.into_iter(); self.start.reserve(iter.size_hint().0); for res in iter { if let Ok(node) = res { self.start.push(node); } else { // Emit the error and terminate self.start.clear(); self.result.insert(0, res); return; } } } fn next<NI: IntoIterator<Item = Result<T, E>>>( &mut self, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, mut cycle_fn: impl FnMut(T) -> E, ) -> Option<Result<T, E>> { if let Some(res) = self.result.pop() { return Some(res); } // Fast path for linear DAG if self.start.len() <= 1 { let node = self.start.pop()?; self.extend(neighbors_fn(&node)); if self.emitted.insert(id_fn(&node)) { return Some(Ok(node)); } else { return Some(Err(cycle_fn(node))); } } // Extract graph nodes based on T's order, and sort them by using ids // (because we wouldn't want to clone T itself) let start_ids = self.start.iter().map(&id_fn).collect_vec(); match look_ahead_sub_graph(mem::take(&mut self.start), &id_fn, &mut neighbors_fn) { Ok((mut node_map, neighbor_ids_map, remainder)) => { self.start = remainder; let sorted_ids = match topo_order_forward_ok( start_ids.iter().map(Ok), |id| *id, |id| neighbor_ids_map[id].iter().map(Ok), |id| cycle_fn(node_map.remove(id).unwrap()), ) { Ok(ids) => ids, Err(err) => return Some(Err(err)), }; self.result.reserve(sorted_ids.len()); for id in sorted_ids { let (id, node) = node_map.remove_entry(id).unwrap(); if self.emitted.insert(id) { self.result.push(Ok(node)); } else { self.result.push(Err(cycle_fn(node))); } } self.result.pop() } Err(err) => Some(Err(err)), } } } /// Splits DAG at the first single fork point, and builds a list of nodes /// reachable from the `start` where neighbors come after the node itself. /// /// This is a building block for lazy DAG iterators similar to /// [`topo_order_reverse_lazy()`]. The `start` list will be updated to include /// the next nodes to visit. /// /// If the split chunk of the graph has cycle, `cycle_fn()` is called with one /// of the nodes involved in the cycle. pub fn topo_order_reverse_chunked<T, ID, E, NI>( start: &mut Vec<T>, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, mut cycle_fn: impl FnMut(T) -> E, ) -> Result<SmallVec<[T; 1]>, E> where T: Ord, ID: Hash + Eq + Clone, NI: IntoIterator<Item = Result<T, E>>, { // Fast path for linear DAG if start.len() <= 1 { let Some(node) = start.pop() else { return Ok(SmallVec::new()); }; let neighbors_iter = neighbors_fn(&node).into_iter(); start.reserve(neighbors_iter.size_hint().0); for neighbor in neighbors_iter { start.push(neighbor?); } return Ok(smallvec_inline![node]); } // Extract graph nodes based on T's order, and sort them by using ids // (because we wouldn't want to clone T itself) let start_ids = start.iter().map(&id_fn).collect_vec(); let (mut node_map, neighbor_ids_map, remainder) = look_ahead_sub_graph(mem::take(start), &id_fn, &mut neighbors_fn)?; *start = remainder; let sorted_ids = topo_order_forward_ok( start_ids.iter().map(Ok), |id| *id, |id| neighbor_ids_map[id].iter().map(Ok), |id| cycle_fn(node_map.remove(id).unwrap()), )?; let sorted_nodes = sorted_ids .iter() .rev() .map(|&id| node_map.remove(id).unwrap()) .collect(); Ok(sorted_nodes) } /// Splits DAG at single fork point, and extracts branchy part as sub graph. /// /// ```text /// o | C /// | o B /// |/ <---- split here (A->B or A->C would create cycle) /// o A /// ``` /// /// If a branch reached to root (empty neighbors), the graph can't be split /// anymore because the other branch may be connected to a descendant of /// the rooted branch. /// /// ```text /// o | C /// | o B /// | <---- can't split here (there may be edge A->B) /// o A /// ``` /// /// We assume the graph is (mostly) topologically ordered by `T: Ord`. #[expect(clippy::type_complexity)] fn look_ahead_sub_graph<T, ID, E, NI>( start: Vec<T>, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, ) -> Result<(HashMap<ID, T>, HashMap<ID, Vec<ID>>, Vec<T>), E> where T: Ord, ID: Hash + Eq + Clone, NI: IntoIterator<Item = Result<T, E>>, { let mut queue: BinaryHeap<T> = start.into(); // Build separate node/neighbors maps since lifetime is different at caller let mut node_map: HashMap<ID, T> = HashMap::new(); let mut neighbor_ids_map: HashMap<ID, Vec<ID>> = HashMap::new(); let mut has_reached_root = false; while queue.len() > 1 || node_map.is_empty() || has_reached_root { let Some(node) = queue.pop() else { break; }; let node_id = id_fn(&node); if node_map.contains_key(&node_id) { continue; } let mut neighbor_ids = Vec::new(); let mut neighbors_iter = neighbors_fn(&node).into_iter().peekable(); has_reached_root |= neighbors_iter.peek().is_none(); for neighbor in neighbors_iter { let neighbor = neighbor?; neighbor_ids.push(id_fn(&neighbor)); queue.push(neighbor); } node_map.insert(node_id.clone(), node); neighbor_ids_map.insert(node_id, neighbor_ids); } assert!(queue.len() <= 1, "order of remainder shouldn't matter"); let remainder = queue.into_vec(); // Omit unvisited neighbors if let Some(unvisited_id) = remainder.first().map(&id_fn) { for neighbor_ids in neighbor_ids_map.values_mut() { neighbor_ids.retain(|id| *id != unvisited_id); } } Ok((node_map, neighbor_ids_map, remainder)) } /// Builds a list of nodes reachable from the `start` where neighbors come after /// the node itself. /// /// Unlike `topo_order_reverse()`, nodes are sorted in reverse `T: Ord` order so /// long as they can respect the topological requirement. /// /// If the graph has cycle, `cycle_fn()` is called with one of the nodes /// involved in the cycle. pub fn topo_order_reverse_ord<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnOnce(T) -> E, ) -> Result<Vec<T>, E> where T: Ord, ID: Hash + Eq + Clone, II: IntoIterator<Item = T>, NI: IntoIterator<Item = T>, { let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node)); topo_order_reverse_ord_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_fn) } /// Builds a list of `Ok` nodes reachable from the `start` where neighbors come /// after the node itself. /// /// Unlike `topo_order_reverse_ok()`, nodes are sorted in reverse `T: Ord` order /// so long as they can respect the topological requirement. /// /// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and /// returns the error. If the graph has cycle, `cycle_fn()` is called with one /// of the nodes involved in the cycle. pub fn topo_order_reverse_ord_ok<T, ID, E, II, NI>( start: II, id_fn: impl Fn(&T) -> ID, mut neighbors_fn: impl FnMut(&T) -> NI, cycle_fn: impl FnOnce(T) -> E, ) -> Result<Vec<T>, E> where T: Ord, ID: Hash + Eq + Clone, II: IntoIterator<Item = Result<T, E>>, NI: IntoIterator<Item = Result<T, E>>, { struct InnerNode<T> { node: Option<T>, indegree: usize, } // DFS to accumulate incoming edges let mut stack: Vec<T> = start.into_iter().try_collect()?; let mut head_node_map: HashMap<ID, T> = HashMap::new(); let mut inner_node_map: HashMap<ID, InnerNode<T>> = HashMap::new(); let mut neighbor_ids_map: HashMap<ID, Vec<ID>> = HashMap::new(); while let Some(node) = stack.pop() { let node_id = id_fn(&node); if neighbor_ids_map.contains_key(&node_id) { continue; // Already visited