Auto merge of #50278 - eddyb:mir-succ-iter, r=nikomatsakis

rustc: return iterators from Terminator(Kind)::successors(_mut).

Minor cleanup (and potentially speedup) prompted by @nnethercote's `SmallVec` experiments.
This PR assumes `.count()` and `.nth(i)` on `iter::Chain<option::IntoIter, slice::Iter(Mut)>` are `O(1)`, but otherwise all of the uses appear to immediately iterate through the successors.

r? @nikomatsakis

Author: bors

src/librustc/mir/cache.rs

@@ -68,7 +68,7 @@ fn calculate_predecessors(mir: &Mir) -> IndexVec<BasicBlock, Vec<BasicBlock>> {
     let mut result = IndexVec::from_elem(vec![], mir.basic_blocks());
     for (bb, data) in mir.basic_blocks().iter_enumerated() {
         if let Some(ref term) = data.terminator {
-            for &tgt in term.successors().iter() {
+            for &tgt in term.successors() {
                 result[tgt].push(bb);
             }
         }

src/librustc/mir/mod.rs

@@ -36,7 +36,7 @@ use hir::{self, InlineAsm};
 use std::borrow::{Cow};
 use rustc_data_structures::sync::ReadGuard;
 use std::fmt::{self, Debug, Formatter, Write};
-use std::{iter, mem, u32};
+use std::{iter, mem, option, u32};
 use std::ops::{Index, IndexMut};
 use std::vec::IntoIter;
 use syntax::ast::{self, Name};
@@ -862,12 +862,17 @@ pub enum TerminatorKind<'tcx> {
     },
 }
 
+pub type Successors<'a> =
+    iter::Chain<option::IntoIter<&'a BasicBlock>, slice::Iter<'a, BasicBlock>>;
+pub type SuccessorsMut<'a> =
+    iter::Chain<option::IntoIter<&'a mut BasicBlock>, slice::IterMut<'a, BasicBlock>>;
+
 impl<'tcx> Terminator<'tcx> {
-    pub fn successors(&self) -> Cow<[BasicBlock]> {
+    pub fn successors(&self) -> Successors {
         self.kind.successors()
     }
 
-    pub fn successors_mut(&mut self) -> Vec<&mut BasicBlock> {
+    pub fn successors_mut(&mut self) -> SuccessorsMut {
         self.kind.successors_mut()
     }
 
@@ -888,72 +893,71 @@ impl<'tcx> TerminatorKind<'tcx> {
         }
     }
 
-    pub fn successors(&self) -> Cow<[BasicBlock]> {
+    pub fn successors(&self) -> Successors {
         use self::TerminatorKind::*;
         match *self {
-            Goto { target: ref b } => slice::from_ref(b).into_cow(),
-            SwitchInt { targets: ref b, .. } => b[..].into_cow(),
-            Resume | Abort | GeneratorDrop => (&[]).into_cow(),
-            Return => (&[]).into_cow(),
-            Unreachable => (&[]).into_cow(),
-            Call { destination: Some((_, t)), cleanup: Some(c), .. } => vec![t, c].into_cow(),
-            Call { destination: Some((_, ref t)), cleanup: None, .. } =>
-                slice::from_ref(t).into_cow(),
-            Call { destination: None, cleanup: Some(ref c), .. } => slice::from_ref(c).into_cow(),
-            Call { destination: None, cleanup: None, .. } => (&[]).into_cow(),
-            Yield { resume: t, drop: Some(c), .. } => vec![t, c].into_cow(),
-            Yield { resume: ref t, drop: None, .. } => slice::from_ref(t).into_cow(),
-            DropAndReplace { target, unwind: Some(unwind), .. } |
-            Drop { target, unwind: Some(unwind), .. } => {
-                vec![target, unwind].into_cow()
+            Resume | Abort | GeneratorDrop | Return | Unreachable |
+            Call { destination: None, cleanup: None, .. } => {
+                None.into_iter().chain(&[])
+            }
+            Goto { target: ref t } |
+            Call { destination: None, cleanup: Some(ref t), .. } |
+            Call { destination: Some((_, ref t)), cleanup: None, .. } |
+            Yield { resume: ref t, drop: None, .. } |
+            DropAndReplace { target: ref t, unwind: None, .. } |
+            Drop { target: ref t, unwind: None, .. } |
+            Assert { target: ref t, cleanup: None, .. } |
+            FalseUnwind { real_target: ref t, unwind: None } => {
+                Some(t).into_iter().chain(&[])
             }
-            DropAndReplace { ref target, unwind: None, .. } |
-            Drop { ref target, unwind: None, .. } => {
-                slice::from_ref(target).into_cow()
+            Call { destination: Some((_, ref t)), cleanup: Some(ref u), .. } |
+            Yield { resume: ref t, drop: Some(ref u), .. } |
+            DropAndReplace { target: ref t, unwind: Some(ref u), .. } |
+            Drop { target: ref t, unwind: Some(ref u), .. } |
+            Assert { target: ref t, cleanup: Some(ref u), .. } |
+            FalseUnwind { real_target: ref t, unwind: Some(ref u) } => {
+                Some(t).into_iter().chain(slice::from_ref(u))
+            }
+            SwitchInt { ref targets, .. } => {
+                None.into_iter().chain(&targets[..])
             }
-            Assert { target, cleanup: Some(unwind), .. } => vec![target, unwind].into_cow(),
-            Assert { ref target, .. } => slice::from_ref(target).into_cow(),
             FalseEdges { ref real_target, ref imaginary_targets } => {
-                let mut s = vec![*real_target];
-                s.extend_from_slice(imaginary_targets);
-                s.into_cow()
+                Some(real_target).into_iter().chain(&imaginary_targets[..])
             }
-            FalseUnwind { real_target: t, unwind: Some(u) } => vec![t, u].into_cow(),
-            FalseUnwind { real_target: ref t, unwind: None } => slice::from_ref(t).into_cow(),
         }
     }
 
-    // FIXME: no mootable cow. I’m honestly not sure what a “cow” between `&mut [BasicBlock]` and
-    // `Vec<&mut BasicBlock>` would look like in the first place.
-    pub fn successors_mut(&mut self) -> Vec<&mut BasicBlock> {
+    pub fn successors_mut(&mut self) -> SuccessorsMut {
         use self::TerminatorKind::*;
         match *self {
-            Goto { target: ref mut b } => vec![b],
-            SwitchInt { targets: ref mut b, .. } => b.iter_mut().collect(),
-            Resume | Abort | GeneratorDrop => Vec::new(),
-            Return => Vec::new(),
-            Unreachable => Vec::new(),
-            Call { destination: Some((_, ref mut t)), cleanup: Some(ref mut c), .. } => vec![t, c],
-            Call { destination: Some((_, ref mut t)), cleanup: None, .. } => vec![t],
-            Call { destination: None, cleanup: Some(ref mut c), .. } => vec![c],
-            Call { destination: None, cleanup: None, .. } => vec![],
-            Yield { resume: ref mut t, drop: Some(ref mut c), .. } => vec![t, c],
-            Yield { resume: ref mut t, drop: None, .. } => vec![t],
-            DropAndReplace { ref mut target, unwind: Some(ref mut unwind), .. } |
-            Drop { ref mut target, unwind: Some(ref mut unwind), .. } => vec![target, unwind],
-            DropAndReplace { ref mut target, unwind: None, .. } |
-            Drop { ref mut target, unwind: None, .. } => {
-                vec![target]
+            Resume | Abort | GeneratorDrop | Return | Unreachable |
+            Call { destination: None, cleanup: None, .. } => {
+                None.into_iter().chain(&mut [])
+            }
+            Goto { target: ref mut t } |
+            Call { destination: None, cleanup: Some(ref mut t), .. } |
+            Call { destination: Some((_, ref mut t)), cleanup: None, .. } |
+            Yield { resume: ref mut t, drop: None, .. } |
+            DropAndReplace { target: ref mut t, unwind: None, .. } |
+            Drop { target: ref mut t, unwind: None, .. } |
+            Assert { target: ref mut t, cleanup: None, .. } |
+            FalseUnwind { real_target: ref mut t, unwind: None } => {
+                Some(t).into_iter().chain(&mut [])
+            }
+            Call { destination: Some((_, ref mut t)), cleanup: Some(ref mut u), .. } |
+            Yield { resume: ref mut t, drop: Some(ref mut u), .. } |
+            DropAndReplace { target: ref mut t, unwind: Some(ref mut u), .. } |
+            Drop { target: ref mut t, unwind: Some(ref mut u), .. } |
+            Assert { target: ref mut t, cleanup: Some(ref mut u), .. } |
+            FalseUnwind { real_target: ref mut t, unwind: Some(ref mut u) } => {
+                Some(t).into_iter().chain(slice::from_ref_mut(u))
+            }
+            SwitchInt { ref mut targets, .. } => {
+                None.into_iter().chain(&mut targets[..])
             }
-            Assert { ref mut target, cleanup: Some(ref mut unwind), .. } => vec![target, unwind],
-            Assert { ref mut target, .. } => vec![target],
             FalseEdges { ref mut real_target, ref mut imaginary_targets } => {
-                let mut s = vec![real_target];
-                s.extend(imaginary_targets.iter_mut());
-                s
+                Some(real_target).into_iter().chain(&mut imaginary_targets[..])
             }
-            FalseUnwind { real_target: ref mut t, unwind: Some(ref mut u) } => vec![t, u],
-            FalseUnwind { ref mut real_target, unwind: None } => vec![real_target],
         }
     }
 
@@ -1073,18 +1077,18 @@ impl<'tcx> BasicBlockData<'tcx> {
 impl<'tcx> Debug for TerminatorKind<'tcx> {
     fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
         self.fmt_head(fmt)?;
-        let successors = self.successors();
+        let successor_count = self.successors().count();
         let labels = self.fmt_successor_labels();
-        assert_eq!(successors.len(), labels.len());
+        assert_eq!(successor_count, labels.len());
 
-        match successors.len() {
+        match successor_count {
             0 => Ok(()),
 
-            1 => write!(fmt, " -> {:?}", successors[0]),
+            1 => write!(fmt, " -> {:?}", self.successors().nth(0).unwrap()),
 
             _ => {
                 write!(fmt, " -> [")?;
-                for (i, target) in successors.iter().enumerate() {
+                for (i, target) in self.successors().enumerate() {
                     if i > 0 {
                         write!(fmt, ", ")?;
                     }
@@ -1943,7 +1947,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.basic_blocks[node].terminator().successors().cloned()
     }
 }
 
@@ -1954,7 +1958,7 @@ impl<'a, 'b> GraphPredecessors<'b> for Mir<'a> {
 
 impl<'a, 'b>  GraphSuccessors<'b> for Mir<'a> {
     type Item = BasicBlock;
-    type Iter = IntoIter<BasicBlock>;
+    type Iter = iter::Cloned<Successors<'b>>;
 }
 
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Ord, PartialOrd)]

src/librustc/mir/traversal.rs

@@ -8,8 +8,6 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-use std::vec;
-
 use rustc_data_structures::bitvec::BitVector;
 use rustc_data_structures::indexed_vec::Idx;
 
@@ -67,7 +65,7 @@ impl<'a, 'tcx> Iterator for Preorder<'a, 'tcx> {
             let data = &self.mir[idx];
 
             if let Some(ref term) = data.terminator {
-                for &succ in term.successors().iter() {
+                for &succ in term.successors() {
                     self.worklist.push(succ);
                 }
             }
@@ -110,7 +108,7 @@ impl<'a, 'tcx> ExactSizeIterator for Preorder<'a, 'tcx> {}
 pub struct Postorder<'a, 'tcx: 'a> {
     mir: &'a Mir<'tcx>,
     visited: BitVector,
-    visit_stack: Vec<(BasicBlock, vec::IntoIter<BasicBlock>)>
+    visit_stack: Vec<(BasicBlock, Successors<'a>)>
 }
 
 impl<'a, 'tcx> Postorder<'a, 'tcx> {
@@ -126,10 +124,7 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> {
 
         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.visit_stack.push((root, term.successors()));
             po.traverse_successor();
         }
 
@@ -186,7 +181,7 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> {
         // two iterations yield `C` and finally `A` for a final traversal of [E, D, B, C, A]
         loop {
             let bb = if let Some(&mut (_, ref mut iter)) = self.visit_stack.last_mut() {
-                if let Some(bb) = iter.next() {
+                if let Some(&bb) = iter.next() {
                     bb
                 } else {
                     break;
@@ -197,8 +192,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, term.successors()));
                 }
             }
         }

src/librustc_mir/borrow_check/nll/explain_borrow/mod.rs

@@ -193,7 +193,6 @@ impl<'gcx, 'tcx> UseFinder<'gcx, 'tcx> {
                         block_data
                             .terminator()
                             .successors()
-                            .iter()
                             .map(|&basic_block| Location {
                                 statement_index: 0,
                                 block: basic_block,

src/librustc_mir/borrow_check/nll/region_infer/dfs.rs

@@ -95,7 +95,6 @@ impl<'tcx> RegionInferenceContext<'tcx> {
                     block_data
                         .terminator()
                         .successors()
-                        .iter()
                         .map(|&basic_block| Location {
                             statement_index: 0,
                             block: basic_block,

src/librustc_mir/dataflow/graphviz.rs

@@ -73,8 +73,8 @@ pub type Node = BasicBlock;
 pub struct Edge { source: BasicBlock, index: usize }
 
 fn outgoing(mir: &Mir, bb: BasicBlock) -> Vec<Edge> {
-    let succ_len = mir[bb].terminator().successors().len();
-    (0..succ_len).map(|index| Edge { source: bb, index: index}).collect()
+    mir[bb].terminator().successors().enumerate()
+        .map(|(index, _)| Edge { source: bb, index: index}).collect()
 }
 
 impl<'a, 'tcx, MWF, P> dot::Labeller<'a> for Graph<'a, 'tcx, MWF, P>
@@ -285,6 +285,6 @@ impl<'a, 'tcx, MWF, P> dot::GraphWalk<'a> for Graph<'a, 'tcx, MWF, P>
 
     fn target(&self, edge: &Edge) -> Node {
         let mir = self.mbcx.mir();
-        mir[edge.source].terminator().successors()[edge.index]
+        *mir[edge.source].terminator().successors().nth(edge.index).unwrap()
     }
 }

src/librustc_mir/transform/inline.rs

@@ -330,7 +330,7 @@ impl<'a, 'tcx> Inliner<'a, 'tcx> {
             }
 
             if !is_drop {
-                for &succ in &term.successors()[..] {
+                for &succ in term.successors() {
                     work_list.push(succ);
                 }
             }

src/librustc_mir/transform/remove_noop_landing_pads.rs

@@ -78,7 +78,7 @@ impl RemoveNoopLandingPads {
             TerminatorKind::SwitchInt { .. } |
             TerminatorKind::FalseEdges { .. } |
             TerminatorKind::FalseUnwind { .. } => {
-                terminator.successors().iter().all(|succ| {
+                terminator.successors().all(|succ| {
                     nop_landing_pads.contains(succ.index())
                 })
             },

src/librustc_mir/transform/simplify.rs

@@ -91,7 +91,7 @@ impl<'a, 'tcx: 'a> CfgSimplifier<'a, 'tcx> {
 
         for (_, data) in traversal::preorder(mir) {
             if let Some(ref term) = data.terminator {
-                for &tgt in term.successors().iter() {
+                for &tgt in term.successors() {
                     pred_count[tgt] += 1;
                 }
             }
@@ -219,10 +219,10 @@ impl<'a, 'tcx: 'a> CfgSimplifier<'a, 'tcx> {
         };
 
         let first_succ = {
-            let successors = terminator.successors();
-            if let Some(&first_succ) = terminator.successors().get(0) {
-                if successors.iter().all(|s| *s == first_succ) {
-                    self.pred_count[first_succ] -= (successors.len()-1) as u32;
+            if let Some(&first_succ) = terminator.successors().nth(0) {
+                if terminator.successors().all(|s| *s == first_succ) {
+                    let count = terminator.successors().count();
+                    self.pred_count[first_succ] -= (count - 1) as u32;
                     first_succ
                 } else {
                     return false

src/librustc_mir/util/graphviz.rs

@@ -125,7 +125,7 @@ fn write_edges<W: Write>(source: BasicBlock, mir: &Mir, w: &mut W) -> io::Result
     let terminator = mir[source].terminator();
     let labels = terminator.kind.fmt_successor_labels();
 
-    for (&target, label) in terminator.successors().iter().zip(labels) {
+    for (&target, label) in terminator.successors().zip(labels) {
         writeln!(w, r#"    {} -> {} [label="{}"];"#, node(source), node(target), label)?;
     }