Use BitMatrix for storage conflicts

3 files changed, 32 insertions, 29 deletions

diff --git a/src/librustc/mir/mod.rs b/src/librustc/mir/mod.rs
index 149d067204f..3e9f7b51cb1 100644
--- a/src/librustc/mir/mod.rs
+++ b/src/librustc/mir/mod.rs
@@ -9,7 +9,7 @@ use crate::hir::def_id::DefId;
 use crate::hir::{self, InlineAsm as HirInlineAsm};
 use crate::mir::interpret::{ConstValue, InterpError, Scalar};
 use crate::mir::visit::MirVisitable;
-use rustc_data_structures::bit_set::BitSet;
+use rustc_data_structures::bit_set::BitMatrix;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::graph::dominators::{dominators, Dominators};
 use rustc_data_structures::graph::{self, GraphPredecessors, GraphSuccessors};
@@ -3008,7 +3008,7 @@ pub struct GeneratorLayout<'tcx> {
     /// Which saved locals are storage-live at the same time. Locals that do not
     /// have conflicts with each other are allowed to overlap in the computed
     /// layout.
-    pub storage_conflicts: IndexVec<GeneratorSavedLocal, BitSet<GeneratorSavedLocal>>,
+    pub storage_conflicts: BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal>,
 
     /// Names and scopes of all the stored generator locals.
     /// NOTE(tmandry) This is *strictly* a temporary hack for codegen
@@ -3586,7 +3586,7 @@ impl<'tcx> TypeFoldable<'tcx> for GeneratorSavedLocal {
     }
 }
 
-impl<'tcx, T: Idx> TypeFoldable<'tcx> for BitSet<T> {
+impl<'tcx, R: Idx, C: Idx> TypeFoldable<'tcx> for BitMatrix<R, C> {
     fn super_fold_with<'gcx: 'tcx, F: TypeFolder<'gcx, 'tcx>>(&self, _: &mut F) -> Self {
         self.clone()
     }
diff --git a/src/librustc/ty/layout.rs b/src/librustc/ty/layout.rs
index 74f352eb3be..279784ef701 100644
--- a/src/librustc/ty/layout.rs
+++ b/src/librustc/ty/layout.rs
@@ -678,12 +678,13 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
 
                 // Next, check every pair of eligible locals to see if they
                 // conflict.
-                for (local_a, conflicts_a) in info.storage_conflicts.iter_enumerated() {
+                for local_a in info.storage_conflicts.rows() {
+                    let conflicts_a = info.storage_conflicts.count(local_a);
                     if ineligible_locals.contains(local_a) {
                         continue;
                     }
 
-                    for local_b in conflicts_a.iter() {
+                    for local_b in info.storage_conflicts.iter(local_a) {
                         // local_a and local_b are storage live at the same time, therefore they
                         // cannot overlap in the generator layout. The only way to guarantee
                         // this is if they are in the same variant, or one is ineligible
@@ -697,8 +698,8 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                         // If they conflict, we will choose one to make ineligible.
                         // This is not always optimal; it's just a greedy heuristic
                         // that seems to produce good results most of the time.
-                        let conflicts_b = &info.storage_conflicts[local_b];
-                        let (remove, other) = if conflicts_a.count() > conflicts_b.count() {
+                        let conflicts_b = info.storage_conflicts.count(local_b);
+                        let (remove, other) = if conflicts_a > conflicts_b {
                             (local_a, local_b)
                         } else {
                             (local_b, local_a)
diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index 1bcffcfb4d2..e9adb976c28 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -59,7 +59,7 @@ use rustc::ty::layout::VariantIdx;
 use rustc::ty::subst::SubstsRef;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::indexed_vec::{Idx, IndexVec};
-use rustc_data_structures::bit_set::BitSet;
+use rustc_data_structures::bit_set::{BitSet, BitMatrix};
 use std::borrow::Cow;
 use std::iter;
 use std::mem;
@@ -408,7 +408,7 @@ struct LivenessInfo {
     /// For every saved local, the set of other saved locals that are
     /// storage-live at the same time as this local. We cannot overlap locals in
     /// the layout which have conflicting storage.
-    storage_conflicts: IndexVec<GeneratorSavedLocal, BitSet<GeneratorSavedLocal>>,
+    storage_conflicts: BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal>,
 
     /// For every suspending block, the locals which are storage-live across
     /// that suspension point.
@@ -551,7 +551,9 @@ fn compute_storage_conflicts(
     ignored: &StorageIgnored,
     storage_live: DataflowResults<'tcx, MaybeStorageLive<'mir, 'tcx>>,
     storage_live_analysis: MaybeStorageLive<'mir, 'tcx>,
-) -> IndexVec<GeneratorSavedLocal, BitSet<GeneratorSavedLocal>> {
+) -> BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal> {
+    assert_eq!(body.local_decls.len(), ignored.0.domain_size());
+    assert_eq!(body.local_decls.len(), stored_locals.domain_size());
     debug!("compute_storage_conflicts({:?})", body.span);
     debug!("ignored = {:?}", ignored.0);
 
@@ -564,18 +566,15 @@ fn compute_storage_conflicts(
     // liveness. Those that do must be in the same variant to remain candidates.
     // FIXME(tmandry): Consider using sparse bitsets here once we have good
     // benchmarks for generators.
-    let mut local_conflicts: IndexVec<Local, _> =
-        // Add conflicts for every ineligible local.
-        iter::repeat(ineligible_locals.clone())
-        .take(body.local_decls.len())
-        .collect();
+    let mut local_conflicts: BitMatrix<Local, Local> =
+        BitMatrix::from_row_n(&ineligible_locals, body.local_decls.len());
 
     for_each_location(body, &storage_live_analysis, &storage_live, |state, loc| {
         let mut eligible_storage_live = state.clone().to_dense();
         eligible_storage_live.intersect(&stored_locals);
 
         for local in eligible_storage_live.iter() {
-            local_conflicts[local].union(&eligible_storage_live);
+            local_conflicts.union_row_with(&eligible_storage_live, local);
         }
 
         if eligible_storage_live.count() > 1 {
@@ -588,19 +587,22 @@ fn compute_storage_conflicts(
     // However, in practice these bitsets are not usually large. The layout code
     // also needs to keep track of how many conflicts each local has, so it's
     // simpler to keep it this way for now.
-    let storage_conflicts: IndexVec<GeneratorSavedLocal, _> = stored_locals
-        .iter()
-        .map(|local_a| {
-            if ineligible_locals.contains(local_a) {
-                // Conflicts with everything.
-                BitSet::new_filled(stored_locals.count())
-            } else {
-                // Keep overlap information only for stored locals.
-                renumber_bitset(&local_conflicts[local_a], stored_locals)
+    let mut storage_conflicts = BitMatrix::new(stored_locals.count(), stored_locals.count());
+    for (idx_a, local_a) in stored_locals.iter().enumerate() {
+        let saved_local_a = GeneratorSavedLocal::new(idx_a);
+        if ineligible_locals.contains(local_a) {
+            // Conflicts with everything.
+            storage_conflicts.insert_all_into_row(saved_local_a);
+        } else {
+            // Keep overlap information only for stored locals.
+            for (idx_b, local_b) in stored_locals.iter().enumerate() {
+                let saved_local_b = GeneratorSavedLocal::new(idx_b);
+                if local_conflicts.contains(local_a, local_b) {
+                    storage_conflicts.insert(saved_local_a, saved_local_b);
+                }
             }
-        })
-        .collect();
-
+        }
+    }
     storage_conflicts
 }
 
@@ -700,7 +702,7 @@ fn compute_layout<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
         variant_fields.push(fields);
     }
     debug!("generator variant_fields = {:?}", variant_fields);
-    debug!("generator storage_conflicts = {:?}", storage_conflicts);
+    debug!("generator storage_conflicts = {:#?}", storage_conflicts);
 
     let layout = GeneratorLayout {
         field_tys: tys,