perf(scc): traverse iteratively and look each successor up once

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<N, FN>
 where
     N: Hash + Eq,
 {
-    preorders: HashMap<N, Option<usize>>,
+    /// 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<N, usize>,
     c: usize,
     successors: FN,
-    p: Vec<N>,
-    s: Vec<N>,
+    /// Preorder numbers of the path nodes whose component is still open.
+    p: Vec<usize>,
+    /// 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<Vec<N>>,
-    scca: HashSet<N>,
 }
 
 impl<N, FN, IN> Params<N, FN>
@@ -27,55 +36,89 @@ where
     FN: FnMut(&N) -> IN,
     IN: IntoIterator<Item = N>,
 {
-    fn new(nodes: &[N], successors: FN) -> Self {
+    fn new(successors: FN) -> Self {
         Self {
-            preorders: nodes
-                .iter()
-                .map(|n| (n.clone(), None))
-                .collect::<HashMap<N, Option<usize>>>(),
+            preorders: FxHashMap::default(),
             c: 0,
             successors,
             p: Vec::new(),
             s: Vec::new(),
             scc: Vec::new(),
-            scca: HashSet::new(),
         }
     }
 }
 
-fn recurse_onto<N, FN, IN>(v: &N, params: &mut Params<N, FN>)
+/// 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<N, FN, IN>(start: &N, params: &mut Params<N, FN>)
 where
     N: Clone + Hash + Eq,
     FN: FnMut(&N) -> IN,
     IN: IntoIterator<Item = N>,
 {
-    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[&params.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::<Vec<_>>();
+                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<N, FN> Params<N, FN>
+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<Item = N>,
 {
-    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<Item = N>,
 {
-    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
 }

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::<Vec<_>>();
+            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");
+}