From 79e0a0faf93922f77f7658d2b1b95c79c9bd9608 Mon Sep 17 00:00:00 2001
From: Will Crichton <wcrichto@cs.stanford.edu>
Date: Tue, 24 Aug 2021 17:50:08 -0700
Subject: [PATCH] Refactor BitSet relational methods into trait with
 parameterized right-hand side

---
 compiler/rustc_index/src/bit_set.rs           | 359 ++++++++++++------
 .../src/borrow_check/region_infer/values.rs   |   4 +-
 compiler/rustc_mir/src/transform/generator.rs |   4 +-
 3 files changed, 238 insertions(+), 129 deletions(-)

diff --git a/compiler/rustc_index/src/bit_set.rs b/compiler/rustc_index/src/bit_set.rs
index 282bed274e6..0b5745072ad 100644
--- a/compiler/rustc_index/src/bit_set.rs
+++ b/compiler/rustc_index/src/bit_set.rs
@@ -16,6 +16,43 @@ pub type Word = u64;
 pub const WORD_BYTES: usize = mem::size_of::<Word>();
 pub const WORD_BITS: usize = WORD_BYTES * 8;
 
+pub trait BitRelations<Rhs> {
+    fn union(&mut self, other: &Rhs) -> bool;
+    fn subtract(&mut self, other: &Rhs) -> bool;
+    fn intersect(&mut self, other: &Rhs) -> bool;
+}
+
+macro_rules! bit_relations_inherent_impls {
+    () => {
+        /// Sets `self = self | other` and returns `true` if `self` changed
+        /// (i.e., if new bits were added).
+        pub fn union<Rhs>(&mut self, other: &Rhs) -> bool
+        where
+            Self: BitRelations<Rhs>,
+        {
+            <Self as BitRelations<Rhs>>::union(self, other)
+        }
+
+        /// Sets `self = self - other` and returns `true` if `self` changed.
+        /// (i.e., if any bits were removed).
+        pub fn subtract<Rhs>(&mut self, other: &Rhs) -> bool
+        where
+            Self: BitRelations<Rhs>,
+        {
+            <Self as BitRelations<Rhs>>::subtract(self, other)
+        }
+
+        /// Sets `self = self & other` and return `true` if `self` changed.
+        /// (i.e., if any bits were removed).
+        pub fn intersect<Rhs>(&mut self, other: &Rhs) -> bool
+        where
+            Self: BitRelations<Rhs>,
+        {
+            <Self as BitRelations<Rhs>>::intersect(self, other)
+        }
+    };
+}
+
 /// A fixed-size bitset type with a dense representation.
 ///
 /// NOTE: Use [`GrowableBitSet`] if you need support for resizing after creation.
@@ -134,25 +171,6 @@ impl<T: Idx> BitSet<T> {
         new_word != word
     }
 
-    /// Sets `self = self | other` and returns `true` if `self` changed
-    /// (i.e., if new bits were added).
-    pub fn union(&mut self, other: &impl UnionIntoBitSet<T>) -> bool {
-        other.union_into(self)
-    }
-
-    /// Sets `self = self - other` and returns `true` if `self` changed.
-    /// (i.e., if any bits were removed).
-    pub fn subtract(&mut self, other: &impl SubtractFromBitSet<T>) -> bool {
-        other.subtract_from(self)
-    }
-
-    /// Sets `self = self & other` and return `true` if `self` changed.
-    /// (i.e., if any bits were removed).
-    pub fn intersect(&mut self, other: &BitSet<T>) -> bool {
-        assert_eq!(self.domain_size, other.domain_size);
-        bitwise(&mut self.words, &other.words, |a, b| a & b)
-    }
-
     /// Gets a slice of the underlying words.
     pub fn words(&self) -> &[Word] {
         &self.words
@@ -208,33 +226,167 @@ impl<T: Idx> BitSet<T> {
 
         not_already
     }
+
+    bit_relations_inherent_impls! {}
 }
 
-/// This is implemented by all the bitsets so that BitSet::union() can be
-/// passed any type of bitset.
-pub trait UnionIntoBitSet<T: Idx> {
-    // Performs `other = other | self`.
-    fn union_into(&self, other: &mut BitSet<T>) -> bool;
-}
-
-/// This is implemented by all the bitsets so that BitSet::subtract() can be
-/// passed any type of bitset.
-pub trait SubtractFromBitSet<T: Idx> {
-    // Performs `other = other - self`.
-    fn subtract_from(&self, other: &mut BitSet<T>) -> bool;
-}
-
-impl<T: Idx> UnionIntoBitSet<T> for BitSet<T> {
-    fn union_into(&self, other: &mut BitSet<T>) -> bool {
+impl<T: Idx> BitRelations<BitSet<T>> for BitSet<T> {
+    fn union(&mut self, other: &BitSet<T>) -> bool {
         assert_eq!(self.domain_size, other.domain_size);
-        bitwise(&mut other.words, &self.words, |a, b| a | b)
+        bitwise(&mut self.words, &other.words, |a, b| a | b)
+    }
+
+    fn subtract(&mut self, other: &BitSet<T>) -> bool {
+        assert_eq!(self.domain_size, other.domain_size);
+        bitwise(&mut self.words, &other.words, |a, b| a & !b)
+    }
+
+    fn intersect(&mut self, other: &BitSet<T>) -> bool {
+        assert_eq!(self.domain_size, other.domain_size);
+        bitwise(&mut self.words, &other.words, |a, b| a & b)
     }
 }
 
-impl<T: Idx> SubtractFromBitSet<T> for BitSet<T> {
-    fn subtract_from(&self, other: &mut BitSet<T>) -> bool {
-        assert_eq!(self.domain_size, other.domain_size);
-        bitwise(&mut other.words, &self.words, |a, b| a & !b)
+fn sequential_update<T: Idx>(mut f: impl FnMut(T) -> bool, it: impl Iterator<Item = T>) -> bool {
+    let mut changed = false;
+    for elem in it {
+        changed |= f(elem);
+    }
+    changed
+}
+
+fn sparse_intersect<T: Idx>(
+    set: &mut SparseBitSet<T>,
+    other_contains: impl Fn(&T) -> bool,
+) -> bool {
+    let mut changed = false;
+    for i in (0..set.len()).rev() {
+        if !other_contains(&set.elems[i]) {
+            set.elems.remove(i);
+            changed = true;
+        }
+    }
+    changed
+}
+
+impl<T: Idx> BitRelations<SparseBitSet<T>> for BitSet<T> {
+    fn union(&mut self, other: &SparseBitSet<T>) -> bool {
+        sequential_update(|elem| self.insert(elem), other.iter().cloned())
+    }
+
+    fn subtract(&mut self, other: &SparseBitSet<T>) -> bool {
+        sequential_update(|elem| self.remove(elem), other.iter().cloned())
+    }
+
+    fn intersect(&mut self, other: &SparseBitSet<T>) -> bool {
+        self.intersect(&other.to_dense())
+    }
+}
+
+impl<T: Idx> BitRelations<BitSet<T>> for SparseBitSet<T> {
+    fn union(&mut self, other: &BitSet<T>) -> bool {
+        sequential_update(|elem| self.insert(elem), other.iter())
+    }
+
+    fn subtract(&mut self, other: &BitSet<T>) -> bool {
+        sequential_update(|elem| self.remove(elem), other.iter())
+    }
+
+    fn intersect(&mut self, other: &BitSet<T>) -> bool {
+        sparse_intersect(self, |el| other.contains(*el))
+    }
+}
+
+impl<T: Idx> BitRelations<SparseBitSet<T>> for SparseBitSet<T> {
+    fn union(&mut self, other: &SparseBitSet<T>) -> bool {
+        sequential_update(|elem| self.insert(elem), other.iter().cloned())
+    }
+
+    fn subtract(&mut self, other: &SparseBitSet<T>) -> bool {
+        sequential_update(|elem| self.insert(elem), other.iter().cloned())
+    }
+
+    fn intersect(&mut self, other: &SparseBitSet<T>) -> bool {
+        sparse_intersect(self, |el| other.contains(*el))
+    }
+}
+
+impl<T: Idx, S> BitRelations<HybridBitSet<T>> for S
+where
+    S: BitRelations<BitSet<T>> + BitRelations<SparseBitSet<T>>,
+{
+    fn union(&mut self, other: &HybridBitSet<T>) -> bool {
+        match other {
+            HybridBitSet::Sparse(sparse) => self.union(sparse),
+            HybridBitSet::Dense(dense) => self.union(dense),
+        }
+    }
+
+    fn subtract(&mut self, other: &HybridBitSet<T>) -> bool {
+        match other {
+            HybridBitSet::Sparse(sparse) => self.subtract(sparse),
+            HybridBitSet::Dense(dense) => self.subtract(dense),
+        }
+    }
+
+    fn intersect(&mut self, other: &HybridBitSet<T>) -> bool {
+        match other {
+            HybridBitSet::Sparse(sparse) => self.intersect(sparse),
+            HybridBitSet::Dense(dense) => self.intersect(dense),
+        }
+    }
+}
+
+impl<T: Idx> BitRelations<HybridBitSet<T>> for HybridBitSet<T> {
+    fn union(&mut self, other: &HybridBitSet<T>) -> bool {
+        match self {
+            HybridBitSet::Sparse(self_sparse) => {
+                match other {
+                    HybridBitSet::Sparse(other_sparse) => self_sparse.union(other_sparse),
+
+                    HybridBitSet::Dense(other_dense) => {
+                        // `self` is sparse and `other` is dense. To
+                        // merge them, we have two available strategies:
+                        // * Densify `self` then merge other
+                        // * Clone other then integrate bits from `self`
+                        // The second strategy requires dedicated method
+                        // since the usual `union` returns the wrong
+                        // result. In the dedicated case the computation
+                        // is slightly faster if the bits of the sparse
+                        // bitset map to only few words of the dense
+                        // representation, i.e. indices are near each
+                        // other.
+                        //
+                        // Benchmarking seems to suggest that the second
+                        // option is worth it.
+                        let mut new_dense = other_dense.clone();
+                        let changed = new_dense.reverse_union_sparse(self_sparse);
+                        *self = HybridBitSet::Dense(new_dense);
+                        changed
+                    }
+                }
+            }
+
+            HybridBitSet::Dense(self_dense) => self_dense.union(other),
+        }
+    }
+
+    fn subtract(&mut self, other: &HybridBitSet<T>) -> bool {
+        // FIXME(willcrichton): should there be an optimized sparse / dense version?
+        match self {
+            HybridBitSet::Sparse(self_sparse) => self_sparse.subtract(other),
+            HybridBitSet::Dense(self_dense) => self_dense.subtract(other),
+        }
+    }
+
+    fn intersect(&mut self, other: &HybridBitSet<T>) -> bool {
+        // FIXME(willcrichton): should there be an optimized sparse / dense version?
+        match self {
+            HybridBitSet::Sparse(self_sparse) => self_sparse.intersect(other),
+            HybridBitSet::Dense(self_dense) => {
+                <BitSet<T> as BitRelations<HybridBitSet<T>>>::intersect(self_dense, other)
+            }
+        }
     }
 }
 
@@ -441,28 +593,8 @@ impl<T: Idx> SparseBitSet<T> {
     fn iter(&self) -> slice::Iter<'_, T> {
         self.elems.iter()
     }
-}
 
-impl<T: Idx> UnionIntoBitSet<T> for SparseBitSet<T> {
-    fn union_into(&self, other: &mut BitSet<T>) -> bool {
-        assert_eq!(self.domain_size, other.domain_size);
-        let mut changed = false;
-        for elem in self.iter() {
-            changed |= other.insert(*elem);
-        }
-        changed
-    }
-}
-
-impl<T: Idx> SubtractFromBitSet<T> for SparseBitSet<T> {
-    fn subtract_from(&self, other: &mut BitSet<T>) -> bool {
-        assert_eq!(self.domain_size, other.domain_size);
-        let mut changed = false;
-        for elem in self.iter() {
-            changed |= other.remove(*elem);
-        }
-        changed
-    }
+    bit_relations_inherent_impls! {}
 }
 
 /// A fixed-size bitset type with a hybrid representation: sparse when there
@@ -579,48 +711,6 @@ impl<T: Idx> HybridBitSet<T> {
         }
     }
 
-    pub fn union(&mut self, other: &HybridBitSet<T>) -> bool {
-        match self {
-            HybridBitSet::Sparse(self_sparse) => {
-                match other {
-                    HybridBitSet::Sparse(other_sparse) => {
-                        // Both sets are sparse. Add the elements in
-                        // `other_sparse` to `self` one at a time. This
-                        // may or may not cause `self` to be densified.
-                        assert_eq!(self.domain_size(), other.domain_size());
-                        let mut changed = false;
-                        for elem in other_sparse.iter() {
-                            changed |= self.insert(*elem);
-                        }
-                        changed
-                    }
-                    HybridBitSet::Dense(other_dense) => {
-                        // `self` is sparse and `other` is dense. To
-                        // merge them, we have two available strategies:
-                        // * Densify `self` then merge other
-                        // * Clone other then integrate bits from `self`
-                        // The second strategy requires dedicated method
-                        // since the usual `union` returns the wrong
-                        // result. In the dedicated case the computation
-                        // is slightly faster if the bits of the sparse
-                        // bitset map to only few words of the dense
-                        // representation, i.e. indices are near each
-                        // other.
-                        //
-                        // Benchmarking seems to suggest that the second
-                        // option is worth it.
-                        let mut new_dense = other_dense.clone();
-                        let changed = new_dense.reverse_union_sparse(self_sparse);
-                        *self = HybridBitSet::Dense(new_dense);
-                        changed
-                    }
-                }
-            }
-
-            HybridBitSet::Dense(self_dense) => self_dense.union(other),
-        }
-    }
-
     /// Converts to a dense set, consuming itself in the process.
     pub fn to_dense(self) -> BitSet<T> {
         match self {
@@ -635,24 +725,8 @@ impl<T: Idx> HybridBitSet<T> {
             HybridBitSet::Dense(dense) => HybridIter::Dense(dense.iter()),
         }
     }
-}
 
-impl<T: Idx> UnionIntoBitSet<T> for HybridBitSet<T> {
-    fn union_into(&self, other: &mut BitSet<T>) -> bool {
-        match self {
-            HybridBitSet::Sparse(sparse) => sparse.union_into(other),
-            HybridBitSet::Dense(dense) => dense.union_into(other),
-        }
-    }
-}
-
-impl<T: Idx> SubtractFromBitSet<T> for HybridBitSet<T> {
-    fn subtract_from(&self, other: &mut BitSet<T>) -> bool {
-        match self {
-            HybridBitSet::Sparse(sparse) => sparse.subtract_from(other),
-            HybridBitSet::Dense(dense) => dense.subtract_from(other),
-        }
-    }
+    bit_relations_inherent_impls! {}
 }
 
 pub enum HybridIter<'a, T: Idx> {
@@ -974,6 +1048,19 @@ impl<R: Idx, C: Idx> SparseBitMatrix<R, C> {
         self.ensure_row(row).insert(column)
     }
 
+    pub fn remove(&mut self, row: R, column: C) -> bool {
+        match self.rows.get_mut(row) {
+            Some(Some(row)) => row.remove(column),
+            _ => false,
+        }
+    }
+
+    pub fn clear(&mut self, row: R) {
+        if let Some(Some(row)) = self.rows.get_mut(row) {
+            row.clear();
+        }
+    }
+
     /// Do the bits from `row` contain `column`? Put another way, is
     /// the matrix cell at `(row, column)` true?  Put yet another way,
     /// if the matrix represents (transitive) reachability, can
@@ -1002,11 +1089,6 @@ impl<R: Idx, C: Idx> SparseBitMatrix<R, C> {
         }
     }
 
-    /// Union a row, `from`, into the `into` row.
-    pub fn union_into_row(&mut self, into: R, from: &HybridBitSet<C>) -> bool {
-        self.ensure_row(into).union(from)
-    }
-
     /// Insert all bits in the given row.
     pub fn insert_all_into_row(&mut self, row: R) {
         self.ensure_row(row).insert_all();
@@ -1025,6 +1107,33 @@ impl<R: Idx, C: Idx> SparseBitMatrix<R, C> {
     pub fn row(&self, row: R) -> Option<&HybridBitSet<C>> {
         if let Some(Some(row)) = self.rows.get(row) { Some(row) } else { None }
     }
+
+    pub fn intersect_row<Set>(&mut self, row: R, set: &Set) -> bool
+    where
+        HybridBitSet<C>: BitRelations<Set>,
+    {
+        match self.rows.get_mut(row) {
+            Some(Some(row)) => row.intersect(set),
+            _ => false,
+        }
+    }
+
+    pub fn subtract_row<Set>(&mut self, row: R, set: &Set) -> bool
+    where
+        HybridBitSet<C>: BitRelations<Set>,
+    {
+        match self.rows.get_mut(row) {
+            Some(Some(row)) => row.subtract(set),
+            _ => false,
+        }
+    }
+
+    pub fn union_row<Set>(&mut self, row: R, set: &Set) -> bool
+    where
+        HybridBitSet<C>: BitRelations<Set>,
+    {
+        self.ensure_row(row).union(set)
+    }
 }
 
 #[inline]
diff --git a/compiler/rustc_mir/src/borrow_check/region_infer/values.rs b/compiler/rustc_mir/src/borrow_check/region_infer/values.rs
index f247d07e1f0..2864abde002 100644
--- a/compiler/rustc_mir/src/borrow_check/region_infer/values.rs
+++ b/compiler/rustc_mir/src/borrow_check/region_infer/values.rs
@@ -160,7 +160,7 @@ impl<N: Idx> LivenessValues<N> {
     /// region. Returns whether any of them are newly added.
     crate fn add_elements(&mut self, row: N, locations: &HybridBitSet<PointIndex>) -> bool {
         debug!("LivenessValues::add_elements(row={:?}, locations={:?})", row, locations);
-        self.points.union_into_row(row, locations)
+        self.points.union_row(row, locations)
     }
 
     /// Adds all the control-flow points to the values for `r`.
@@ -294,7 +294,7 @@ impl<N: Idx> RegionValues<N> {
     /// the region `to` in `self`.
     crate fn merge_liveness<M: Idx>(&mut self, to: N, from: M, values: &LivenessValues<M>) {
         if let Some(set) = values.points.row(from) {
-            self.points.union_into_row(to, set);
+            self.points.union_row(to, set);
         }
     }
 
diff --git a/compiler/rustc_mir/src/transform/generator.rs b/compiler/rustc_mir/src/transform/generator.rs
index 963f93a1ace..acdaa5b4568 100644
--- a/compiler/rustc_mir/src/transform/generator.rs
+++ b/compiler/rustc_mir/src/transform/generator.rs
@@ -626,7 +626,7 @@ fn compute_storage_conflicts(
     // Locals that are always live or ones that need to be stored across
     // suspension points are not eligible for overlap.
     let mut ineligible_locals = always_live_locals.into_inner();
-    ineligible_locals.intersect(saved_locals);
+    ineligible_locals.intersect(&**saved_locals);
 
     // Compute the storage conflicts for all eligible locals.
     let mut visitor = StorageConflictVisitor {
@@ -701,7 +701,7 @@ impl<'body, 'tcx, 's> StorageConflictVisitor<'body, 'tcx, 's> {
         }
 
         let mut eligible_storage_live = flow_state.clone();
-        eligible_storage_live.intersect(&self.saved_locals);
+        eligible_storage_live.intersect(&**self.saved_locals);
 
         for local in eligible_storage_live.iter() {
             self.local_conflicts.union_row_with(&eligible_storage_live, local);