Use Sparse bitsets instead of dense ones for NLL results

Fixes #48170

Author: spastorino

src/librustc_data_structures/bitvec.rs

@@ -8,7 +8,11 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use std::collections::BTreeMap;
+use std::collections::btree_map::Entry;
+use std::marker::PhantomData;
 use std::iter::FromIterator;
+use indexed_vec::{Idx, IndexVec};
 
 type Word = u128;
 const WORD_BITS: usize = 128;
@@ -257,6 +261,199 @@ impl BitMatrix {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct SparseBitMatrix<R, C> where R: Idx, C: Idx {
+    vector: IndexVec<R, SparseBitSet<C>>,
+}
+
+impl<R: Idx, C: Idx> SparseBitMatrix<R, C> {
+    /// Create a new `rows x columns` matrix, initially empty.
+    pub fn new(rows: R, _columns: C) -> SparseBitMatrix<R, C> {
+        SparseBitMatrix {
+            vector: IndexVec::from_elem_n(SparseBitSet::new(), rows.index()),
+        }
+    }
+
+    /// Sets the cell at `(row, column)` to true. Put another way, insert
+    /// `column` to the bitset for `row`.
+    ///
+    /// Returns true if this changed the matrix, and false otherwise.
+    pub fn add(&mut self, row: R, column: C) -> bool {
+        self.vector[row].insert(column)
+    }
+
+    /// 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
+    /// `row` reach `column`?
+    pub fn contains(&self, row: R, column: C) -> bool {
+        self.vector[row].contains(column)
+    }
+
+    /// Add the bits from row `read` to the bits from row `write`,
+    /// return true if anything changed.
+    ///
+    /// This is used when computing transitive reachability because if
+    /// you have an edge `write -> read`, because in that case
+    /// `write` can reach everything that `read` can (and
+    /// potentially more).
+    pub fn merge(&mut self, read: R, write: R) -> bool {
+        let mut changed = false;
+
+        if read != write {
+            let (bit_set_read, bit_set_write) = self.vector.pick2_mut(read, write);
+
+            for read_val in bit_set_read.iter() {
+                changed = changed | bit_set_write.insert(read_val);
+            }
+        }
+
+        changed
+    }
+
+    /// Iterates through all the columns set to true in a given row of
+    /// the matrix.
+    pub fn iter<'a>(&'a self, row: R) -> impl Iterator<Item = C> + 'a {
+        self.vector[row].iter()
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct SparseBitSet<I: Idx> {
+    chunk_bits: BTreeMap<u32, Word>,
+    _marker: PhantomData<I>,
+}
+
+#[derive(Copy, Clone)]
+pub struct SparseChunk<I> {
+    key: u32,
+    bits: Word,
+    _marker: PhantomData<I>,
+}
+
+impl<I: Idx> SparseChunk<I> {
+    pub fn one(index: I) -> Self {
+        let index = index.index();
+        let key_usize = index / 128;
+        let key = key_usize as u32;
+        assert_eq!(key as usize, key_usize);
+        SparseChunk {
+            key,
+            bits: 1 << (index % 128),
+            _marker: PhantomData
+        }
+    }
+
+    pub fn any(&self) -> bool {
+        self.bits != 0
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = I> {
+        let base = self.key as usize * 128;
+        let mut bits = self.bits;
+        (0..128).map(move |i| {
+            let current_bits = bits;
+            bits >>= 1;
+            (i, current_bits)
+        }).take_while(|&(_, bits)| bits != 0)
+          .filter_map(move |(i, bits)| {
+            if (bits & 1) != 0 {
+                Some(I::new(base + i))
+            } else {
+                None
+            }
+        })
+    }
+}
+
+impl<I: Idx> SparseBitSet<I> {
+    pub fn new() -> Self {
+        SparseBitSet {
+            chunk_bits: BTreeMap::new(),
+            _marker: PhantomData
+        }
+    }
+
+    pub fn capacity(&self) -> usize {
+        self.chunk_bits.len() * 128
+    }
+
+    pub fn contains_chunk(&self, chunk: SparseChunk<I>) -> SparseChunk<I> {
+        SparseChunk {
+            bits: self.chunk_bits.get(&chunk.key).map_or(0, |bits| bits & chunk.bits),
+            ..chunk
+        }
+    }
+
+    pub fn insert_chunk(&mut self, chunk: SparseChunk<I>) -> SparseChunk<I> {
+        if chunk.bits == 0 {
+            return chunk;
+        }
+        let bits = self.chunk_bits.entry(chunk.key).or_insert(0);
+        let old_bits = *bits;
+        let new_bits = old_bits | chunk.bits;
+        *bits = new_bits;
+        let changed = new_bits ^ old_bits;
+        SparseChunk {
+            bits: changed,
+            ..chunk
+        }
+    }
+
+    pub fn remove_chunk(&mut self, chunk: SparseChunk<I>) -> SparseChunk<I> {
+        if chunk.bits == 0 {
+            return chunk;
+        }
+        let changed = match self.chunk_bits.entry(chunk.key) {
+            Entry::Occupied(mut bits) => {
+                let old_bits = *bits.get();
+                let new_bits = old_bits & !chunk.bits;
+                if new_bits == 0 {
+                    bits.remove();
+                } else {
+                    bits.insert(new_bits);
+                }
+                new_bits ^ old_bits
+            }
+            Entry::Vacant(_) => 0
+        };
+        SparseChunk {
+            bits: changed,
+            ..chunk
+        }
+    }
+
+    pub fn clear(&mut self) {
+        self.chunk_bits.clear();
+    }
+
+    pub fn chunks<'a>(&'a self) -> impl Iterator<Item = SparseChunk<I>> + 'a {
+        self.chunk_bits.iter().map(|(&key, &bits)| {
+            SparseChunk {
+                key,
+                bits,
+                _marker: PhantomData
+            }
+        })
+    }
+
+    pub fn contains(&self, index: I) -> bool {
+        self.contains_chunk(SparseChunk::one(index)).any()
+    }
+
+    pub fn insert(&mut self, index: I) -> bool {
+        self.insert_chunk(SparseChunk::one(index)).any()
+    }
+
+    pub fn remove(&mut self, index: I) -> bool {
+        self.remove_chunk(SparseChunk::one(index)).any()
+    }
+
+    pub fn iter<'a>(&'a self) -> impl Iterator<Item = I> + 'a {
+        self.chunks().flat_map(|chunk| chunk.iter())
+    }
+}
+
 #[inline]
 fn words(elements: usize) -> usize {
     (elements + WORD_BITS - 1) / WORD_BITS

src/librustc_data_structures/indexed_vec.rs

@@ -482,6 +482,21 @@ impl<I: Idx, T> IndexVec<I, T> {
     pub fn get_mut(&mut self, index: I) -> Option<&mut T> {
         self.raw.get_mut(index.index())
     }
+
+    /// Return mutable references to two distinct elements, a and b. Panics if a == b.
+    #[inline]
+    pub fn pick2_mut(&mut self, a: I, b: I) -> (&mut T, &mut T) {
+        let (ai, bi) = (a.index(), b.index());
+        assert!(ai != bi);
+
+        if ai < bi {
+            let (c1, c2) = self.raw.split_at_mut(bi);
+            (&mut c1[ai], &mut c2[0])
+        } else {
+            let (c2, c1) = self.pick2_mut(b, a);
+            (c1, c2)
+        }
+    }
 }
 
 impl<I: Idx, T: Clone> IndexVec<I, T> {

src/librustc_mir/borrow_check/nll/region_infer/values.rs

@@ -9,7 +9,7 @@
 // except according to those terms.
 
 use std::rc::Rc;
-use rustc_data_structures::bitvec::BitMatrix;
+use rustc_data_structures::bitvec::SparseBitMatrix;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::indexed_vec::Idx;
 use rustc_data_structures::indexed_vec::IndexVec;
@@ -132,7 +132,7 @@ impl RegionValueElements {
 }
 
 /// A newtype for the integers that represent one of the possible
-/// elements in a region. These are the rows in the `BitMatrix` that
+/// elements in a region. These are the rows in the `SparseBitMatrix` that
 /// is used to store the values of all regions. They have the following
 /// convention:
 ///
@@ -184,18 +184,18 @@ impl ToElementIndex for RegionElementIndex {
 }
 
 /// Stores the values for a set of regions. These are stored in a
-/// compact `BitMatrix` representation, with one row per region
+/// compact `SparseBitMatrix` representation, with one row per region
 /// variable. The columns consist of either universal regions or
 /// points in the CFG.
 #[derive(Clone)]
 pub(super) struct RegionValues {
     elements: Rc<RegionValueElements>,
-    matrix: BitMatrix,
+    matrix: SparseBitMatrix<RegionVid, RegionElementIndex>,
 
     /// If cause tracking is enabled, maps from a pair (r, e)
     /// consisting of a region `r` that contains some element `e` to
     /// the reason that the element is contained. There should be an
-    /// entry for every bit set to 1 in `BitMatrix`.
+    /// entry for every bit set to 1 in `SparseBitMatrix`.
     causes: Option<CauseMap>,
 }
 
@@ -214,7 +214,8 @@ impl RegionValues {
 
         Self {
             elements: elements.clone(),
-            matrix: BitMatrix::new(num_region_variables, elements.num_elements()),
+            matrix: SparseBitMatrix::new(RegionVid::new(num_region_variables),
+                                         RegionElementIndex::new(elements.num_elements())),
             causes: if track_causes.0 {
                 Some(CauseMap::default())
             } else {
@@ -238,7 +239,7 @@ impl RegionValues {
     where
         F: FnOnce(&CauseMap) -> Cause,
     {
-        if self.matrix.add(r.index(), i.index()) {
+        if self.matrix.add(r, i) {
             debug!("add(r={:?}, i={:?})", r, self.elements.to_element(i));
 
             if let Some(causes) = &mut self.causes {
@@ -289,7 +290,7 @@ impl RegionValues {
         constraint_location: Location,
         constraint_span: Span,
     ) -> bool {
-        // We could optimize this by improving `BitMatrix::merge` so
+        // We could optimize this by improving `SparseBitMatrix::merge` so
         // it does not always merge an entire row. That would
         // complicate causal tracking though.
         debug!(
@@ -315,7 +316,7 @@ impl RegionValues {
     /// True if the region `r` contains the given element.
     pub(super) fn contains<E: ToElementIndex>(&self, r: RegionVid, elem: E) -> bool {
         let i = self.elements.index(elem);
-        self.matrix.contains(r.index(), i.index())
+        self.matrix.contains(r, i)
     }
 
     /// Iterate over the value of the region `r`, yielding up element
@@ -326,8 +327,8 @@ impl RegionValues {
         r: RegionVid,
     ) -> impl Iterator<Item = RegionElementIndex> + 'a {
         self.matrix
-            .iter(r.index())
-            .map(move |i| RegionElementIndex::new(i))
+            .iter(r)
+            .map(move |i| i)
     }
 
     /// Returns just the universal regions that are contained in a given region's value.