add a `depth_first_search` helper function

1 files changed, 35 insertions, 1 deletions

diff --git a/src/librustc_data_structures/graph/iterate/mod.rs b/src/librustc_data_structures/graph/iterate/mod.rs
index c09364b0a53..5612778ce07 100644
--- a/src/librustc_data_structures/graph/iterate/mod.rs
+++ b/src/librustc_data_structures/graph/iterate/mod.rs
@@ -1,5 +1,6 @@
 use super::super::indexed_vec::IndexVec;
-use super::{DirectedGraph, WithSuccessors, WithNumNodes};
+use super::{DirectedGraph, WithNumNodes, WithSuccessors};
+use crate::bit_set::BitSet;
 
 #[cfg(test)]
 mod test;
@@ -51,3 +52,36 @@ pub fn reverse_post_order<G: DirectedGraph + WithSuccessors + WithNumNodes>(
     vec.reverse();
     vec
 }
+
+/// A "depth-first search" iterator for a directed graph.
+pub struct DepthFirstSearch<'graph, G>
+where
+    G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
+{
+    graph: &'graph G,
+    stack: Vec<G::Node>,
+    visited: BitSet<G::Node>,
+}
+
+impl<G> DepthFirstSearch<'graph, G>
+where
+    G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
+{
+    pub fn new(graph: &'graph G, start_node: G::Node) -> Self {
+        Self { graph, stack: vec![start_node], visited: BitSet::new_empty(graph.num_nodes()) }
+    }
+}
+
+impl<G> Iterator for DepthFirstSearch<'_, G>
+where
+    G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
+{
+    type Item = G::Node;
+
+    fn next(&mut self) -> Option<G::Node> {
+        let DepthFirstSearch { stack, visited, graph } = self;
+        let n = stack.pop()?;
+        stack.extend(graph.successors(n).filter(|&m| visited.insert(m)));
+        Some(n)
+    }
+}