From efe634f14c20661216c57322b319c1cac1f9704f Mon Sep 17 00:00:00 2001
From: Dylan MacKenzie <ecstaticmorse@gmail.com>
Date: Wed, 8 Jul 2020 09:47:14 -0700
Subject: [PATCH] Add `reachable` and friends to `mir::traversal` module

---
 src/librustc_middle/mir/traversal.rs          | 17 +++++++++++++++++
 src/librustc_mir/dataflow/framework/engine.rs |  9 +++++++++
 src/librustc_mir/transform/generator.rs       |  4 +---
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/src/librustc_middle/mir/traversal.rs b/src/librustc_middle/mir/traversal.rs
index ed8129b1e09a5..36e277d1a88f3 100644
--- a/src/librustc_middle/mir/traversal.rs
+++ b/src/librustc_middle/mir/traversal.rs
@@ -292,3 +292,20 @@ impl<'a, 'tcx> Iterator for ReversePostorder<'a, 'tcx> {
 }
 
 impl<'a, 'tcx> ExactSizeIterator for ReversePostorder<'a, 'tcx> {}
+
+/// Returns an iterator over all basic blocks reachable from the `START_BLOCK` in no particular
+/// order.
+///
+/// This is clearer than writing `preorder` in cases where the order doesn't matter.
+pub fn reachable<'a, 'tcx>(
+    body: &'a Body<'tcx>,
+) -> impl 'a + Iterator<Item = (BasicBlock, &'a BasicBlockData<'tcx>)> {
+    preorder(body)
+}
+
+/// Returns a `BitSet` containing all basic blocks reachable from the `START_BLOCK`.
+pub fn reachable_as_bitset(body: &Body<'tcx>) -> BitSet<BasicBlock> {
+    let mut iter = preorder(body);
+    (&mut iter).for_each(drop);
+    iter.visited
+}
diff --git a/src/librustc_mir/dataflow/framework/engine.rs b/src/librustc_mir/dataflow/framework/engine.rs
index 243b3679f2984..eb38811791168 100644
--- a/src/librustc_mir/dataflow/framework/engine.rs
+++ b/src/librustc_mir/dataflow/framework/engine.rs
@@ -52,6 +52,15 @@ where
         visit_results(body, blocks, self, vis)
     }
 
+    pub fn visit_reachable_with(
+        &self,
+        body: &'mir mir::Body<'tcx>,
+        vis: &mut impl ResultsVisitor<'mir, 'tcx, FlowState = BitSet<A::Idx>>,
+    ) {
+        let blocks = mir::traversal::reachable(body);
+        visit_results(body, blocks.map(|(bb, _)| bb), self, vis)
+    }
+
     pub fn visit_in_rpo_with(
         &self,
         body: &'mir mir::Body<'tcx>,
diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index 523d3c9af3f68..8618cc126c563 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -624,9 +624,7 @@ fn compute_storage_conflicts(
         local_conflicts: BitMatrix::from_row_n(&ineligible_locals, body.local_decls.len()),
     };
 
-    // Visit only reachable basic blocks. The exact order is not important.
-    let reachable_blocks = traversal::preorder(body).map(|(bb, _)| bb);
-    requires_storage.visit_with(body, reachable_blocks, &mut visitor);
+    requires_storage.visit_reachable_with(body, &mut visitor);
 
     let local_conflicts = visitor.local_conflicts;