From 9923c580c3146ba59d9404fc147f622131350064 Mon Sep 17 00:00:00 2001 From: tachsin Date: Wed, 9 Sep 2026 20:54:07 +0300 Subject: [PATCH] fix(count_paths): walk an explicit stack so deep graphs do not overflow `cached_count_paths` recursed once per node, so the depth of the recursion was the length of the longest path. A chain of 200 000 nodes overflows the stack, and that graph has exactly one path and no loop anywhere in it, although the documentation implies a loop-free graph is safe. Walk an explicit stack instead. The nodes are held in an `FxIndexMap` that the stack refers to by index, which keeps `T: Clone` off the bounds and makes the node's state part of the same lookup that reads its count. That state also settles what a loop does now. There is no stack left to overflow, so an undetected loop would instead grow the working set until memory ran out, which is a worse failure than the documented one. A node reached again while its own count is still unknown lies on a loop by definition, so it is recognised at no extra cost and reported as a panic naming the problem. The documentation changes accordingly, from a note that loops overflow the stack to a `# Panics` section. Counting itself is unchanged: each node is still expanded once and its count reused, which the added test pins down by counting the calls. --- src/directed/count_paths.rs | 102 +++++++++++++++++++++++++++--------- tests/count_paths.rs | 44 ++++++++++++++++ 2 files changed, 121 insertions(+), 25 deletions(-) diff --git a/src/directed/count_paths.rs b/src/directed/count_paths.rs index 29b4533f..747647ad 100644 --- a/src/directed/count_paths.rs +++ b/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( - 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( + node: T, + counts: &mut FxIndexMap>, + stack: &mut Vec<(usize, IN::IntoIter, usize)>, successors: &mut FN, success: &mut FS, - cache: &mut FxHashMap, -) -> usize +) -> Option where T: Eq + Hash, FN: FnMut(&T) -> IN, IN: IntoIterator, 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(start: T, successors: &mut FN, success: &mut FS) -> usize +where + T: Eq + Hash, + FN: FnMut(&T) -> IN, + IN: IntoIterator, + 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> = 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,6 +108,10 @@ 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(start: T, mut successors: FN, mut success: FS) -> usize where T: Eq + Hash, @@ -62,10 +119,5 @@ where IN: IntoIterator, FS: FnMut(&T) -> bool, { - cached_count_paths( - start, - &mut successors, - &mut success, - &mut FxHashMap::default(), - ) + count_paths_from(start, &mut successors, &mut success) } diff --git a/tests/count_paths.rs b/tests/count_paths.rs index cf68c8a2..debbd5ab 100644 --- a/tests/count_paths.rs +++ b/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); +}