Stop invalidating predecessors cache when accessing unique basic block, invalidate cache when accessing unique terminator

3 files changed, 41 insertions, 8 deletions

diff --git a/src/librustc/mir/mod.rs b/src/librustc/mir/mod.rs
index d0903e04d16..0d07caa6cda 100644
--- a/src/librustc/mir/mod.rs
+++ b/src/librustc/mir/mod.rs
@@ -203,17 +203,38 @@ impl<'tcx> Body<'tcx> {
 
     #[inline]
     pub fn basic_blocks_mut(&mut self) -> &mut IndexVec<BasicBlock, BasicBlockData<'tcx>> {
-        self.predecessors_cache = None;
-//        self.cache.invalidate();
         &mut self.basic_blocks
     }
 
+    pub fn basic_block_terminator_opt_mut(&mut self, bb: BasicBlock) -> &mut Option<Terminator<'tcx>> {
+        self.predecessors_cache = None;
+        &mut self.basic_blocks[bb].terminator
+    }
+
+    pub fn basic_block_terminator_mut(&mut self, bb: BasicBlock) -> &mut Terminator<'tcx> {
+        self.predecessors_cache = None;
+/*
+        let data = &mut self.basic_blocks[bb];
+        if let Some(cache) = self.predecessors_cache.as_mut() {
+            for successor in data.terminator().successors() {
+                let successor_vec = &mut cache[*successor];
+                for i in (0..successor_vec.len()).rev() {
+                    if successor_vec[i] == bb {
+                        successor_vec.swap_remove(i);
+                        break;
+                    }
+                }
+            }
+        }
+*/
+
+        self.basic_blocks[bb].terminator_mut()
+    }
+
     #[inline]
     pub fn basic_blocks_and_local_decls_mut(
         &mut self,
     ) -> (&mut IndexVec<BasicBlock, BasicBlockData<'tcx>>, &mut LocalDecls<'tcx>) {
-        self.predecessors_cache = None;
-//        self.cache.invalidate();
         (&mut self.basic_blocks, &mut self.local_decls)
     }
 
@@ -1358,6 +1379,10 @@ impl<'tcx> BasicBlockData<'tcx> {
         BasicBlockData { statements: vec![], terminator, is_cleanup: false }
     }
 
+    pub fn terminator_opt(&self) -> &Option<Terminator<'tcx>> {
+        &self.terminator
+    }
+
     /// Accessor for terminator.
     ///
     /// Terminator may not be None after construction of the basic block is complete. This accessor
@@ -1366,10 +1391,17 @@ impl<'tcx> BasicBlockData<'tcx> {
         self.terminator.as_ref().expect("invalid terminator state")
     }
 
-    pub fn terminator_mut(&mut self) -> &mut Terminator<'tcx> {
+    // This cannot be public since changing the terminator will break the predecessors cache in Body
+    // To do so outside of this module, use Body::basic_block_terminator_mut(BasicBlock)
+    fn terminator_mut(&mut self) -> &mut Terminator<'tcx> {
         self.terminator.as_mut().expect("invalid terminator state")
     }
 
+    // This can be public since changing the kind will not break the predecessors cache in Body
+    pub fn terminator_kind_mut(&mut self) -> &mut TerminatorKind<'tcx> {
+        &mut self.terminator_mut().kind
+    }
+
     pub fn retain_statements<F>(&mut self, mut f: F)
     where
         F: FnMut(&mut Statement<'_>) -> bool,
diff --git a/src/librustc/mir/traversal.rs b/src/librustc/mir/traversal.rs
index f129dd3abef..917f1fae52a 100644
--- a/src/librustc/mir/traversal.rs
+++ b/src/librustc/mir/traversal.rs
@@ -55,7 +55,7 @@ impl<'a, 'tcx> Iterator for Preorder<'a, 'tcx> {
 
             let data = &self.body[idx];
 
-            if let Some(ref term) = data.terminator {
+            if let Some(ref term) = data.terminator_opt() {
                 self.worklist.extend(term.successors());
             }
 
@@ -117,7 +117,7 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> {
 
         let data = &po.body[root];
 
-        if let Some(ref term) = data.terminator {
+        if let Some(ref term) = data.terminator_opt() {
             po.visited.insert(root);
             po.visit_stack.push((root, term.successors()));
             po.traverse_successor();
@@ -186,7 +186,7 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> {
             };
 
             if self.visited.insert(bb) {
-                if let Some(term) = &self.body[bb].terminator {
+                if let Some(term) = &self.body[bb].terminator_opt() {
                     self.visit_stack.push((bb, term.successors()));
                 }
             }
diff --git a/src/librustc/mir/visit.rs b/src/librustc/mir/visit.rs
index 145593f1c4d..fe146528dae 100644
--- a/src/librustc/mir/visit.rs
+++ b/src/librustc/mir/visit.rs
@@ -793,6 +793,7 @@ macro_rules! make_mir_visitor {
             fn visit_location(&mut self, body: & $($mutability)? Body<'tcx>, location: Location) {
                 let basic_block = & $($mutability)? body[location.block];
                 if basic_block.statements.len() == location.statement_index {
+                    // TODO(nashenas88) how to ensure we clear the cache only in the mutable case...
                     if let Some(ref $($mutability)? terminator) = basic_block.terminator {
                         self.visit_terminator(terminator, location)
                     }