Compute reachable locals as part of non_ssa_locals

1 files changed, 40 insertions, 82 deletions

diff --git a/compiler/rustc_middle/src/mir/traversal.rs b/compiler/rustc_middle/src/mir/traversal.rs
index 1d55a38f30b..b8b74da401c 100644
--- a/compiler/rustc_middle/src/mir/traversal.rs
+++ b/compiler/rustc_middle/src/mir/traversal.rs
@@ -1,5 +1,4 @@
 use super::*;
-use crate::mir::visit::Visitor;
 
 /// Preorder traversal of a graph.
 ///
@@ -105,36 +104,46 @@ impl<'a, 'tcx> Iterator for Preorder<'a, 'tcx> {
 /// ```
 ///
 /// A Postorder traversal of this graph is `D B C A` or `D C B A`
-pub struct Postorder<'a, 'tcx> {
+pub struct Postorder<'a, 'tcx, C> {
     basic_blocks: &'a IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
     visited: BitSet<BasicBlock>,
     visit_stack: Vec<(BasicBlock, Successors<'a>)>,
     root_is_start_block: bool,
+    extra: C,
 }
 
-impl<'a, 'tcx> Postorder<'a, 'tcx> {
+impl<'a, 'tcx, C> Postorder<'a, 'tcx, C>
+where
+    C: Customization<'tcx>,
+{
     pub fn new(
         basic_blocks: &'a IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
         root: BasicBlock,
-    ) -> Postorder<'a, 'tcx> {
+        extra: C,
+    ) -> Postorder<'a, 'tcx, C> {
         let mut po = Postorder {
             basic_blocks,
             visited: BitSet::new_empty(basic_blocks.len()),
             visit_stack: Vec::new(),
             root_is_start_block: root == START_BLOCK,
+            extra,
         };
 
-        let data = &po.basic_blocks[root];
-
-        if let Some(ref term) = data.terminator {
-            po.visited.insert(root);
-            po.visit_stack.push((root, term.successors()));
-            po.traverse_successor();
-        }
+        po.visit(root);
+        po.traverse_successor();
 
         po
     }
 
+    fn visit(&mut self, bb: BasicBlock) {
+        if !self.visited.insert(bb) {
+            return;
+        }
+        let data = &self.basic_blocks[bb];
+        let successors = C::successors(data, self.extra);
+        self.visit_stack.push((bb, successors));
+    }
+
     fn traverse_successor(&mut self) {
         // This is quite a complex loop due to 1. the borrow checker not liking it much
         // and 2. what exactly is going on is not clear
@@ -184,16 +193,15 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> {
         // since we've already visited `E`, that child isn't added to the stack. The last
         // two iterations yield `B` and finally `A` for a final traversal of [E, D, C, B, A]
         while let Some(bb) = self.visit_stack.last_mut().and_then(|(_, iter)| iter.next_back()) {
-            if self.visited.insert(bb) {
-                if let Some(term) = &self.basic_blocks[bb].terminator {
-                    self.visit_stack.push((bb, term.successors()));
-                }
-            }
+            self.visit(bb);
         }
     }
 }
 
-impl<'tcx> Iterator for Postorder<'_, 'tcx> {
+impl<'tcx, C> Iterator for Postorder<'_, 'tcx, C>
+where
+    C: Customization<'tcx>,
+{
     type Item = BasicBlock;
 
     fn next(&mut self) -> Option<BasicBlock> {
@@ -233,73 +241,23 @@ pub fn postorder<'a, 'tcx>(
     reverse_postorder(body).rev()
 }
 
-struct UsedLocals(BitSet<Local>);
-
-impl<'tcx> Visitor<'tcx> for UsedLocals {
-    fn visit_local(
-        &mut self,
-        local: Local,
-        _ctx: crate::mir::visit::PlaceContext,
-        _location: Location,
-    ) {
-        self.0.insert(local);
-    }
-}
-
-struct MonoReachablePostorder<'a, 'tcx> {
-    basic_blocks: &'a IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
-    visited: BitSet<BasicBlock>,
-    visit_stack: Vec<(BasicBlock, Successors<'a>)>,
-    locals: UsedLocals,
-    tcx: TyCtxt<'tcx>,
-    instance: Instance<'tcx>,
+/// Lets us plug in some additional logic and data into a Postorder traversal. Or not.
+pub trait Customization<'tcx>: Copy {
+    fn successors<'a>(_: &'a BasicBlockData<'tcx>, _: Self) -> Successors<'a>;
 }
 
-impl<'a, 'tcx> MonoReachablePostorder<'a, 'tcx> {
-    fn new(
-        body: &'a Body<'tcx>,
-        tcx: TyCtxt<'tcx>,
-        instance: Instance<'tcx>,
-    ) -> MonoReachablePostorder<'a, 'tcx> {
-        let basic_blocks = &body.basic_blocks;
-        let mut po = MonoReachablePostorder {
-            basic_blocks,
-            visited: BitSet::new_empty(basic_blocks.len()),
-            visit_stack: Vec::new(),
-            locals: UsedLocals(BitSet::new_empty(body.local_decls.len())),
-            tcx,
-            instance,
-        };
-
-        po.visit(START_BLOCK);
-        po.traverse_successor();
-        po
-    }
-
-    fn visit(&mut self, bb: BasicBlock) {
-        if !self.visited.insert(bb) {
-            return;
-        }
-        let data = &self.basic_blocks[bb];
-        self.locals.visit_basic_block_data(bb, data);
-        let successors = data.mono_successors(self.tcx, self.instance);
-        self.visit_stack.push((bb, successors));
-    }
-
-    fn traverse_successor(&mut self) {
-        while let Some(bb) = self.visit_stack.last_mut().and_then(|(_, iter)| iter.next_back()) {
-            self.visit(bb);
-        }
+impl<'tcx> Customization<'tcx> for () {
+    fn successors<'a>(data: &'a BasicBlockData<'tcx>, _: ()) -> Successors<'a> {
+        data.terminator().successors()
     }
 }
 
-impl<'tcx> Iterator for MonoReachablePostorder<'_, 'tcx> {
-    type Item = BasicBlock;
-
-    fn next(&mut self) -> Option<BasicBlock> {
-        let (bb, _) = self.visit_stack.pop()?;
-        self.traverse_successor();
-        Some(bb)
+impl<'tcx> Customization<'tcx> for (TyCtxt<'tcx>, Instance<'tcx>) {
+    fn successors<'a>(
+        data: &'a BasicBlockData<'tcx>,
+        (tcx, instance): (TyCtxt<'tcx>, Instance<'tcx>),
+    ) -> Successors<'a> {
+        data.mono_successors(tcx, instance)
     }
 }
 
@@ -307,14 +265,14 @@ pub fn mono_reachable_reverse_postorder<'a, 'tcx>(
     body: &'a Body<'tcx>,
     tcx: TyCtxt<'tcx>,
     instance: Instance<'tcx>,
-) -> (Vec<BasicBlock>, BitSet<Local>) {
-    let mut iter = MonoReachablePostorder::new(body, tcx, instance);
+) -> Vec<BasicBlock> {
+    let mut iter = Postorder::new(&body.basic_blocks, START_BLOCK, (tcx, instance));
     let mut items = Vec::with_capacity(body.basic_blocks.len());
     while let Some(block) = iter.next() {
         items.push(block);
     }
     items.reverse();
-    (items, iter.locals.0)
+    items
 }
 
 /// Returns an iterator over all basic blocks reachable from the `START_BLOCK` in no particular