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

17 files changed, 112 insertions, 73 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)
                     }
diff --git a/src/librustc_codegen_ssa/mir/mod.rs b/src/librustc_codegen_ssa/mir/mod.rs
index 6041232489d..66ebe4095dc 100644
--- a/src/librustc_codegen_ssa/mir/mod.rs
+++ b/src/librustc_codegen_ssa/mir/mod.rs
@@ -265,7 +265,7 @@ fn create_funclets<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
 
         let funclet;
         let ret_llbb;
-        match mir[bb].terminator.as_ref().map(|t| &t.kind) {
+        match mir[bb].terminator_opt().as_ref().map(|t| &t.kind) {
             // This is a basic block that we're aborting the program for,
             // notably in an `extern` function. These basic blocks are inserted
             // so that we assert that `extern` functions do indeed not panic,
diff --git a/src/librustc_mir/borrow_check/error_reporting.rs b/src/librustc_mir/borrow_check/error_reporting.rs
index a555e0b74c2..140ee2583fa 100644
--- a/src/librustc_mir/borrow_check/error_reporting.rs
+++ b/src/librustc_mir/borrow_check/error_reporting.rs
@@ -497,7 +497,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                             ..
                         },
                         ..
-                    }) = bbd.terminator {
+                    }) = bbd.terminator_opt() {
                         if let Some(source)
                             = BorrowedContentSource::from_call(func.ty(self.body, tcx), tcx)
                         {
diff --git a/src/librustc_mir/build/cfg.rs b/src/librustc_mir/build/cfg.rs
index 3ed6b4ff346..734135c324a 100644
--- a/src/librustc_mir/build/cfg.rs
+++ b/src/librustc_mir/build/cfg.rs
@@ -64,7 +64,7 @@ impl<'tcx> CFG<'tcx> {
                      source_info: SourceInfo,
                      kind: TerminatorKind<'tcx>) {
         debug!("terminating block {:?} <- {:?}", block, kind);
-        debug_assert!(self.block_data(block).terminator.is_none(),
+        debug_assert!(self.block_data(block).terminator_opt().is_none(),
                       "terminate: block {:?}={:?} already has a terminator set",
                       block,
                       self.block_data(block));
diff --git a/src/librustc_mir/build/mod.rs b/src/librustc_mir/build/mod.rs
index eb9b401f272..1e4b94632fa 100644
--- a/src/librustc_mir/build/mod.rs
+++ b/src/librustc_mir/build/mod.rs
@@ -731,7 +731,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
     fn finish(self) -> Body<'tcx> {
         for (index, block) in self.cfg.basic_blocks.iter().enumerate() {
-            if block.terminator.is_none() {
+            if block.terminator_opt().is_none() {
                 span_bug!(self.fn_span, "no terminator on block {:?}", index);
             }
         }
diff --git a/src/librustc_mir/dataflow/impls/borrows.rs b/src/librustc_mir/dataflow/impls/borrows.rs
index 402e5aeacbf..f2c0b714961 100644
--- a/src/librustc_mir/dataflow/impls/borrows.rs
+++ b/src/librustc_mir/dataflow/impls/borrows.rs
@@ -98,7 +98,7 @@ fn precompute_borrows_out_of_scope<'tcx>(
             // Add successor BBs to the work list, if necessary.
             let bb_data = &body[bb];
             assert!(hi == bb_data.statements.len());
-            for &succ_bb in bb_data.terminator.as_ref().unwrap().successors() {
+            for &succ_bb in bb_data.terminator().successors() {
                 visited.entry(succ_bb)
                     .and_modify(|lo| {
                         // `succ_bb` has been seen before. If it wasn't
diff --git a/src/librustc_mir/lints.rs b/src/librustc_mir/lints.rs
index 158b730b9bd..70c1d120629 100644
--- a/src/librustc_mir/lints.rs
+++ b/src/librustc_mir/lints.rs
@@ -79,7 +79,7 @@ fn check_fn_for_unconditional_recursion(
 
         let block = &basic_blocks[bb];
 
-        if let Some(ref terminator) = block.terminator {
+        if let Some(ref terminator) = block.terminator_opt() {
             match terminator.kind {
                 TerminatorKind::Call { ref func, .. } => {
                     let func_ty = func.ty(body, tcx);
diff --git a/src/librustc_mir/transform/add_call_guards.rs b/src/librustc_mir/transform/add_call_guards.rs
index bf3df1ae2fd..46e6df82318 100644
--- a/src/librustc_mir/transform/add_call_guards.rs
+++ b/src/librustc_mir/transform/add_call_guards.rs
@@ -46,8 +46,9 @@ impl AddCallGuards {
 
         let cur_len = body.basic_blocks().len();
 
-        for block in body.basic_blocks_mut() {
-            match block.terminator {
+        for bb in body.basic_blocks().indices() {
+            let is_cleanup = body.basic_blocks()[bb].is_cleanup;
+            match body.basic_block_terminator_opt_mut(bb) {
                 Some(Terminator {
                     kind: TerminatorKind::Call {
                         destination: Some((_, ref mut destination)),
@@ -60,9 +61,9 @@ impl AddCallGuards {
                     // It's a critical edge, break it
                     let call_guard = BasicBlockData {
                         statements: vec![],
-                        is_cleanup: block.is_cleanup,
+                        is_cleanup: is_cleanup,
                         terminator: Some(Terminator {
-                            source_info,
+                            source_info: *source_info,
                             kind: TerminatorKind::Goto { target: *destination }
                         })
                     };
diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index 5d88629435b..4c090d1b065 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -368,7 +368,7 @@ impl MutVisitor<'tcx> for TransformVisitor<'tcx> {
                 VariantIdx::new(RETURNED) // state for returned
             };
             data.statements.push(self.set_discr(state, source_info));
-            data.terminator.as_mut().unwrap().kind = TerminatorKind::Return;
+            *data.terminator_kind_mut() = TerminatorKind::Return;
         }
 
         self.super_basic_block_data(block, data);
@@ -852,10 +852,10 @@ fn insert_switch<'tcx>(
         is_cleanup: false,
     });
 
-    let blocks = body.basic_blocks_mut().iter_mut();
-
-    for target in blocks.flat_map(|b| b.terminator_mut().successors_mut()) {
-        *target = BasicBlock::new(target.index() + 1);
+    for bb in body.basic_blocks_mut().indices() {
+        for target in body.basic_block_terminator_mut(bb).successors_mut() {
+            *target = BasicBlock::new(target.index() + 1);
+        }
     }
 }
 
@@ -941,7 +941,7 @@ fn create_generator_drop_shim<'tcx>(
     insert_switch(&mut body, cases, &transform, TerminatorKind::Return);
 
     for block in body.basic_blocks_mut() {
-        let kind = &mut block.terminator_mut().kind;
+        let kind = block.terminator_kind_mut();
         if let TerminatorKind::GeneratorDrop = *kind {
             *kind = TerminatorKind::Return;
         }
diff --git a/src/librustc_mir/transform/inline.rs b/src/librustc_mir/transform/inline.rs
index ebfadd0cfd3..0ae07f081ae 100644
--- a/src/librustc_mir/transform/inline.rs
+++ b/src/librustc_mir/transform/inline.rs
@@ -379,7 +379,7 @@ impl Inliner<'tcx> {
                    callsite: CallSite<'tcx>,
                    caller_body: &mut Body<'tcx>,
                    mut callee_body: Body<'tcx>) -> bool {
-        let terminator = caller_body[callsite.bb].terminator.take().unwrap();
+        let terminator = caller_body.basic_block_terminator_opt_mut(callsite.bb).take().unwrap();
         match terminator.kind {
             // FIXME: Handle inlining of diverging calls
             TerminatorKind::Call { args, destination: Some(destination), cleanup, .. } => {
@@ -496,12 +496,12 @@ impl Inliner<'tcx> {
                     kind: TerminatorKind::Goto { target: BasicBlock::new(bb_len) }
                 };
 
-                caller_body[callsite.bb].terminator = Some(terminator);
+                *caller_body.basic_block_terminator_opt_mut(callsite.bb) = Some(terminator);
 
                 true
             }
             kind => {
-                caller_body[callsite.bb].terminator = Some(Terminator {
+                *caller_body.basic_block_terminator_opt_mut(callsite.bb) = Some(Terminator {
                     source_info: terminator.source_info,
                     kind,
                 });
diff --git a/src/librustc_mir/transform/promote_consts.rs b/src/librustc_mir/transform/promote_consts.rs
index 591f3ca4400..a569dc65228 100644
--- a/src/librustc_mir/transform/promote_consts.rs
+++ b/src/librustc_mir/transform/promote_consts.rs
@@ -866,7 +866,7 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
             let terminator = if self.keep_original {
                 self.source[loc.block].terminator().clone()
             } else {
-                let terminator = self.source[loc.block].terminator_mut();
+                let terminator = self.source.basic_block_terminator_mut(loc.block);
                 let target = match terminator.kind {
                     TerminatorKind::Call { destination: Some((_, target)), .. } => target,
                     ref kind => {
@@ -891,7 +891,7 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
                     let last = self.promoted.basic_blocks().last().unwrap();
                     let new_target = self.new_block();
 
-                    *self.promoted[last].terminator_mut() = Terminator {
+                    *self.promoted.basic_block_terminator_mut(last) = Terminator {
                         kind: TerminatorKind::Call {
                             func,
                             args,
@@ -976,11 +976,12 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
                     }
                 },
                 Candidate::Argument { bb, index } => {
-                    let terminator = blocks[bb].terminator_mut();
-                    match terminator.kind {
+                    let data = &mut blocks[bb];
+                    let terminator_span = data.terminator().source_info.span;
+                    match data.terminator_kind_mut() {
                         TerminatorKind::Call { ref mut args, .. } => {
                             let ty = args[index].ty(local_decls, self.tcx);
-                            let span = terminator.source_info.span;
+                            let span = terminator_span;
                             let operand = Operand::Copy(promoted_place(ty, span));
                             mem::replace(&mut args[index], operand)
                         }
@@ -1103,8 +1104,8 @@ pub fn promote_candidates<'tcx>(
 
     // Eliminate assignments to, and drops of promoted temps.
     let promoted = |index: Local| temps[index] == TempState::PromotedOut;
-    for block in body.basic_blocks_mut() {
-        block.statements.retain(|statement| {
+    for bb in body.basic_blocks().indices() {
+        body.basic_blocks_mut()[bb].statements.retain(|statement| {
             match &statement.kind {
                 StatementKind::Assign(box(place, _)) => {
                     if let Some(index) = place.as_local() {
@@ -1120,7 +1121,7 @@ pub fn promote_candidates<'tcx>(
                 _ => true
             }
         });
-        let terminator = block.terminator_mut();
+        let terminator = body.basic_block_terminator_mut(bb);
         match &terminator.kind {
             TerminatorKind::Drop { location: place, target, .. } => {
                 if let Some(index) = place.as_local() {
diff --git a/src/librustc_mir/transform/remove_noop_landing_pads.rs b/src/librustc_mir/transform/remove_noop_landing_pads.rs
index 130393e2c4c..9e0f80dcf9a 100644
--- a/src/librustc_mir/transform/remove_noop_landing_pads.rs
+++ b/src/librustc_mir/transform/remove_noop_landing_pads.rs
@@ -103,7 +103,7 @@ impl RemoveNoopLandingPads {
         let postorder: Vec<_> = traversal::postorder(body).map(|(bb, _)| bb).collect();
         for bb in postorder {
             debug!("  processing {:?}", bb);
-            for target in body[bb].terminator_mut().successors_mut() {
+            for target in body.basic_block_terminator_mut(bb).successors_mut() {
                 if *target != resume_block && nop_landing_pads.contains(*target) {
                     debug!("    folding noop jump to {:?} to resume block", target);
                     *target = resume_block;
@@ -111,7 +111,7 @@ impl RemoveNoopLandingPads {
                 }
             }
 
-            match body[bb].terminator_mut().unwind_mut() {
+            match body.basic_block_terminator_mut(bb).unwind_mut() {
                 Some(unwind) => {
                     if *unwind == Some(resume_block) {
                         debug!("    removing noop landing pad");
diff --git a/src/librustc_mir/transform/simplify.rs b/src/librustc_mir/transform/simplify.rs
index c9185d14148..5f5014d6d46 100644
--- a/src/librustc_mir/transform/simplify.rs
+++ b/src/librustc_mir/transform/simplify.rs
@@ -63,7 +63,7 @@ impl<'tcx> MirPass<'tcx> for SimplifyCfg {
 }
 
 pub struct CfgSimplifier<'a, 'tcx> {
-    basic_blocks: &'a mut IndexVec<BasicBlock, BasicBlockData<'tcx>>,
+    body: &'a mut Body<'tcx>,
     pred_count: IndexVec<BasicBlock, u32>
 }
 
@@ -76,21 +76,23 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
         pred_count[START_BLOCK] = 1;
 
         for (_, data) in traversal::preorder(body) {
-            if let Some(ref term) = data.terminator {
+            if let Some(ref term) = data.terminator_opt() {
                 for &tgt in term.successors() {
                     pred_count[tgt] += 1;
                 }
             }
         }
 
-        let basic_blocks = body.basic_blocks_mut();
-
         CfgSimplifier {
-            basic_blocks,
+            body,
             pred_count,
         }
     }
 
+    fn basic_blocks(&mut self) -> &mut IndexVec<BasicBlock, BasicBlockData<'tcx>> {
+        self.body.basic_blocks_mut()
+    }
+
     pub fn simplify(mut self) {
         self.strip_nops();
 
@@ -101,14 +103,14 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
 
             self.collapse_goto_chain(&mut start, &mut changed);
 
-            for bb in self.basic_blocks.indices() {
+            for bb in self.body.basic_blocks().indices() {
                 if self.pred_count[bb] == 0 {
                     continue
                 }
 
                 debug!("simplifying {:?}", bb);
 
-                let mut terminator = self.basic_blocks[bb].terminator.take()
+                let mut terminator = self.body.basic_block_terminator_opt_mut(bb).take()
                     .expect("invalid terminator state");
 
                 for successor in terminator.successors_mut() {
@@ -123,9 +125,8 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
                     inner_changed |= self.merge_successor(&mut new_stmts, &mut terminator);
                     changed |= inner_changed;
                 }
-
-                self.basic_blocks[bb].statements.extend(new_stmts);
-                self.basic_blocks[bb].terminator = Some(terminator);
+                self.basic_blocks()[bb].statements.extend(new_stmts);
+                *self.body.basic_block_terminator_opt_mut(bb) = Some(terminator);
 
                 changed |= inner_changed;
             }
@@ -135,17 +136,17 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
 
         if start != START_BLOCK {
             debug_assert!(self.pred_count[START_BLOCK] == 0);
-            self.basic_blocks.swap(START_BLOCK, start);
+            self.basic_blocks().swap(START_BLOCK, start);
             self.pred_count.swap(START_BLOCK, start);
 
             // pred_count == 1 if the start block has no predecessor _blocks_.
             if self.pred_count[START_BLOCK] > 1 {
-                for (bb, data) in self.basic_blocks.iter_enumerated_mut() {
+                for bb in self.basic_blocks().indices() {
                     if self.pred_count[bb] == 0 {
                         continue;
                     }
 
-                    for target in data.terminator_mut().successors_mut() {
+                    for target in self.body.basic_block_terminator_mut(bb).successors_mut() {
                         if *target == start {
                             *target = START_BLOCK;
                         }
@@ -157,7 +158,7 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
 
     // Collapse a goto chain starting from `start`
     fn collapse_goto_chain(&mut self, start: &mut BasicBlock, changed: &mut bool) {
-        let mut terminator = match self.basic_blocks[*start] {
+        let mut terminator = match self.basic_blocks()[*start] {
             BasicBlockData {
                 ref statements,
                 terminator: ref mut terminator @ Some(Terminator {
@@ -176,7 +177,7 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
             }
             _ => unreachable!()
         };
-        self.basic_blocks[*start].terminator = terminator;
+        self.basic_blocks()[*start].terminator = terminator;
 
         debug!("collapsing goto chain from {:?} to {:?}", *start, target);
 
@@ -208,7 +209,7 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
         };
 
         debug!("merging block {:?} into {:?}", target, terminator);
-        *terminator = match self.basic_blocks[target].terminator.take() {
+        *terminator = match self.body.basic_block_terminator_opt_mut(target).take() {
             Some(terminator) => terminator,
             None => {
                 // unreachable loop - this should not be possible, as we
@@ -216,7 +217,7 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
                 return false
             }
         };
-        new_stmts.extend(self.basic_blocks[target].statements.drain(..));
+        new_stmts.extend(self.basic_blocks()[target].statements.drain(..));
         self.pred_count[target] = 0;
 
         true
@@ -249,7 +250,7 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
     }
 
     fn strip_nops(&mut self) {
-        for blk in self.basic_blocks.iter_mut() {
+        for blk in self.basic_blocks().iter_mut() {
             blk.statements.retain(|stmt| if let StatementKind::Nop = stmt.kind {
                 false
             } else {
@@ -265,24 +266,27 @@ pub fn remove_dead_blocks(body: &mut Body<'_>) {
         seen.insert(bb.index());
     }
 
-    let basic_blocks = body.basic_blocks_mut();
-
-    let num_blocks = basic_blocks.len();
-    let mut replacements : Vec<_> = (0..num_blocks).map(BasicBlock::new).collect();
-    let mut used_blocks = 0;
-    for alive_index in seen.iter() {
-        replacements[alive_index] = BasicBlock::new(used_blocks);
-        if alive_index != used_blocks {
-            // Swap the next alive block data with the current available slot. Since alive_index is
-            // non-decreasing this is a valid operation.
-            basic_blocks.raw.swap(alive_index, used_blocks);
+    let mut replacements: Vec<BasicBlock>;
+    {
+        let basic_blocks = body.basic_blocks_mut();
+
+        let num_blocks = basic_blocks.len();
+        replacements = (0..num_blocks).map(BasicBlock::new).collect();
+        let mut used_blocks = 0;
+        for alive_index in seen.iter() {
+            replacements[alive_index] = BasicBlock::new(used_blocks);
+            if alive_index != used_blocks {
+                // Swap the next alive block data with the current available slot. Since alive_index is
+                // non-decreasing this is a valid operation.
+                basic_blocks.raw.swap(alive_index, used_blocks);
+            }
+            used_blocks += 1;
         }
-        used_blocks += 1;
+        basic_blocks.raw.truncate(used_blocks);
     }
-    basic_blocks.raw.truncate(used_blocks);
 
-    for block in basic_blocks {
-        for target in block.terminator_mut().successors_mut() {
+    for bb in body.basic_blocks().indices() {
+        for target in body.basic_block_terminator_mut(bb).successors_mut() {
             *target = replacements[target.index()];
         }
     }
diff --git a/src/librustc_mir/transform/simplify_branches.rs b/src/librustc_mir/transform/simplify_branches.rs
index 0a509666d34..1e45ad5d428 100644
--- a/src/librustc_mir/transform/simplify_branches.rs
+++ b/src/librustc_mir/transform/simplify_branches.rs
@@ -21,8 +21,8 @@ impl<'tcx> MirPass<'tcx> for SimplifyBranches {
 
     fn run_pass(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, body: &mut Body<'tcx>) {
         let param_env = tcx.param_env(src.def_id());
-        for block in body.basic_blocks_mut() {
-            let terminator = block.terminator_mut();
+        for bb in body.basic_blocks().indices() {
+            let terminator = body.basic_block_terminator_mut(bb);
             terminator.kind = match terminator.kind {
                 TerminatorKind::SwitchInt {
                     discr: Operand::Constant(ref c), switch_ty, ref values, ref targets, ..
diff --git a/src/librustc_mir/util/patch.rs b/src/librustc_mir/util/patch.rs
index a5f7e540157..554ae2d3c6b 100644
--- a/src/librustc_mir/util/patch.rs
+++ b/src/librustc_mir/util/patch.rs
@@ -141,7 +141,7 @@ impl<'tcx> MirPatch<'tcx> {
         for (src, patch) in self.patch_map.into_iter_enumerated() {
             if let Some(patch) = patch {
                 debug!("MirPatch: patching block {:?}", src);
-                body[src].terminator_mut().kind = patch;
+                body.basic_block_terminator_mut(src).kind = patch;
             }
         }