split polonius context into per-phase data

- describe how that data flows during borrowck
- prepares for recording some liveness data for diagnostics, not just
  for the main analysis

5 files changed, 70 insertions, 50 deletions

diff --git a/compiler/rustc_borrowck/src/polonius/liveness_constraints.rs b/compiler/rustc_borrowck/src/polonius/liveness_constraints.rs
index 75ee29c9d0d..6ab09f731c0 100644
--- a/compiler/rustc_borrowck/src/polonius/liveness_constraints.rs
+++ b/compiler/rustc_borrowck/src/polonius/liveness_constraints.rs
@@ -1,7 +1,6 @@
 use std::collections::BTreeMap;
 
 use rustc_index::bit_set::SparseBitMatrix;
-use rustc_index::interval::SparseIntervalMatrix;
 use rustc_middle::mir::{Body, Location};
 use rustc_middle::ty::relate::{self, Relate, RelateResult, TypeRelation};
 use rustc_middle::ty::{self, RegionVid, Ty, TyCtxt, TypeVisitable};
@@ -9,12 +8,12 @@ use rustc_mir_dataflow::points::PointIndex;
 
 use super::{
     ConstraintDirection, LocalizedOutlivesConstraint, LocalizedOutlivesConstraintSet,
-    PoloniusContext,
+    PoloniusLivenessContext,
 };
 use crate::region_infer::values::LivenessValues;
 use crate::universal_regions::UniversalRegions;
 
-impl PoloniusContext {
+impl PoloniusLivenessContext {
     /// Record the variance of each region contained within the given value.
     pub(crate) fn record_live_region_variance<'tcx>(
         &mut self,
@@ -30,25 +29,6 @@ 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);
-    }
 }
 
 /// Propagate loans throughout the CFG: for each statement in the MIR, create localized outlives
diff --git a/compiler/rustc_borrowck/src/polonius/mod.rs b/compiler/rustc_borrowck/src/polonius/mod.rs
index 502c868194a..11414a5719d 100644
--- a/compiler/rustc_borrowck/src/polonius/mod.rs
+++ b/compiler/rustc_borrowck/src/polonius/mod.rs
@@ -32,6 +32,16 @@
 //! - <https://smallcultfollowing.com/babysteps/blog/2023/09/22/polonius-part-1/>
 //! - <https://smallcultfollowing.com/babysteps/blog/2023/09/29/polonius-part-2/>
 //!
+//!
+//! Data flows like this:
+//! 1) during MIR typeck, record liveness data needed later: live region variances, as well as the
+//!    usual NLL liveness data (just computed on more locals). That's the [PoloniusLivenessContext].
+//! 2) once that is done, variance data is transferred, and the NLL region liveness is converted to
+//!    the polonius shape. That's the main [PoloniusContext].
+//! 3) during region inference, that data and the NLL outlives constraints are used to create the
+//!    localized outlives constraints, as described above.
+//! 4) transfer these constraints back to the main borrowck procedure: it handles computing errors
+//!    and diagnostics, debugging and MIR dumping concerns.
 
 mod constraints;
 mod dump;
@@ -43,6 +53,7 @@ mod typeck_constraints;
 use std::collections::BTreeMap;
 
 use rustc_index::bit_set::SparseBitMatrix;
+use rustc_index::interval::SparseIntervalMatrix;
 use rustc_middle::mir::Body;
 use rustc_middle::ty::{RegionVid, TyCtxt};
 use rustc_mir_dataflow::points::PointIndex;
@@ -57,11 +68,21 @@ use crate::{BorrowSet, RegionInferenceContext};
 
 pub(crate) type LiveLoans = SparseBitMatrix<PointIndex, BorrowIndex>;
 
-/// This struct holds the data needed to create the Polonius localized constraints.
+/// This struct holds the liveness data created during MIR typeck, and which will be used later in
+/// the process, to compute the polonius localized constraints.
+#[derive(Default)]
+pub(crate) struct PoloniusLivenessContext {
+    /// 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.
+    live_region_variances: BTreeMap<RegionVid, ConstraintDirection>,
+}
+
+/// This struct holds the data needed to create the Polonius localized constraints. Its data is
+/// transferred and converted from the [PoloniusLivenessContext] at the end of MIR typeck.
 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: Option<SparseBitMatrix<PointIndex, RegionVid>>,
+    live_regions: 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.
@@ -83,8 +104,27 @@ enum ConstraintDirection {
 }
 
 impl PoloniusContext {
-    pub(crate) fn new() -> PoloniusContext {
-        Self { live_region_variances: BTreeMap::new(), live_regions: None }
+    /// 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 create_from_liveness(
+        liveness_context: PoloniusLivenessContext,
+        num_regions: usize,
+        points_per_live_region: &SparseIntervalMatrix<RegionVid, PointIndex>,
+    ) -> PoloniusContext {
+        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);
+            }
+        }
+
+        PoloniusContext {
+            live_regions: live_regions_per_point,
+            live_region_variances: liveness_context.live_region_variances,
+        }
     }
 
     /// Computes live loans using the set of loans model for `-Zpolonius=next`.
@@ -112,13 +152,10 @@ impl PoloniusContext {
             &mut localized_outlives_constraints,
         );
 
-        let live_regions = self.live_regions.as_ref().expect(
-            "live regions per-point data should have been created at the end of MIR typeck",
-        );
         create_liveness_constraints(
             body,
             regioncx.liveness_constraints(),
-            live_regions,
+            &self.live_regions,
             &self.live_region_variances,
             regioncx.universal_regions(),
             &mut localized_outlives_constraints,
diff --git a/compiler/rustc_borrowck/src/type_check/liveness/mod.rs b/compiler/rustc_borrowck/src/type_check/liveness/mod.rs
index 4e0b2a4e296..23b8bbb38e8 100644
--- a/compiler/rustc_borrowck/src/type_check/liveness/mod.rs
+++ b/compiler/rustc_borrowck/src/type_check/liveness/mod.rs
@@ -14,7 +14,7 @@ use tracing::debug;
 
 use super::TypeChecker;
 use crate::constraints::OutlivesConstraintSet;
-use crate::polonius::PoloniusContext;
+use crate::polonius::PoloniusLivenessContext;
 use crate::region_infer::values::LivenessValues;
 use crate::universal_regions::UniversalRegions;
 
@@ -70,7 +70,7 @@ pub(super) fn generate<'a, 'tcx>(
         typeck.tcx(),
         &mut typeck.constraints.liveness_constraints,
         &typeck.universal_regions,
-        &mut typeck.polonius_context,
+        &mut typeck.polonius_liveness,
         body,
     );
 }
@@ -147,11 +147,11 @@ fn record_regular_live_regions<'tcx>(
     tcx: TyCtxt<'tcx>,
     liveness_constraints: &mut LivenessValues,
     universal_regions: &UniversalRegions<'tcx>,
-    polonius_context: &mut Option<PoloniusContext>,
+    polonius_liveness: &mut Option<PoloniusLivenessContext>,
     body: &Body<'tcx>,
 ) {
     let mut visitor =
-        LiveVariablesVisitor { tcx, liveness_constraints, universal_regions, polonius_context };
+        LiveVariablesVisitor { tcx, liveness_constraints, universal_regions, polonius_liveness };
     for (bb, data) in body.basic_blocks.iter_enumerated() {
         visitor.visit_basic_block_data(bb, data);
     }
@@ -162,7 +162,7 @@ struct LiveVariablesVisitor<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     liveness_constraints: &'a mut LivenessValues,
     universal_regions: &'a UniversalRegions<'tcx>,
-    polonius_context: &'a mut Option<PoloniusContext>,
+    polonius_liveness: &'a mut Option<PoloniusLivenessContext>,
 }
 
 impl<'a, 'tcx> Visitor<'tcx> for LiveVariablesVisitor<'a, 'tcx> {
@@ -214,8 +214,8 @@ impl<'a, 'tcx> LiveVariablesVisitor<'a, 'tcx> {
         });
 
         // When using `-Zpolonius=next`, we record the variance of each live region.
-        if let Some(polonius_context) = self.polonius_context {
-            polonius_context.record_live_region_variance(self.tcx, self.universal_regions, value);
+        if let Some(polonius_liveness) = self.polonius_liveness {
+            polonius_liveness.record_live_region_variance(self.tcx, self.universal_regions, value);
         }
     }
 }
diff --git a/compiler/rustc_borrowck/src/type_check/liveness/trace.rs b/compiler/rustc_borrowck/src/type_check/liveness/trace.rs
index c564d85616e..62d49a62744 100644
--- a/compiler/rustc_borrowck/src/type_check/liveness/trace.rs
+++ b/compiler/rustc_borrowck/src/type_check/liveness/trace.rs
@@ -580,8 +580,8 @@ impl<'tcx> LivenessContext<'_, '_, '_, 'tcx> {
         });
 
         // When using `-Zpolonius=next`, we record the variance of each live region.
-        if let Some(polonius_context) = typeck.polonius_context {
-            polonius_context.record_live_region_variance(
+        if let Some(polonius_liveness) = typeck.polonius_liveness.as_mut() {
+            polonius_liveness.record_live_region_variance(
                 typeck.infcx.tcx,
                 typeck.universal_regions,
                 value,
diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs
index a2ef5588f48..45120324a0c 100644
--- a/compiler/rustc_borrowck/src/type_check/mod.rs
+++ b/compiler/rustc_borrowck/src/type_check/mod.rs
@@ -48,8 +48,8 @@ use crate::borrow_set::BorrowSet;
 use crate::constraints::{OutlivesConstraint, OutlivesConstraintSet};
 use crate::diagnostics::UniverseInfo;
 use crate::member_constraints::MemberConstraintSet;
-use crate::polonius::PoloniusContext;
 use crate::polonius::legacy::{PoloniusFacts, PoloniusLocationTable};
+use crate::polonius::{PoloniusContext, PoloniusLivenessContext};
 use crate::region_infer::TypeTest;
 use crate::region_infer::values::{LivenessValues, PlaceholderIndex, PlaceholderIndices};
 use crate::renumber::RegionCtxt;
@@ -148,8 +148,8 @@ 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() {
-        Some(PoloniusContext::new())
+    let polonius_liveness = if infcx.tcx.sess.opts.unstable_opts.polonius.is_next_enabled() {
+        Some(PoloniusLivenessContext::default())
     } else {
         None
     };
@@ -168,7 +168,7 @@ pub(crate) fn type_check<'a, 'tcx>(
         polonius_facts,
         borrow_set,
         constraints: &mut constraints,
-        polonius_context: &mut polonius_context,
+        polonius_liveness,
     };
 
     typeck.check_user_type_annotations();
@@ -185,11 +185,14 @@ 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);
-    }
+    // We're done with typeck, we can finalize the polonius liveness context for region inference.
+    let polonius_context = typeck.polonius_liveness.take().map(|liveness_context| {
+        PoloniusContext::create_from_liveness(
+            liveness_context,
+            infcx.num_region_vars(),
+            typeck.constraints.liveness_constraints.points(),
+        )
+    });
 
     MirTypeckResults {
         constraints,
@@ -564,8 +567,8 @@ struct TypeChecker<'a, 'tcx> {
     polonius_facts: &'a mut Option<PoloniusFacts>,
     borrow_set: &'a BorrowSet<'tcx>,
     constraints: &'a mut MirTypeckRegionConstraints<'tcx>,
-    /// When using `-Zpolonius=next`, the helper data used to create polonius constraints.
-    polonius_context: &'a mut Option<PoloniusContext>,
+    /// When using `-Zpolonius=next`, the liveness helper data used to create polonius constraints.
+    polonius_liveness: Option<PoloniusLivenessContext>,
 }
 
 /// Holder struct for passing results from MIR typeck to the rest of the non-lexical regions