Fix MIR passes to work with new StatementKind::Assert

Author: Mark-Simulacrum

src/librustc/mir/cache.rs

@@ -17,7 +17,8 @@ use rustc_serialize as serialize;
 
 #[derive(Clone, Debug)]
 pub struct Cache {
-    predecessors: RefCell<Option<IndexVec<Block, Vec<Block>>>>
+    predecessors: RefCell<Option<IndexVec<Block, Vec<Block>>>>,
+    successors: RefCell<Option<IndexVec<Block, Vec<Block>>>>,
 }
 
 
@@ -37,13 +38,15 @@ impl serialize::Decodable for Cache {
 impl Cache {
     pub fn new() -> Self {
         Cache {
-            predecessors: RefCell::new(None)
+            predecessors: RefCell::new(None),
+            successors: RefCell::new(None)
         }
     }
 
     pub fn invalidate(&self) {
         // FIXME: consider being more fine-grained
         *self.predecessors.borrow_mut() = None;
+        *self.successors.borrow_mut() = None;
     }
 
     pub fn predecessors(&self, mir: &Mir) -> Ref<IndexVec<Block, Vec<Block>>> {
@@ -53,11 +56,25 @@ impl Cache {
 
         Ref::map(self.predecessors.borrow(), |p| p.as_ref().unwrap())
     }
+
+    pub fn successors(&self, mir: &Mir) -> Ref<IndexVec<Block, Vec<Block>>> {
+        if self.successors.borrow().is_none() {
+            *self.successors.borrow_mut() = Some(calculate_successors(mir));
+        }
+
+        Ref::map(self.successors.borrow(), |p| p.as_ref().unwrap())
+    }
 }
 
 fn calculate_predecessors(mir: &Mir) -> IndexVec<Block, Vec<Block>> {
     let mut result = IndexVec::from_elem(vec![], mir.basic_blocks());
     for (bb, data) in mir.basic_blocks().iter_enumerated() {
+        for stmt in &data.statements {
+            for &tgt in stmt.successors().iter() {
+                result[tgt].push(bb);
+            }
+        }
+
         if let Some(ref term) = data.terminator {
             for &tgt in term.successors().iter() {
                 result[tgt].push(bb);
@@ -67,3 +84,22 @@ fn calculate_predecessors(mir: &Mir) -> IndexVec<Block, Vec<Block>> {
 
     result
 }
+
+fn calculate_successors(mir: &Mir) -> IndexVec<Block, Vec<Block>> {
+    let mut result = IndexVec::from_elem(vec![], mir.basic_blocks());
+    for (bb, data) in mir.basic_blocks().iter_enumerated() {
+        for stmt in &data.statements {
+            for &tgt in stmt.successors().iter() {
+                result[bb].push(tgt);
+            }
+        }
+
+        if let Some(ref term) = data.terminator {
+            for &tgt in term.successors().iter() {
+                result[bb].push(tgt);
+            }
+        }
+    }
+
+    result
+}

src/librustc/mir/mod.rs

@@ -162,11 +162,21 @@ impl<'tcx> Mir<'tcx> {
         self.cache.predecessors(self)
     }
 
+    #[inline]
+    pub fn successors(&self) -> Ref<IndexVec<Block, Vec<Block>>> {
+        self.cache.successors(self)
+    }
+
     #[inline]
     pub fn predecessors_for(&self, bb: Block) -> Ref<Vec<Block>> {
         Ref::map(self.predecessors(), |p| &p[bb])
     }
 
+    #[inline]
+    pub fn successors_for(&self, bb: Block) -> Ref<Vec<Block>> {
+        Ref::map(self.successors(), |p| &p[bb])
+    }
+
     #[inline]
     pub fn dominators(&self) -> Dominators<Block> {
         dominators(self)
@@ -724,6 +734,24 @@ impl<'tcx> Statement<'tcx> {
     pub fn make_nop(&mut self) {
         self.kind = StatementKind::Nop
     }
+
+    pub fn successors(&self) -> Cow<[Block]> {
+        match self.kind {
+            StatementKind::Assert { cleanup: Some(unwind), .. } => {
+                vec![unwind].into_cow()
+            }
+            _ => (&[]).into_cow(),
+        }
+    }
+
+    pub fn successors_mut(&mut self) -> Vec<&mut Block> {
+        match self.kind {
+            StatementKind::Assert { cleanup: Some(ref mut unwind), .. } => {
+                vec![unwind]
+            }
+            _ => Vec::new(),
+        }
+    }
 }
 
 #[derive(Clone, Debug, RustcEncodable, RustcDecodable)]
@@ -1302,7 +1330,7 @@ impl<'tcx> ControlFlowGraph for Mir<'tcx> {
     fn successors<'graph>(&'graph self, node: Self::Node)
                           -> <Self as GraphSuccessors<'graph>>::Iter
     {
-        self.basic_blocks[node].terminator().successors().into_owned().into_iter()
+        self.successors_for(node).clone().into_iter()
     }
 }
 

src/librustc/mir/traversal.rs

@@ -64,15 +64,11 @@ impl<'a, 'tcx> Iterator for Preorder<'a, 'tcx> {
                 continue;
             }
 
-            let data = &self.mir[idx];
-
-            if let Some(ref term) = data.terminator {
-                for &succ in term.successors().iter() {
-                    self.worklist.push(succ);
-                }
+            for &succ in self.mir.successors_for(idx).iter() {
+                self.worklist.push(succ);
             }
 
-            return Some((idx, data));
+            return Some((idx, &self.mir[idx]));
         }
 
         None
@@ -111,17 +107,11 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> {
             visit_stack: Vec::new()
         };
 
+        po.visited.insert(root.index());
 
-        let data = &po.mir[root];
-
-        if let Some(ref term) = data.terminator {
-            po.visited.insert(root.index());
-
-            let succs = term.successors().into_owned().into_iter();
-
-            po.visit_stack.push((root, succs));
-            po.traverse_successor();
-        }
+        let succs = ControlFlowGraph::successors(&po.mir, root);
+        po.visit_stack.push((root, succs));
+        po.traverse_successor();
 
         po
     }
@@ -186,10 +176,7 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> {
             };
 
             if self.visited.insert(bb.index()) {
-                if let Some(ref term) = self.mir[bb].terminator {
-                    let succs = term.successors().into_owned().into_iter();
-                    self.visit_stack.push((bb, succs));
-                }
+                self.visit_stack.push((bb, ControlFlowGraph::successors(&self.mir, bb)));
             }
         }
     }

src/librustc_borrowck/borrowck/mir/dataflow/graphviz.rs

@@ -125,7 +125,7 @@ pub type Node = Block;
 pub struct Edge { source: Block, index: usize }
 
 fn outgoing(mir: &Mir, bb: Block) -> Vec<Edge> {
-    let succ_len = mir[bb].terminator().successors().len();
+    let succ_len = mir.successors_for(bb).len();
     (0..succ_len).map(|index| Edge { source: bb, index: index}).collect()
 }
 

src/librustc_borrowck/borrowck/mir/dataflow/impls.rs

@@ -474,6 +474,7 @@ impl<'a, 'tcx> BitDenotation for MovingOutStatements<'a, 'tcx> {
             mir::StatementKind::StorageLive(_) |
             mir::StatementKind::StorageDead(_) |
             mir::StatementKind::InlineAsm { .. } |
+            mir::StatementKind::Assert { .. } |
             mir::StatementKind::Nop => {}
         }
     }

src/librustc_borrowck/borrowck/mir/dataflow/mod.rs

@@ -439,14 +439,12 @@ impl<'a, 'tcx: 'a, D> DataflowAnalysis<'a, 'tcx, D>
             mir::TerminatorKind::Resume |
             mir::TerminatorKind::Unreachable => {}
             mir::TerminatorKind::Goto { ref target } |
-            mir::TerminatorKind::Assert { ref target, cleanup: None, .. } |
             mir::TerminatorKind::Drop { ref target, location: _, unwind: None } |
             mir::TerminatorKind::DropAndReplace {
                 ref target, value: _, location: _, unwind: None
             } => {
                 self.propagate_bits_into_entry_set_for(in_out, changed, target);
             }
-            mir::TerminatorKind::Assert { ref target, cleanup: Some(ref unwind), .. } |
             mir::TerminatorKind::Drop { ref target, location: _, unwind: Some(ref unwind) } |
             mir::TerminatorKind::DropAndReplace {
                 ref target, value: _, location: _, unwind: Some(ref unwind)

src/librustc_borrowck/borrowck/mir/dataflow/sanity_check.rs

@@ -105,6 +105,7 @@ fn each_block<'a, 'tcx, O>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
             mir::StatementKind::StorageLive(_) |
             mir::StatementKind::StorageDead(_) |
             mir::StatementKind::InlineAsm { .. } |
+            mir::StatementKind::Assert { .. } |
             mir::StatementKind::Nop => continue,
             mir::StatementKind::SetDiscriminant{ .. } =>
                 span_bug!(stmt.source_info.span,

src/librustc_borrowck/borrowck/mir/gather_moves.rs

@@ -413,6 +413,7 @@ impl<'a, 'tcx> MoveDataBuilder<'a, 'tcx> {
                           "SetDiscriminant should not exist during borrowck");
             }
             StatementKind::InlineAsm { .. } |
+            StatementKind::Assert { .. } |
             StatementKind::Nop => {}
         }
     }
@@ -464,7 +465,6 @@ impl<'a, 'tcx> MoveDataBuilder<'a, 'tcx> {
                 self.gather_move(loc, &Lvalue::Local(RETURN_POINTER));
             }
 
-            TerminatorKind::Assert { .. } |
             TerminatorKind::SwitchInt { .. } => {
                 // branching terminators - these don't move anything
             }

src/librustc_borrowck/borrowck/mir/mod.rs

@@ -364,6 +364,7 @@ fn drop_flag_effects_for_location<'a, 'tcx, F>(
             mir::StatementKind::StorageLive(_) |
             mir::StatementKind::StorageDead(_) |
             mir::StatementKind::InlineAsm { .. } |
+            mir::StatementKind::Assert { .. } |
             mir::StatementKind::Nop => {}
         },
         None => {

src/librustc_mir/build/expr/as_lvalue.rs

@@ -75,15 +75,16 @@ impl<'a, 'gcx, 'tcx> Builder<'a, 'gcx, 'tcx> {
                     len: Operand::Consume(len),
                     index: idx.clone()
                 };
-                this.cfg.push(block, Statement {
-                    source_info: this.source_info(expr_span),
+                let stmt = Statement {
+                    source_info: source_info,
                     kind: StatementKind::Assert {
                         cond: Operand::Consume(lt),
                         expected: true,
                         msg: msg,
                         cleanup: this.diverge_cleanup(),
                     }
-                });
+                };
+                this.cfg.push(block, stmt);
                 block.and(slice.index(idx))
             }
             ExprKind::SelfRef => {