From 5f9a33efeb3f7ac3ce26120f590953ba09ce924f Mon Sep 17 00:00:00 2001 From: tachsin Date: Wed, 9 Sep 2026 00:11:55 +0300 Subject: [PATCH] perf(topological_sort): traverse iteratively and drop the redundant sets `visit` recurses once per node, so a deep graph overflows the stack: a chain of 200 000 nodes is enough, and so is a random DAG of 60 000. Walk an explicit stack instead. The bookkeeping around it was doing the same work several times over: - `marked` and `temp` are nested: a node is in `temp` from the moment it is reached and in `marked` once it is finished, so `marked` is a subset of `temp` and every node was hashed into both. One map from node to "finished?" answers both questions in a single lookup. - The outer loop repeatedly took an arbitrary key out of a `HashSet` of roots that `visit` was emptying as it went, so each call rescanned the buckets left behind by the ones before it. Walking the caller's slice needs no set at all, and with it goes a clone of every root. - Both sets used the default hasher, unlike the rest of the crate. `topological_sort_into_groups` keeps its structure and just moves to the same hasher. `successors` is still called exactly once per node, which the existing `complexity` test checks. Orders are checked for validity against 200 random DAGs. On 60 000 nodes: about 76% off `topological_sort`, and about 33% off `topological_sort_into_groups`. --- src/directed/topological_sort.rs | 78 +++++++++++++++++++------------- tests/topological_sort.rs | 17 +++++++ 2 files changed, 64 insertions(+), 31 deletions(-) diff --git a/src/directed/topological_sort.rs b/src/directed/topological_sort.rs index f15f190d..0c758eca 100644 --- a/src/directed/topological_sort.rs +++ b/src/directed/topological_sort.rs @@ -1,6 +1,7 @@ //! Find a topological order in a directed graph if one exists. -use std::collections::{HashMap, HashSet, VecDeque}; +use rustc_hash::{FxHashMap, FxHashSet}; +use std::collections::VecDeque; use std::hash::Hash; use std::mem; @@ -79,30 +80,27 @@ where FN: FnMut(&N) -> IN, IN: IntoIterator, { - let mut marked = HashSet::with_capacity(roots.len()); - let mut temp = HashSet::new(); + let mut visited = FxHashMap::default(); let mut sorted = VecDeque::with_capacity(roots.len()); - let mut roots: HashSet = roots.iter().cloned().collect::>(); - while let Some(node) = roots.iter().next().cloned() { - temp.clear(); - visit( - &node, - &mut successors, - &mut roots, - &mut marked, - &mut temp, - &mut sorted, - )?; + for root in roots { + visit(root, &mut successors, &mut visited, &mut sorted)?; } Ok(sorted.into_iter().collect()) } +/// Explore the graph below `node` in depth-first order, prepending each node to `sorted` once +/// everything reachable from it has been placed. +/// +/// `visited` maps every node reached so far to whether it is finished. A node that has been +/// reached but is not finished is still on the path being explored, so meeting it again closes +/// a cycle. +/// +/// The traversal keeps its own stack rather than recursing, since the depth of a graph is easily +/// enough to exhaust the call stack. fn visit( - node: &N, + start: &N, successors: &mut FN, - unmarked: &mut HashSet, - marked: &mut HashSet, - temp: &mut HashSet, + visited: &mut FxHashMap, sorted: &mut VecDeque, ) -> Result<(), N> where @@ -110,19 +108,37 @@ where FN: FnMut(&N) -> IN, IN: IntoIterator, { - unmarked.remove(node); - if marked.contains(node) { - return Ok(()); + match visited.get(start) { + Some(true) => return Ok(()), + Some(false) => return Err(start.clone()), + None => (), } - if temp.contains(node) { - return Err(node.clone()); - } - temp.insert(node.clone()); - for n in successors(node) { - visit(&n, successors, unmarked, marked, temp, sorted)?; + visited.insert(start.clone(), false); + let mut stack: Vec<(N, IN::IntoIter)> = vec![(start.clone(), 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(node) = successor { + match visited.get(&node) { + // Finished: everything below it is already in place. + Some(true) => (), + // Reached but not finished, so it is still on the path being explored. + Some(false) => return Err(node), + None => { + visited.insert(node.clone(), false); + let successors = successors(&node).into_iter(); + stack.push((node, successors)); + } + } + } else { + // Everything reachable from this node has been placed, so it comes before all of it. + let (node, _) = stack.pop().unwrap(); + if let Some(finished) = visited.get_mut(&node) { + *finished = true; + } + sorted.push_front(node); + } } - marked.insert(node.clone()); - sorted.push_front(node.clone()); Ok(()) } @@ -166,8 +182,8 @@ where if nodes.is_empty() { return Ok(Vec::new()); } - let mut succs_map = HashMap::>::with_capacity(nodes.len()); - let mut preds_map = HashMap::::with_capacity(nodes.len()); + let mut succs_map = FxHashMap::>::default(); + let mut preds_map = FxHashMap::::default(); for node in nodes { succs_map.insert(node.clone(), successors(node).into_iter().collect()); preds_map.insert(node.clone(), 0); diff --git a/tests/topological_sort.rs b/tests/topological_sort.rs index 8499a9df..76acd97a 100644 --- a/tests/topological_sort.rs +++ b/tests/topological_sort.rs @@ -114,3 +114,20 @@ fn tsig_self_edge() { Err((vec![vec![0], vec![1, 2]], vec![3])) ); } + +#[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 sorted = tsort(&nodes, |&i| (i + 1 < N).then_some(i + 1)).expect("cycle reported"); + assert_eq!(sorted, nodes); + }) + .expect("cannot spawn thread") + .join() + .expect("traversal exhausted the stack"); +}