optimization: use a single DepthFirstSearch instead of hashsets

Extend the `DepthFirstSearch` iterator so that it can be re-used and
extended with add'l start nodes. Then replace the FxHashSets of nodes
we were using in the fallback analysis with a single iterator. This
way we won't re-walk portions of the graph that are reached more than
once, and we also do less allocation etc.

Author: nikomatsakis

compiler/rustc_data_structures/src/graph/iterate/mod.rs

@@ -83,11 +83,58 @@ impl<G> DepthFirstSearch<'graph, G>
 where
     G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
 {
-    pub fn new(graph: &'graph G, start_node: G::Node) -> Self {
-        let mut visited = BitSet::new_empty(graph.num_nodes());
-        visited.insert(start_node);
-        Self { graph, stack: vec![start_node], visited }
+    pub fn new(graph: &'graph G) -> Self {
+        Self { graph, stack: vec![], visited: BitSet::new_empty(graph.num_nodes()) }
+    }
+
+    /// Version of `push_start_node` that is convenient for chained
+    /// use.
+    pub fn with_start_node(mut self, start_node: G::Node) -> Self {
+        self.push_start_node(start_node);
+        self
+    }
+
+    /// Pushes another start node onto the stack. If the node
+    /// has not already been visited, then you will be able to
+    /// walk its successors (and so forth) after the current
+    /// contents of the stack are drained. If multiple start nodes
+    /// are added into the walk, then their mutual successors
+    /// will all be walked. You can use this method once the
+    /// iterator has been completely drained to add additional
+    /// start nodes.
+    pub fn push_start_node(&mut self, start_node: G::Node) {
+        if self.visited.insert(start_node) {
+            self.stack.push(start_node);
+        }
+    }
+
+    /// Searches all nodes reachable from the current start nodes.
+    /// This is equivalent to just invoke `next` repeatedly until
+    /// you get a `None` result.
+    pub fn complete_search(&mut self) {
+        while let Some(_) = self.next() {}
     }
+
+    /// Returns true if node has been visited thus far.
+    /// A node is considered "visited" once it is pushed
+    /// onto the internal stack; it may not yet have been yielded
+    /// from the iterator. This method is best used after
+    /// the iterator is completely drained.
+    pub fn visited(&self, node: G::Node) -> bool {
+        self.visited.contains(node)
+    }
+}
+
+impl<G> std::fmt::Debug for DepthFirstSearch<'_, G>
+where
+    G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
+{
+    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut f = fmt.debug_set();
+        for n in self.visited.iter() {
+            f.entry(&n);
+        }
+        f.finish()
     }
 }
 

compiler/rustc_data_structures/src/graph/iterate/tests.rs

@@ -25,6 +25,14 @@ fn is_cyclic() {
 fn dfs() {
     let graph = TestGraph::new(0, &[(0, 1), (0, 2), (1, 3), (2, 3), (3, 0)]);
 
-    let result: Vec<usize> = DepthFirstSearch::new(&graph, 0).collect();
+    let result: Vec<usize> = DepthFirstSearch::new(&graph).with_start_node(0).collect();
     assert_eq!(result, vec![0, 2, 3, 1]);
 }
+
+#[test]
+fn dfs_debug() {
+    let graph = TestGraph::new(0, &[(0, 1), (0, 2), (1, 3), (2, 3), (3, 0)]);
+    let mut dfs = DepthFirstSearch::new(&graph).with_start_node(0);
+    while let Some(_) = dfs.next() {}
+    assert_eq!(format!("{{0, 1, 2, 3}}"), format!("{:?}", dfs));
+}

compiler/rustc_data_structures/src/graph/mod.rs

@@ -32,7 +32,7 @@ where
     where
         Self: WithNumNodes,
     {
-        iterate::DepthFirstSearch::new(self, from)
+        iterate::DepthFirstSearch::new(self).with_start_node(from)
     }
 }
 

compiler/rustc_typeck/src/check/fallback.rs

@@ -1,6 +1,9 @@
 use crate::check::FnCtxt;
 use rustc_data_structures::{
-    fx::FxHashMap, graph::vec_graph::VecGraph, graph::WithSuccessors, stable_set::FxHashSet,
+    fx::FxHashMap,
+    graph::WithSuccessors,
+    graph::{iterate::DepthFirstSearch, vec_graph::VecGraph},
+    stable_set::FxHashSet,
 };
 use rustc_middle::ty::{self, Ty};
 
@@ -275,7 +278,7 @@ impl<'tcx> FnCtxt<'_, 'tcx> {
         // type variable. These will typically default to `!`, unless
         // we find later that they are *also* reachable from some
         // other type variable outside this set.
-        let mut roots_reachable_from_diverging = FxHashSet::default();
+        let mut roots_reachable_from_diverging = DepthFirstSearch::new(&coercion_graph);
         let mut diverging_vids = vec![];
         let mut non_diverging_vids = vec![];
         for unsolved_vid in unsolved_vids {
@@ -288,16 +291,21 @@ impl<'tcx> FnCtxt<'_, 'tcx> {
             );
             if diverging_roots.contains(&root_vid) {
                 diverging_vids.push(unsolved_vid);
+                roots_reachable_from_diverging.push_start_node(root_vid);
+
                 debug!(
                     "calculate_diverging_fallback: root_vid={:?} reaches {:?}",
                     root_vid,
                     coercion_graph.depth_first_search(root_vid).collect::<Vec<_>>()
                 );
-                roots_reachable_from_diverging.extend(coercion_graph.depth_first_search(root_vid));
+
+                // drain the iterator to visit all nodes reachable from this node
+                roots_reachable_from_diverging.complete_search();
             } else {
                 non_diverging_vids.push(unsolved_vid);
             }
         }
+
         debug!(
             "calculate_diverging_fallback: roots_reachable_from_diverging={:?}",
             roots_reachable_from_diverging,
@@ -307,13 +315,14 @@ impl<'tcx> FnCtxt<'_, 'tcx> {
         // diverging variable, and then compute the set reachable from
         // N0, which we call N. These are the *non-diverging* type
         // variables. (Note that this set consists of "root variables".)
-        let mut roots_reachable_from_non_diverging = FxHashSet::default();
+        let mut roots_reachable_from_non_diverging = DepthFirstSearch::new(&coercion_graph);
         for &non_diverging_vid in &non_diverging_vids {
             let root_vid = self.infcx.root_var(non_diverging_vid);
-            if roots_reachable_from_diverging.contains(&root_vid) {
+            if roots_reachable_from_diverging.visited(root_vid) {
                 continue;
             }
-            roots_reachable_from_non_diverging.extend(coercion_graph.depth_first_search(root_vid));
+            roots_reachable_from_non_diverging.push_start_node(root_vid);
+            roots_reachable_from_non_diverging.complete_search();
         }
         debug!(
             "calculate_diverging_fallback: roots_reachable_from_non_diverging={:?}",
@@ -329,7 +338,7 @@ impl<'tcx> FnCtxt<'_, 'tcx> {
             let root_vid = self.infcx.root_var(diverging_vid);
             let can_reach_non_diverging = coercion_graph
                 .depth_first_search(root_vid)
-                .any(|n| roots_reachable_from_non_diverging.contains(&n));
+                .any(|n| roots_reachable_from_non_diverging.visited(n));
             if can_reach_non_diverging {
                 debug!("fallback to (): {:?}", diverging_vid);
                 diverging_fallback.insert(diverging_ty, self.tcx.types.unit);