Use GrowableBitSet in the main loop of rustc_mir_transform::sroa

compiler/rustc_index/src/bit_set.rs

@@ -114,24 +114,6 @@ impl<T: Idx> DenseBitSet<T> {
         result
     }
 
-    /// Replaces this bitset with one having the same elements, but a larger domain size.
-    #[inline]
-    pub fn enlarge(self, new_domain_size: usize) -> DenseBitSet<T> {
-        // We could also support shrinking, but it's hard to imagine a real use-case for it.
-        assert!(self.domain_size <= new_domain_size);
-        let new_num_words = num_words(new_domain_size);
-
-        let DenseBitSet { domain_size: _, mut words, marker } = self;
-
-        if new_num_words != words.len() {
-            let mut words_vec = words.into_vec();
-            words_vec.resize(new_num_words, 0);
-            words = words_vec.into_boxed_slice()
-        }
-
-        DenseBitSet { domain_size: new_domain_size, words, marker }
-    }
-
     /// Clear all elements.
     #[inline]
     pub fn clear(&mut self) {
@@ -1231,21 +1213,19 @@ impl<'a, T: Idx> Iterator for MixedBitIter<'a, T> {
 /// just be `usize`.
 #[derive(Debug, PartialEq)]
 pub struct GrowableBitSet<T: Idx> {
-    domain_size: usize,
     words: Vec<Word>,
     marker: PhantomData<T>,
 }
 
 // Manually implemented to provide `clone_from`.
 impl<T: Idx> Clone for GrowableBitSet<T> {
     fn clone(&self) -> Self {
-        let &GrowableBitSet { domain_size, ref words, marker } = self;
-        GrowableBitSet { domain_size, words: words.clone(), marker }
+        let &GrowableBitSet { ref words, marker } = self;
+        GrowableBitSet { words: words.clone(), marker }
     }
 
     fn clone_from(&mut self, source: &Self) {
-        let GrowableBitSet { domain_size, words, marker } = source;
-        self.domain_size.clone_from(domain_size);
+        let GrowableBitSet { words, marker } = source;
         self.words.clone_from(words);
         self.marker.clone_from(marker);
     }
@@ -1258,28 +1238,25 @@ impl<T: Idx> Default for GrowableBitSet<T> {
 }
 
 impl<T: Idx> GrowableBitSet<T> {
-    /// Ensure that the set can hold at least `min_domain_size` elements.
-    pub fn ensure(&mut self, min_domain_size: usize) {
-        if self.domain_size < min_domain_size {
-            self.domain_size = min_domain_size;
-        }
+    /// Ensure that the set has allocated and initialized at least `min_num_bits` bits.
+    fn ensure(&mut self, min_num_bits: usize) {
+        let min_num_words = num_words(min_num_bits);
+        self.ensure_words(min_num_words);
+    }
 
-        let min_num_words = num_words(min_domain_size);
+    /// Ensures that the set has allocated and initialized at least `min_num_words` words.
+    fn ensure_words(&mut self, min_num_words: usize) {
         if self.words.len() < min_num_words {
             self.words.resize(min_num_words, 0)
         }
     }
 
     pub fn new_empty() -> GrowableBitSet<T> {
-        GrowableBitSet { domain_size: 0, words: vec![], marker: PhantomData }
+        GrowableBitSet { words: vec![], marker: PhantomData }
     }
 
     pub fn with_capacity(capacity: usize) -> GrowableBitSet<T> {
-        GrowableBitSet {
-            domain_size: capacity,
-            words: vec![0; num_words(capacity)],
-            marker: PhantomData,
-        }
+        GrowableBitSet { words: Vec::with_capacity(num_words(capacity)), marker: PhantomData }
     }
 
     /// Returns `true` if the set has changed.
@@ -1309,6 +1286,16 @@ impl<T: Idx> GrowableBitSet<T> {
     pub fn iter(&self) -> BitIter<'_, T> {
         BitIter::new(&self.words)
     }
+
+    /// Mutates `self = self | other`.
+    #[inline]
+    pub fn union(&mut self, other: &GrowableBitSet<T>) {
+        // Eagerly grow `self` to be at least as large as `other`.
+        // This is simpler than trying to check whether `other` has any nonzero
+        // bits beyond our current size.
+        self.ensure_words(other.words.len());
+        update_words(&mut self.words[..other.words.len()], &other.words, |a, b| a | b);
+    }
 }
 
 /// A fixed-size 2D bit matrix type with a dense representation.

compiler/rustc_index/src/bit_set/tests.rs

@@ -554,6 +554,35 @@ fn grow() {
     }
 }
 
+#[test]
+fn growable_union() {
+    // Create two input sets with partly-overlapping values, and different sizes.
+    let mut twos = GrowableBitSet::<usize>::new_empty();
+    for i in (0usize..100).map(|x| x * 2) {
+        twos.insert(i);
+    }
+
+    let mut threes = GrowableBitSet::<usize>::new_empty();
+    for i in (0usize..100).map(|x| x * 3) {
+        threes.insert(i);
+    }
+
+    // Double-check that we did end up with input sets of different sizes.
+    assert_ne!(twos.words.len(), threes.words.len());
+
+    // Perform a union in both directions, and check that the resulting contents are correct.
+    for (mut lhs, rhs) in [(twos.clone(), threes.clone()), (threes.clone(), twos.clone())] {
+        lhs.union(&rhs);
+
+        for i in 0..400 {
+            assert_eq!(
+                lhs.contains(i),
+                (i.is_multiple_of(2) && i < 200) || (i.is_multiple_of(3) && i < 300)
+            );
+        }
+    }
+}
+
 #[test]
 fn matrix_intersection() {
     let mut matrix: BitMatrix<usize, usize> = BitMatrix::new(200, 200);

compiler/rustc_mir_dataflow/src/value_analysis.rs

@@ -5,7 +5,7 @@ use std::ops::Range;
 use rustc_abi::{FieldIdx, VariantIdx};
 use rustc_data_structures::fx::{FxHashMap, FxIndexSet, StdEntry};
 use rustc_index::IndexVec;
-use rustc_index::bit_set::DenseBitSet;
+use rustc_index::bit_set::GrowableBitSet;
 use rustc_middle::mir::visit::{PlaceContext, Visitor};
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, Ty, TyCtxt, Unnormalized};
@@ -1039,9 +1039,9 @@ pub fn iter_fields<'tcx>(
 }
 
 /// Returns all locals with projections that have their reference or address taken.
-pub fn excluded_locals(body: &Body<'_>) -> DenseBitSet<Local> {
+pub fn excluded_locals(body: &Body<'_>) -> GrowableBitSet<Local> {
     struct Collector {
-        result: DenseBitSet<Local>,
+        result: GrowableBitSet<Local>,
     }
 
     impl<'tcx> Visitor<'tcx> for Collector {
@@ -1054,7 +1054,7 @@ pub fn excluded_locals(body: &Body<'_>) -> DenseBitSet<Local> {
         }
     }
 
-    let mut collector = Collector { result: DenseBitSet::new_empty(body.local_decls.len()) };
+    let mut collector = Collector { result: GrowableBitSet::new_empty() };
     collector.visit_body(body);
     collector.result
 }

compiler/rustc_mir_transform/src/sroa.rs

@@ -2,7 +2,7 @@ use rustc_abi::FieldIdx;
 use rustc_data_structures::flat_map_in_place::FlatMapInPlace;
 use rustc_hir::attrs::lang_items::LangItem;
 use rustc_index::IndexVec;
-use rustc_index::bit_set::DenseBitSet;
+use rustc_index::bit_set::{DenseBitSet, GrowableBitSet};
 use rustc_middle::bug;
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
@@ -40,7 +40,6 @@ impl<'tcx> crate::MirPass<'tcx> for ScalarReplacementOfAggregates {
             let all_dead_locals = replace_flattened_locals(tcx, body, replacements);
             if !all_dead_locals.is_empty() {
                 excluded.union(&all_dead_locals);
-                excluded = excluded.enlarge(body.local_decls.len());
             } else {
                 break;
             }
@@ -57,7 +56,7 @@ impl<'tcx> crate::MirPass<'tcx> for ScalarReplacementOfAggregates {
 ///   client code.
 fn escaping_locals<'tcx>(
     tcx: TyCtxt<'tcx>,
-    excluded: &DenseBitSet<Local>,
+    excluded: &GrowableBitSet<Local>,
     body: &Body<'tcx>,
 ) -> DenseBitSet<Local> {
     let is_excluded_ty = |ty: Ty<'tcx>| {
@@ -208,9 +207,11 @@ fn replace_flattened_locals<'tcx>(
     tcx: TyCtxt<'tcx>,
     body: &mut Body<'tcx>,
     replacements: ReplacementMap<'tcx>,
-) -> DenseBitSet<Local> {
-    let mut all_dead_locals = DenseBitSet::new_empty(replacements.fragments.len());
-    for (local, replacements) in replacements.fragments.iter_enumerated() {
+) -> GrowableBitSet<Local> {
+    // Start with an empty GrowableBitSet, to avoid allocation if nothing is dead.
+    // Then fill the set in descending order so that it allocates at most once.
+    let mut all_dead_locals = GrowableBitSet::new_empty();
+    for (local, replacements) in replacements.fragments.iter_enumerated().rev() {
         if replacements.is_some() {
             all_dead_locals.insert(local);
         }
@@ -249,7 +250,7 @@ struct ReplacementVisitor<'tcx, 'll> {
     /// Work to do.
     replacements: &'ll ReplacementMap<'tcx>,
     /// This is used to check that we are not leaving references to replaced locals behind.
-    all_dead_locals: DenseBitSet<Local>,
+    all_dead_locals: GrowableBitSet<Local>,
     patch: MirPatch<'tcx>,
 }