diff --git a/src/directed/strongly_connected_components.rs b/src/directed/strongly_connected_components.rs index 9645b644..ad82961b 100644 --- a/src/directed/strongly_connected_components.rs +++ b/src/directed/strongly_connected_components.rs @@ -5,20 +5,29 @@ //! algorithm](https://en.wikipedia.org/wiki/Path-based_strong_component_algorithm) //! is used. -use std::collections::{HashMap, HashSet}; +use rustc_hash::FxHashMap; use std::hash::Hash; +/// Marks a node whose component has already been closed. Preorder numbers count nodes, so a +/// real one can never reach this value. +const ASSIGNED: usize = usize::MAX; + struct Params where N: Hash + Eq, { - preorders: HashMap>, + /// Preorder number of every node the traversal has entered, or [`ASSIGNED`] once the node's + /// component has been emitted. Keeping both states in one map means a successor costs a + /// single lookup instead of one lookup per state. + preorders: FxHashMap, c: usize, successors: FN, - p: Vec, - s: Vec, + /// Preorder numbers of the path nodes whose component is still open. + p: Vec, + /// Nodes entered but not yet assigned to a component, with their preorder numbers. Nodes + /// are pushed as they are entered, so the preorder numbers increase from bottom to top. + s: Vec<(N, usize)>, scc: Vec>, - scca: HashSet, } impl Params @@ -27,55 +36,89 @@ where FN: FnMut(&N) -> IN, IN: IntoIterator, { - fn new(nodes: &[N], successors: FN) -> Self { + fn new(successors: FN) -> Self { Self { - preorders: nodes - .iter() - .map(|n| (n.clone(), None)) - .collect::>>(), + preorders: FxHashMap::default(), c: 0, successors, p: Vec::new(), s: Vec::new(), scc: Vec::new(), - scca: HashSet::new(), } } } -fn recurse_onto(v: &N, params: &mut Params) +/// Explore the graph from `start`, emitting each strongly connected component as it is closed. +/// +/// The traversal keeps its own stack rather than recursing: the depth is bounded only by the +/// number of nodes, which is enough to exhaust the call stack on graphs of very ordinary size. +fn traverse_from(start: &N, params: &mut Params) where N: Clone + Hash + Eq, FN: FnMut(&N) -> IN, IN: IntoIterator, { - params.preorders.insert(v.clone(), Some(params.c)); - params.c += 1; - params.s.push(v.clone()); - params.p.push(v.clone()); - for w in (params.successors)(v) { - if !params.scca.contains(&w) { - if let Some(pw) = params.preorders.get(&w).and_then(|w| *w) { - while params.preorders[¶ms.p[params.p.len() - 1]].unwrap() > pw { - params.p.pop(); + // Each entry is a node's preorder number and the successors of it left to look at. + let mut stack: Vec<(usize, IN::IntoIter)> = Vec::new(); + let pv = params.enter(start.clone()); + stack.push((pv, (params.successors)(start).into_iter())); + + while let Some(top) = stack.last_mut() { + // The borrow of `stack` ends here, so that the body below is free to push onto it. + let successor = top.1.next(); + if let Some(w) = successor { + match params.preorders.get(&w) { + // Already in a component of its own: it cannot be part of this one. + Some(&ASSIGNED) => (), + // Already on the current path, so everything entered after `w` belongs to the + // same component as `w`: those paths can no longer be closed separately. + Some(&pw) => { + while params.p.last().is_some_and(|&p| p > pw) { + params.p.pop(); + } + } + None => { + let pw = params.enter(w.clone()); + stack.push((pw, (params.successors)(&w).into_iter())); } - } else { - recurse_onto(&w, params); } - } - } - if params.p[params.p.len() - 1] == *v { - params.p.pop(); - let mut component = Vec::new(); - while let Some(node) = params.s.pop() { - component.push(node.clone()); - params.scca.insert(node.clone()); - params.preorders.remove(&node); - if node == *v { - break; + } else { + // Every successor has been looked at: if this node still heads an open path, it is + // the root of a component made of everything entered since. + let (pv, _) = stack.pop().unwrap(); + if params.p.last() == Some(&pv) { + params.p.pop(); + // `s` is ordered by preorder number, so the component is exactly its tail. + let first = params.s.partition_point(|&(_, p)| p < pv); + let component = params + .s + .drain(first..) + .map(|(node, _)| node) + .collect::>(); + for node in &component { + // Cannot fail: every node in `s` was given a preorder number on entry. + if let Some(preorder) = params.preorders.get_mut(node) { + *preorder = ASSIGNED; + } + } + params.scc.push(component); } } - params.scc.push(component); + } +} + +impl Params +where + N: Clone + Hash + Eq, +{ + /// Record `node` as entered and return the preorder number it was given. + fn enter(&mut self, node: N) -> usize { + let preorder = self.c; + self.c += 1; + self.preorders.insert(node.clone(), preorder); + self.p.push(preorder); + self.s.push((node, preorder)); + preorder } } @@ -92,8 +135,8 @@ where FN: FnMut(&N) -> IN, IN: IntoIterator, { - let mut params = Params::new(&[], successors); - recurse_onto(start, &mut params); + let mut params = Params::new(successors); + traverse_from(start, &mut params); params.scc } @@ -129,9 +172,11 @@ where FN: FnMut(&N) -> IN, IN: IntoIterator, { - let mut params = Params::new(nodes, successors); - while let Some(node) = params.preorders.keys().find(|_| true).cloned() { - recurse_onto(&node, &mut params); + let mut params = Params::new(successors); + for node in nodes { + if !params.preorders.contains_key(node) { + traverse_from(node, &mut params); + } } params.scc } diff --git a/tests/strongly_connected_components.rs b/tests/strongly_connected_components.rs index 6162cb9c..326414e3 100644 --- a/tests/strongly_connected_components.rs +++ b/tests/strongly_connected_components.rs @@ -119,3 +119,21 @@ fn loops() { c.sort(); assert_eq!(c, vec![vec![0], vec![42]]); } + +#[test] +fn deep_graph_does_not_exhaust_the_stack() { + // A long chain used to need one stack frame per node. The thread is given a deliberately + // small stack so that a return to a recursive traversal fails here rather than silently. + const N: usize = 200_000; + std::thread::Builder::new() + .stack_size(1 << 20) + .spawn(|| { + let nodes = (0..N).collect::>(); + let components = + strongly_connected_components(&nodes, |&i| (i + 1 < N).then_some(i + 1)); + assert_eq!(components.len(), N); + }) + .expect("cannot spawn thread") + .join() + .expect("traversal exhausted the stack"); +}