fix(count_paths): walk an explicit stack so deep graphs do not overflow

src/directed/count_paths.rs

@@ -2,40 +2,93 @@
 
 use std::hash::Hash;
 
-use rustc_hash::FxHashMap;
+use crate::FxIndexMap;
+use indexmap::map::Entry::{Occupied, Vacant};
 
-fn cached_count_paths<T, FN, IN, FS>(
-    start: T,
+/// Account for `node` in the search.
+///
+/// Returns the number of paths from it when that is already settled. Otherwise the node is
+/// recorded as being worked on and a frame is pushed for it, and `None` is returned; note that
+/// the stack is therefore only ever pushed to when this returns `None`.
+///
+/// # Panics
+///
+/// If `node` is reached again while its own count is still unknown, it lies on a loop.
+fn enter<T, FN, IN, FS>(
+    node: T,
+    counts: &mut FxIndexMap<T, Option<usize>>,
+    stack: &mut Vec<(usize, IN::IntoIter, usize)>,
     successors: &mut FN,
     success: &mut FS,
-    cache: &mut FxHashMap<T, usize>,
-) -> usize
+) -> Option<usize>
 where
     T: Eq + Hash,
     FN: FnMut(&T) -> IN,
     IN: IntoIterator<Item = T>,
     FS: FnMut(&T) -> bool,
 {
-    if let Some(&n) = cache.get(&start) {
-        return n;
+    match counts.entry(node) {
+        Occupied(e) => match *e.get() {
+            Some(count) => Some(count),
+            None => panic!("the graph given to count_paths contains a loop"),
+        },
+        Vacant(e) => {
+            if success(e.key()) {
+                e.insert(Some(1));
+                Some(1)
+            } else {
+                let index = e.index();
+                let successors = successors(e.key()).into_iter();
+                e.insert(None);
+                stack.push((index, successors, 0));
+                None
+            }
+        }
     }
+}
 
-    let count = if success(&start) {
-        1
-    } else {
-        successors(&start)
-            .into_iter()
-            .map(|successor| cached_count_paths(successor, successors, success, cache))
-            .sum()
-    };
+fn count_paths_from<T, FN, IN, FS>(start: T, successors: &mut FN, success: &mut FS) -> usize
+where
+    T: Eq + Hash,
+    FN: FnMut(&T) -> IN,
+    IN: IntoIterator<Item = T>,
+    FS: FnMut(&T) -> bool,
+{
+    // The nodes live here rather than in the stack, which refers to them by index, so that `T`
+    // need not be `Clone`. A `None` count marks a node that is still being worked out.
+    let mut counts: FxIndexMap<T, Option<usize>> = FxIndexMap::default();
+    // Each frame is a node, the successors of it left to look at, and the number of paths
+    // found through the ones already done.
+    let mut stack: Vec<(usize, IN::IntoIter, usize)> = Vec::new();
 
-    cache.insert(start, count);
+    if let Some(count) = enter(start, &mut counts, &mut stack, successors, success) {
+        return count;
+    }
 
-    count
+    let mut total = 0;
+    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(successor) = successor {
+            if let Some(count) = enter(successor, &mut counts, &mut stack, successors, success) {
+                // Nothing was pushed, so the frame on top is still the one being counted.
+                stack.last_mut().unwrap().2 += count;
+            }
+        } else {
+            // Every successor is accounted for, so this node's own count is settled.
+            let (index, _, count) = stack.pop().unwrap();
+            *counts.get_index_mut(index).unwrap().1 = Some(count);
+            match stack.last_mut() {
+                Some(parent) => parent.2 += count,
+                // The starting node, which is the last frame to be popped.
+                None => total = count,
+            }
+        }
+    }
+    total
 }
 
-/// Count the total number of possible paths to reach a destination. There must be no loops
-/// in the graph, or the function will overflow its stack.
+/// Count the total number of possible paths to reach a destination.
 ///
 /// # Example
 ///
@@ -55,17 +108,16 @@ where
 /// );
 /// assert_eq!(n, 3432);
 /// ```
+///
+/// # Panics
+///
+/// If the graph contains a loop, since the number of paths through it is then unbounded.
 pub fn count_paths<T, FN, IN, FS>(start: T, mut successors: FN, mut success: FS) -> usize
 where
     T: Eq + Hash,
     FN: FnMut(&T) -> IN,
     IN: IntoIterator<Item = T>,
     FS: FnMut(&T) -> bool,
 {
-    cached_count_paths(
-        start,
-        &mut successors,
-        &mut success,
-        &mut FxHashMap::default(),
-    )
+    count_paths_from(start, &mut successors, &mut success)
 }

tests/count_paths.rs

@@ -13,3 +13,47 @@ fn grid() {
     );
     assert_eq!(n, 3432);
 }
+
+#[test]
+fn deep_graph_does_not_exhaust_the_stack() {
+    // A single chain: one path, but formerly one stack frame per node. The thread is given a
+    // deliberately small stack so that a return to a recursive walk fails here rather than
+    // silently on someone else's machine.
+    const N: usize = 200_000;
+    std::thread::Builder::new()
+        .stack_size(1 << 20)
+        .spawn(|| {
+            let n = count_paths(0usize, |&n| (n + 1 < N).then_some(n + 1), |&n| n == N - 1);
+            assert_eq!(n, 1);
+        })
+        .expect("cannot spawn thread")
+        .join()
+        .expect("counting exhausted the stack");
+}
+
+#[test]
+fn counts_are_shared_between_paths() {
+    // Every node of a diamond lattice is reachable by several routes; the count is only
+    // correct, and only cheap, if each node is counted once and reused.
+    let mut visited = 0;
+    let n = count_paths(
+        (0u32, 0u32),
+        |&(x, y)| {
+            visited += 1;
+            [(x + 1, y), (x, y + 1)]
+                .into_iter()
+                .filter(|&(x, y)| x < 20 && y < 20)
+        },
+        |&c| c == (19, 19),
+    );
+    // Central binomial coefficient C(38, 19).
+    assert_eq!(n, 35_345_263_800);
+    // One expansion per node, not one per path.
+    assert!(visited <= 400, "successors called {visited} times");
+}
+
+#[test]
+#[should_panic(expected = "loop")]
+fn a_loop_is_reported() {
+    count_paths(0u32, |&n| [(n + 1) % 4], |_| false);
+}