From 7fd0db7dd319cfb73664c8a068474dc8759ebabf Mon Sep 17 00:00:00 2001
From: Niko Matsakis <niko@alum.mit.edu>
Date: Tue, 11 Jun 2019 16:29:27 -0400
Subject: [PATCH] add a `depth_first_search` helper function

---
 .../graph/iterate/mod.rs                      | 36 ++++++++++++++++++-
 src/librustc_data_structures/graph/mod.rs     |  7 ++++
 .../graph/vec_graph/test.rs                   |  7 ++++
 3 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/src/librustc_data_structures/graph/iterate/mod.rs b/src/librustc_data_structures/graph/iterate/mod.rs
index c09364b0a539..5612778ce07e 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)
+    }
+}
diff --git a/src/librustc_data_structures/graph/mod.rs b/src/librustc_data_structures/graph/mod.rs
index 45e7e5db38f0..2787fa3c6b1e 100644
--- a/src/librustc_data_structures/graph/mod.rs
+++ b/src/librustc_data_structures/graph/mod.rs
@@ -30,6 +30,13 @@ where
         &'graph self,
         node: Self::Node,
     ) -> <Self as GraphSuccessors<'graph>>::Iter;
+
+    fn depth_first_search(&self, from: Self::Node) -> iterate::DepthFirstSearch<'_, Self>
+    where
+        Self: WithNumNodes,
+    {
+        iterate::DepthFirstSearch::new(self, from)
+    }
 }
 
 pub trait GraphSuccessors<'graph> {
diff --git a/src/librustc_data_structures/graph/vec_graph/test.rs b/src/librustc_data_structures/graph/vec_graph/test.rs
index 4a168ee1d44f..c54a72264f66 100644
--- a/src/librustc_data_structures/graph/vec_graph/test.rs
+++ b/src/librustc_data_structures/graph/vec_graph/test.rs
@@ -44,3 +44,10 @@ fn succesors() {
     assert_eq!(graph.successors(5), &[1]);
     assert_eq!(graph.successors(6), &[]);
 }
+
+#[test]
+fn dfs() {
+    let graph = create_graph();
+    let dfs: Vec<_> = graph.depth_first_search(0).collect();
+    assert_eq!(dfs, vec![0, 1, 3, 4, 2]);
+}