[MIR] Change SimplifyCfg pass to use BitVec

src/librustc_data_structures/bitvec.rs

@@ -50,6 +50,45 @@ impl BitVector {
         let extra_words = self.data.len() - num_words;
         self.data.extend((0..extra_words).map(|_| 0));
     }
+
+    /// Iterates over indexes of set bits in a sorted order
+    pub fn iter<'a>(&'a self) -> BitVectorIter<'a> {
+        BitVectorIter {
+            iter: self.data.iter(),
+            current: 0,
+            idx: 0
+        }
+    }
+}
+
+pub struct BitVectorIter<'a> {
+    iter: ::std::slice::Iter<'a, u64>,
+    current: u64,
+    idx: usize
+}
+
+impl<'a> Iterator for BitVectorIter<'a> {
+    type Item = usize;
+    fn next(&mut self) -> Option<usize> {
+        while self.current == 0 {
+            self.current = if let Some(&i) = self.iter.next() {
+                if i == 0 {
+                    self.idx += 64;
+                    continue;
+                } else {
+                    self.idx = u64s(self.idx) * 64;
+                    i
+                }
+            } else {
+                return None;
+            }
+        }
+        let offset = self.current.trailing_zeros() as usize;
+        self.current >>= offset;
+        self.current >>= 1; // shift otherwise overflows for 0b1000_0000_…_0000
+        self.idx += offset + 1;
+        return Some(self.idx - 1);
+    }
 }
 
 /// A "bit matrix" is basically a square matrix of booleans
@@ -153,6 +192,46 @@ fn word_mask(index: usize) -> (usize, u64) {
     (word, mask)
 }
 
+#[test]
+fn bitvec_iter_works() {
+    let mut bitvec = BitVector::new(100);
+    bitvec.insert(1);
+    bitvec.insert(10);
+    bitvec.insert(19);
+    bitvec.insert(62);
+    bitvec.insert(63);
+    bitvec.insert(64);
+    bitvec.insert(65);
+    bitvec.insert(66);
+    bitvec.insert(99);
+    assert_eq!(bitvec.iter().collect::<Vec<_>>(), [1, 10, 19, 62, 63, 64, 65, 66, 99]);
+}
+
+#[test]
+fn bitvec_iter_works_2() {
+    let mut bitvec = BitVector::new(300);
+    bitvec.insert(1);
+    bitvec.insert(10);
+    bitvec.insert(19);
+    bitvec.insert(62);
+    bitvec.insert(66);
+    bitvec.insert(99);
+    bitvec.insert(299);
+    assert_eq!(bitvec.iter().collect::<Vec<_>>(), [1, 10, 19, 62, 66, 99, 299]);
+
+}
+
+#[test]
+fn bitvec_iter_works_3() {
+    let mut bitvec = BitVector::new(319);
+    bitvec.insert(0);
+    bitvec.insert(127);
+    bitvec.insert(191);
+    bitvec.insert(255);
+    bitvec.insert(319);
+    assert_eq!(bitvec.iter().collect::<Vec<_>>(), [0, 127, 191, 255, 319]);
+}
+
 #[test]
 fn union_two_vecs() {
     let mut vec1 = BitVector::new(65);

src/librustc_mir/transform/mod.rs

@@ -13,4 +13,3 @@ pub mod simplify_cfg;
 pub mod erase_regions;
 pub mod no_landing_pads;
 pub mod type_check;
-mod util;

src/librustc_mir/transform/simplify_cfg.rs

@@ -8,10 +8,10 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use rustc_data_structures::bitvec::BitVector;
 use rustc::middle::const_eval::ConstVal;
 use rustc::middle::infer;
 use rustc::mir::repr::*;
-use transform::util;
 use rustc::mir::transform::MirPass;
 
 pub struct SimplifyCfg;
@@ -22,23 +22,21 @@ impl SimplifyCfg {
     }
 
     fn remove_dead_blocks(&self, mir: &mut Mir) {
-        let mut seen = vec![false; mir.basic_blocks.len()];
-
+        let mut seen = BitVector::new(mir.basic_blocks.len());
         // These blocks are always required.
-        seen[START_BLOCK.index()] = true;
-        seen[END_BLOCK.index()] = true;
+        seen.insert(START_BLOCK.index());
+        seen.insert(END_BLOCK.index());
 
-        let mut worklist = vec![START_BLOCK];
+        let mut worklist = Vec::with_capacity(4);
+        worklist.push(START_BLOCK);
         while let Some(bb) = worklist.pop() {
             for succ in mir.basic_block_data(bb).terminator().successors().iter() {
-                if !seen[succ.index()] {
-                    seen[succ.index()] = true;
+                if seen.insert(succ.index()) {
                     worklist.push(*succ);
                 }
             }
         }
-
-        util::retain_basic_blocks(mir, &seen);
+        retain_basic_blocks(mir, &seen);
     }
 
     fn remove_goto_chains(&self, mir: &mut Mir) -> bool {
@@ -90,12 +88,12 @@ impl SimplifyCfg {
         for bb in mir.all_basic_blocks() {
             let basic_block = mir.basic_block_data_mut(bb);
             let mut terminator = basic_block.terminator_mut();
-
             *terminator = match *terminator {
                 Terminator::If { ref targets, .. } if targets.0 == targets.1 => {
                     changed = true;
                     Terminator::Goto { target: targets.0 }
                 }
+
                 Terminator::If { ref targets, cond: Operand::Constant(Constant {
                     literal: Literal::Value {
                         value: ConstVal::Bool(cond)
@@ -108,6 +106,7 @@ impl SimplifyCfg {
                         Terminator::Goto { target: targets.1 }
                     }
                 }
+
                 Terminator::SwitchInt { ref targets, .. }  if targets.len() == 1 => {
                     Terminator::Goto { target: targets[0] }
                 }
@@ -131,3 +130,27 @@ impl MirPass for SimplifyCfg {
         mir.basic_blocks.shrink_to_fit();
     }
 }
+
+/// Mass removal of basic blocks to keep the ID-remapping cheap.
+fn retain_basic_blocks(mir: &mut Mir, keep: &BitVector) {
+    let num_blocks = mir.basic_blocks.len();
+
+    let mut replacements: Vec<_> = (0..num_blocks).map(BasicBlock::new).collect();
+    let mut used_blocks = 0;
+    for alive_index in keep.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.
+            mir.basic_blocks.swap(alive_index, used_blocks);
+        }
+        used_blocks += 1;
+    }
+    mir.basic_blocks.truncate(used_blocks);
+
+    for bb in mir.all_basic_blocks() {
+        for target in mir.basic_block_data_mut(bb).terminator_mut().successors_mut() {
+            *target = replacements[target.index()];
+        }
+    }
+}