finish filling polonius context

transpose liveness matrix and record live regions at the end of MIR typeck

4 files changed, 40 insertions, 8 deletions

diff --git a/compiler/rustc_borrowck/src/polonius/liveness_constraints.rs b/compiler/rustc_borrowck/src/polonius/liveness_constraints.rs
index 6ed13313463..804a166014f 100644
--- a/compiler/rustc_borrowck/src/polonius/liveness_constraints.rs
+++ b/compiler/rustc_borrowck/src/polonius/liveness_constraints.rs
@@ -1,7 +1,10 @@
 use std::collections::BTreeMap;
 
+use rustc_index::bit_set::SparseBitMatrix;
+use rustc_index::interval::SparseIntervalMatrix;
 use rustc_middle::ty::relate::{self, Relate, RelateResult, TypeRelation};
 use rustc_middle::ty::{self, RegionVid, Ty, TyCtxt, TypeVisitable};
+use rustc_mir_dataflow::points::PointIndex;
 
 use super::{ConstraintDirection, PoloniusContext};
 use crate::universal_regions::UniversalRegions;
@@ -22,6 +25,25 @@ impl PoloniusContext {
         };
         extractor.relate(value, value).expect("Can't have a type error relating to itself");
     }
+
+    /// Unlike NLLs, in polonius we traverse the cfg to look for regions live across an edge, so we
+    /// need to transpose the "points where each region is live" matrix to a "live regions per point"
+    /// matrix.
+    // FIXME: avoid this conversion by always storing liveness data in this shape in the rest of
+    // borrowck.
+    pub(crate) fn record_live_regions_per_point(
+        &mut self,
+        num_regions: usize,
+        points_per_live_region: &SparseIntervalMatrix<RegionVid, PointIndex>,
+    ) {
+        let mut live_regions_per_point = SparseBitMatrix::new(num_regions);
+        for region in points_per_live_region.rows() {
+            for point in points_per_live_region.row(region).unwrap().iter() {
+                live_regions_per_point.insert(point, region);
+            }
+        }
+        self.live_regions = Some(live_regions_per_point);
+    }
 }
 
 /// Extracts variances for regions contained within types. Follows the same structure as
diff --git a/compiler/rustc_borrowck/src/polonius/mod.rs b/compiler/rustc_borrowck/src/polonius/mod.rs
index d86edd00725..5a95ac758fd 100644
--- a/compiler/rustc_borrowck/src/polonius/mod.rs
+++ b/compiler/rustc_borrowck/src/polonius/mod.rs
@@ -57,7 +57,7 @@ use crate::universal_regions::UniversalRegions;
 pub(crate) struct PoloniusContext {
     /// The set of regions that are live at a given point in the CFG, used to create localized
     /// outlives constraints between regions that are live at connected points in the CFG.
-    live_regions: SparseBitMatrix<PointIndex, RegionVid>,
+    live_regions: Option<SparseBitMatrix<PointIndex, RegionVid>>,
 
     /// The expected edge direction per live region: the kind of directed edge we'll create as
     /// liveness constraints depends on the variance of types with respect to each contained region.
@@ -79,11 +79,8 @@ enum ConstraintDirection {
 }
 
 impl PoloniusContext {
-    pub(crate) fn new(num_regions: usize) -> PoloniusContext {
-        Self {
-            live_region_variances: BTreeMap::new(),
-            live_regions: SparseBitMatrix::new(num_regions),
-        }
+    pub(crate) fn new() -> PoloniusContext {
+        Self { live_region_variances: BTreeMap::new(), live_regions: None }
     }
 
     /// Creates a constraint set for `-Zpolonius=next` by:
diff --git a/compiler/rustc_borrowck/src/region_infer/values.rs b/compiler/rustc_borrowck/src/region_infer/values.rs
index 0b0757f16ab..e567f3a8b0d 100644
--- a/compiler/rustc_borrowck/src/region_infer/values.rs
+++ b/compiler/rustc_borrowck/src/region_infer/values.rs
@@ -99,6 +99,14 @@ impl LivenessValues {
         }
     }
 
+    /// Returns the liveness matrix of points where each region is live. Panics if the liveness
+    /// values have been created without any per-point data (that is, for promoteds).
+    pub(crate) fn points(&self) -> &SparseIntervalMatrix<RegionVid, PointIndex> {
+        self.points
+            .as_ref()
+            .expect("this `LivenessValues` wasn't created using `with_specific_points`")
+    }
+
     /// Iterate through each region that has a value in this set.
     pub(crate) fn regions(&self) -> impl Iterator<Item = RegionVid> + '_ {
         self.points.as_ref().expect("use with_specific_points").rows()
diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs
index ea13e9b0fcf..3968900d047 100644
--- a/compiler/rustc_borrowck/src/type_check/mod.rs
+++ b/compiler/rustc_borrowck/src/type_check/mod.rs
@@ -150,8 +150,7 @@ pub(crate) fn type_check<'a, 'tcx>(
     debug!(?normalized_inputs_and_output);
 
     let mut polonius_context = if infcx.tcx.sess.opts.unstable_opts.polonius.is_next_enabled() {
-        let num_regions = infcx.num_region_vars();
-        Some(PoloniusContext::new(num_regions))
+        Some(PoloniusContext::new())
     } else {
         None
     };
@@ -187,6 +186,12 @@ pub(crate) fn type_check<'a, 'tcx>(
     let opaque_type_values =
         opaque_types::take_opaques_and_register_member_constraints(&mut typeck);
 
+    if let Some(polonius_context) = typeck.polonius_context.as_mut() {
+        let num_regions = infcx.num_region_vars();
+        let points_per_live_region = typeck.constraints.liveness_constraints.points();
+        polonius_context.record_live_regions_per_point(num_regions, points_per_live_region);
+    }
+
     MirTypeckResults {
         constraints,
         universal_region_relations,