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 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)
+ }
+}
diff --git a/src/librustc_data_structures/graph/mod.rs b/src/librustc_data_structures/graph/mod.rs
index 45e7e5db38f..2787fa3c6b1 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 4a168ee1d44..c54a72264f6 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]);
+}