Use a BitVector instead of Vec<bool> for recording cleanup blocks

Also adds a FromIterator impl for BitVector to allow construction of a
BitVector from an iterator yeilding bools.

2 files changed, 34 insertions, 2 deletions

diff --git a/src/librustc_data_structures/bitvec.rs b/src/librustc_data_structures/bitvec.rs
index e45c6cfc6dc..092b406ae9e 100644
--- a/src/librustc_data_structures/bitvec.rs
+++ b/src/librustc_data_structures/bitvec.rs
@@ -8,6 +8,8 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use std::iter::FromIterator;
+
 /// A very simple BitVector type.
 #[derive(Clone)]
 pub struct BitVector {
@@ -51,7 +53,9 @@ impl BitVector {
     pub fn grow(&mut self, num_bits: usize) {
         let num_words = u64s(num_bits);
         let extra_words = self.data.len() - num_words;
-        self.data.extend((0..extra_words).map(|_| 0));
+        if extra_words > 0 {
+            self.data.extend((0..extra_words).map(|_| 0));
+        }
     }
 
     /// Iterates over indexes of set bits in a sorted order
@@ -94,6 +98,27 @@ impl<'a> Iterator for BitVectorIter<'a> {
     }
 }
 
+impl FromIterator<bool> for BitVector {
+    fn from_iter<I>(iter: I) -> BitVector where I: IntoIterator<Item=bool> {
+        let iter = iter.into_iter();
+        let (len, _) = iter.size_hint();
+        // Make the minimum length for the bitvector 64 bits since that's
+        // the smallest non-zero size anyway.
+        let len = if len < 64 { 64 } else { len };
+        let mut bv = BitVector::new(len);
+        for (idx, val) in iter.enumerate() {
+            if idx > len {
+                bv.grow(idx);
+            }
+            if val {
+                bv.insert(idx);
+            }
+        }
+
+        bv
+    }
+}
+
 /// A "bit matrix" is basically a square matrix of booleans
 /// represented as one gigantic bitvector. In other words, it is as if
 /// you have N bitvectors, each of length N. Note that `elements` here is `N`/
diff --git a/src/librustc_mir/transform/break_critical_edges.rs b/src/librustc_mir/transform/break_critical_edges.rs
index e78c8e8fd73..e1fb5dfd437 100644
--- a/src/librustc_mir/transform/break_critical_edges.rs
+++ b/src/librustc_mir/transform/break_critical_edges.rs
@@ -13,6 +13,8 @@ use rustc::mir::repr::*;
 use rustc::mir::transform::{MirPass, Pass};
 use syntax::ast::NodeId;
 
+use rustc_data_structures::bitvec::BitVector;
+
 use traversal;
 
 pub struct BreakCriticalEdges;
@@ -60,6 +62,9 @@ fn break_critical_edges(mir: &mut Mir) {
         }
     }
 
+    let cleanup_map : BitVector = mir.basic_blocks
+        .iter().map(|bb| bb.is_cleanup).collect();
+
     // We need a place to store the new blocks generated
     let mut new_blocks = Vec::new();
 
@@ -84,7 +89,9 @@ fn break_critical_edges(mir: &mut Mir) {
                             scope: term_scope,
                             kind: TerminatorKind::Goto { target: *tgt }
                         };
-                        let data = BasicBlockData::new(Some(goto));
+                        let mut data = BasicBlockData::new(Some(goto));
+                        data.is_cleanup = cleanup_map.contains(tgt.index());
+
                         // Get the index it will be when inserted into the MIR
                         let idx = cur_len + new_blocks.len();
                         new_blocks.push(data);