Auto merge of #53629 - nnethercote:lazier-SparseBitMatrix, r=nikomatsakis

Lazier sparse bit matrix

A small NLL win.

r? @nikomatsakis

1 files changed, 38 insertions, 34 deletions

diff --git a/src/librustc_data_structures/bitvec.rs b/src/librustc_data_structures/bitvec.rs
index 49ab3e58812..642d24f4831 100644
--- a/src/librustc_data_structures/bitvec.rs
+++ b/src/librustc_data_structures/bitvec.rs
@@ -318,31 +318,41 @@ impl<R: Idx, C: Idx> BitMatrix<R, C> {
     }
 }
 
-/// A moderately sparse bit matrix: rows are appended lazily, but columns
-/// within appended rows are instantiated fully upon creation.
+/// A moderately sparse bit matrix, in which rows are instantiated lazily.
+///
+/// Initially, every row has no explicit representation. If any bit within a
+/// row is set, the entire row is instantiated as
+/// `Some(<full-column-width-BitArray>)`. Furthermore, any previously
+/// uninstantiated rows prior to it will be instantiated as `None`. Those prior
+/// rows may themselves become fully instantiated later on if any of their bits
+/// are set.
 #[derive(Clone, Debug)]
 pub struct SparseBitMatrix<R, C>
 where
     R: Idx,
     C: Idx,
 {
-    columns: usize,
-    vector: IndexVec<R, BitArray<C>>,
+    num_columns: usize,
+    rows: IndexVec<R, Option<BitArray<C>>>,
 }
 
 impl<R: Idx, C: Idx> SparseBitMatrix<R, C> {
     /// Create a new empty sparse bit matrix with no rows or columns.
-    pub fn new(columns: usize) -> Self {
+    pub fn new(num_columns: usize) -> Self {
         Self {
-            columns,
-            vector: IndexVec::new(),
+            num_columns,
+            rows: IndexVec::new(),
         }
     }
 
-    fn ensure_row(&mut self, row: R) {
-        let columns = self.columns;
-        self.vector
-            .ensure_contains_elem(row, || BitArray::new(columns));
+    fn ensure_row(&mut self, row: R) -> &mut BitArray<C> {
+        // Instantiate any missing rows up to and including row `row` with an
+        // empty BitArray.
+        self.rows.ensure_contains_elem(row, || None);
+
+        // Then replace row `row` with a full BitArray if necessary.
+        let num_columns = self.num_columns;
+        self.rows[row].get_or_insert_with(|| BitArray::new(num_columns))
     }
 
     /// Sets the cell at `(row, column)` to true. Put another way, insert
@@ -350,8 +360,7 @@ impl<R: Idx, C: Idx> SparseBitMatrix<R, C> {
     ///
     /// Returns true if this changed the matrix, and false otherwise.
     pub fn add(&mut self, row: R, column: C) -> bool {
-        self.ensure_row(row);
-        self.vector[row].insert(column)
+        self.ensure_row(row).insert(column)
     }
 
     /// Do the bits from `row` contain `column`? Put another way, is
@@ -359,7 +368,7 @@ impl<R: Idx, C: Idx> SparseBitMatrix<R, C> {
     /// if the matrix represents (transitive) reachability, can
     /// `row` reach `column`?
     pub fn contains(&self, row: R, column: C) -> bool {
-        self.vector.get(row).map_or(false, |r| r.contains(column))
+        self.row(row).map_or(false, |r| r.contains(column))
     }
 
     /// Add the bits from row `read` to the bits from row `write`,
@@ -370,49 +379,44 @@ impl<R: Idx, C: Idx> SparseBitMatrix<R, C> {
     /// `write` can reach everything that `read` can (and
     /// potentially more).
     pub fn merge(&mut self, read: R, write: R) -> bool {
-        if read == write || self.vector.get(read).is_none() {
+        if read == write || self.row(read).is_none() {
             return false;
         }
 
         self.ensure_row(write);
-        let (bitvec_read, bitvec_write) = self.vector.pick2_mut(read, write);
-        bitvec_write.merge(bitvec_read)
+        if let (Some(bitvec_read), Some(bitvec_write)) = self.rows.pick2_mut(read, write) {
+            bitvec_write.merge(bitvec_read)
+        } else {
+            unreachable!()
+        }
     }
 
     /// Merge a row, `from`, into the `into` row.
     pub fn merge_into(&mut self, into: R, from: &BitArray<C>) -> bool {
-        self.ensure_row(into);
-        self.vector[into].merge(from)
+        self.ensure_row(into).merge(from)
     }
 
     /// Add all bits to the given row.
     pub fn add_all(&mut self, row: R) {
-        self.ensure_row(row);
-        self.vector[row].insert_all();
-    }
-
-    /// Number of elements in the matrix.
-    pub fn len(&self) -> usize {
-        self.vector.len()
+        self.ensure_row(row).insert_all();
     }
 
     pub fn rows(&self) -> impl Iterator<Item = R> {
-        self.vector.indices()
+        self.rows.indices()
     }
 
     /// 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.get(row).into_iter().flat_map(|r| r.iter())
-    }
-
-    /// Iterates through each row and the accompanying bit set.
-    pub fn iter_enumerated<'a>(&'a self) -> impl Iterator<Item = (R, &'a BitArray<C>)> + 'a {
-        self.vector.iter_enumerated()
+        self.row(row).into_iter().flat_map(|r| r.iter())
     }
 
     pub fn row(&self, row: R) -> Option<&BitArray<C>> {
-        self.vector.get(row)
+        if let Some(Some(row)) = self.rows.get(row) {
+            Some(row)
+        } else {
+            None
+        }
     }
 }