Auto merge of #134670 - lqd:polonius-next-episode-4, r=jackh726

Compute liveness constraints in location-sensitive polonius

This continues the location-sensitive prototype. In this episode, we build the liveness constraints.

Reminder of the approach we're taking: we need variance data to create liveness edges in the forward/backward/both directions (respectively in the cases of covariance, contravariance, invariance) in the localized constraint graph.

This PR:
- introduces the holder for that, and for the liveness data in the correct shape: the transpose of what we're using today, "live regions per points".
- records use/drop live region variance during tracing
- records regular live region variance at the end of liveness
- records the correctly shaped live region per point matrix
- uses all of the above to compute the liveness constraints

(There's still technically one tiny part of the liveness owl left to do, but I'll leave it for a future PR: we also need to disable the NLL optimization that avoids computing liveness for locals whose types contain a region outliving a free region -- the existing constraints make it effectively live at all points; this doesn't work under polonius)

r? `@jackh726` cc `@matthewjasper`

9 files changed, 498 insertions, 128 deletions

diff --git a/compiler/rustc_borrowck/src/nll.rs b/compiler/rustc_borrowck/src/nll.rs
index abe27555b18..4428e695844 100644
--- a/compiler/rustc_borrowck/src/nll.rs
+++ b/compiler/rustc_borrowck/src/nll.rs
@@ -100,19 +100,23 @@ pub(crate) fn compute_regions<'a, 'tcx>(
     let elements = Rc::new(DenseLocationMap::new(body));
 
     // Run the MIR type-checker.
-    let MirTypeckResults { constraints, universal_region_relations, opaque_type_values } =
-        type_check::type_check(
-            infcx,
-            body,
-            promoted,
-            universal_regions,
-            location_table,
-            borrow_set,
-            &mut all_facts,
-            flow_inits,
-            move_data,
-            Rc::clone(&elements),
-        );
+    let MirTypeckResults {
+        constraints,
+        universal_region_relations,
+        opaque_type_values,
+        mut polonius_context,
+    } = type_check::type_check(
+        infcx,
+        body,
+        promoted,
+        universal_regions,
+        location_table,
+        borrow_set,
+        &mut all_facts,
+        flow_inits,
+        move_data,
+        Rc::clone(&elements),
+    );
 
     // Create the region inference context, taking ownership of the
     // region inference data that was contained in `infcx`, and the
@@ -141,12 +145,9 @@ pub(crate) fn compute_regions<'a, 'tcx>(
 
     // If requested for `-Zpolonius=next`, convert NLL constraints to localized outlives
     // constraints.
-    let localized_outlives_constraints =
-        if infcx.tcx.sess.opts.unstable_opts.polonius.is_next_enabled() {
-            Some(polonius::create_localized_constraints(&mut regioncx, body))
-        } else {
-            None
-        };
+    let localized_outlives_constraints = polonius_context
+        .as_mut()
+        .map(|polonius_context| polonius_context.create_localized_constraints(&mut regioncx, body));
 
     // If requested: dump NLL facts, and run legacy polonius analysis.
     let polonius_output = all_facts.as_ref().and_then(|all_facts| {
diff --git a/compiler/rustc_borrowck/src/polonius/liveness_constraints.rs b/compiler/rustc_borrowck/src/polonius/liveness_constraints.rs
new file mode 100644
index 00000000000..75ee29c9d0d
--- /dev/null
+++ b/compiler/rustc_borrowck/src/polonius/liveness_constraints.rs
@@ -0,0 +1,336 @@
+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};
+use rustc_mir_dataflow::points::PointIndex;
+
+use super::{
+    ConstraintDirection, LocalizedOutlivesConstraint, LocalizedOutlivesConstraintSet,
+    PoloniusContext,
+};
+use crate::region_infer::values::LivenessValues;
+use crate::universal_regions::UniversalRegions;
+
+impl PoloniusContext {
+    /// Record the variance of each region contained within the given value.
+    pub(crate) fn record_live_region_variance<'tcx>(
+        &mut self,
+        tcx: TyCtxt<'tcx>,
+        universal_regions: &UniversalRegions<'tcx>,
+        value: impl TypeVisitable<TyCtxt<'tcx>> + Relate<TyCtxt<'tcx>>,
+    ) {
+        let mut extractor = VarianceExtractor {
+            tcx,
+            ambient_variance: ty::Variance::Covariant,
+            directions: &mut self.live_region_variances,
+            universal_regions,
+        };
+        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
+/// constraints for loans that are propagated to the next statements.
+pub(super) fn create_liveness_constraints<'tcx>(
+    body: &Body<'tcx>,
+    liveness: &LivenessValues,
+    live_regions: &SparseBitMatrix<PointIndex, RegionVid>,
+    live_region_variances: &BTreeMap<RegionVid, ConstraintDirection>,
+    universal_regions: &UniversalRegions<'tcx>,
+    localized_outlives_constraints: &mut LocalizedOutlivesConstraintSet,
+) {
+    for (block, bb) in body.basic_blocks.iter_enumerated() {
+        let statement_count = bb.statements.len();
+        for statement_index in 0..=statement_count {
+            let current_location = Location { block, statement_index };
+            let current_point = liveness.point_from_location(current_location);
+
+            if statement_index < statement_count {
+                // Intra-block edges, straight line constraints from each point to its successor
+                // within the same block.
+                let next_location = Location { block, statement_index: statement_index + 1 };
+                let next_point = liveness.point_from_location(next_location);
+                propagate_loans_between_points(
+                    current_point,
+                    next_point,
+                    live_regions,
+                    live_region_variances,
+                    universal_regions,
+                    localized_outlives_constraints,
+                );
+            } else {
+                // Inter-block edges, from the block's terminator to each successor block's entry
+                // point.
+                for successor_block in bb.terminator().successors() {
+                    let next_location = Location { block: successor_block, statement_index: 0 };
+                    let next_point = liveness.point_from_location(next_location);
+                    propagate_loans_between_points(
+                        current_point,
+                        next_point,
+                        live_regions,
+                        live_region_variances,
+                        universal_regions,
+                        localized_outlives_constraints,
+                    );
+                }
+            }
+        }
+    }
+}
+
+/// Propagate loans within a region between two points in the CFG, if that region is live at both
+/// the source and target points.
+fn propagate_loans_between_points(
+    current_point: PointIndex,
+    next_point: PointIndex,
+    live_regions: &SparseBitMatrix<PointIndex, RegionVid>,
+    live_region_variances: &BTreeMap<RegionVid, ConstraintDirection>,
+    universal_regions: &UniversalRegions<'_>,
+    localized_outlives_constraints: &mut LocalizedOutlivesConstraintSet,
+) {
+    // Universal regions are semantically live at all points.
+    // Note: we always have universal regions but they're not always (or often) involved in the
+    // subset graph. For now, we emit all their edges unconditionally, but some of these subgraphs
+    // will be disconnected from the rest of the graph and thus, unnecessary.
+    //
+    // FIXME: only emit the edges of universal regions that existential regions can reach.
+    for region in universal_regions.universal_regions_iter() {
+        localized_outlives_constraints.push(LocalizedOutlivesConstraint {
+            source: region,
+            from: current_point,
+            target: region,
+            to: next_point,
+        });
+    }
+
+    let Some(current_live_regions) = live_regions.row(current_point) else {
+        // There are no constraints to add: there are no live regions at the current point.
+        return;
+    };
+    let Some(next_live_regions) = live_regions.row(next_point) else {
+        // There are no constraints to add: there are no live regions at the next point.
+        return;
+    };
+
+    for region in next_live_regions.iter() {
+        if !current_live_regions.contains(region) {
+            continue;
+        }
+
+        // `region` is indeed live at both points, add a constraint between them, according to
+        // variance.
+        if let Some(&direction) = live_region_variances.get(&region) {
+            add_liveness_constraint(
+                region,
+                current_point,
+                next_point,
+                direction,
+                localized_outlives_constraints,
+            );
+        } else {
+            // Note: there currently are cases related to promoted and const generics, where we
+            // don't yet have variance information (possibly about temporary regions created when
+            // typeck sanitizes the promoteds). Until that is done, we conservatively fallback to
+            // maximizing reachability by adding a bidirectional edge here. This will not limit
+            // traversal whatsoever, and thus propagate liveness when needed.
+            //
+            // FIXME: add the missing variance information and remove this fallback bidirectional
+            // edge.
+            let fallback = ConstraintDirection::Bidirectional;
+            add_liveness_constraint(
+                region,
+                current_point,
+                next_point,
+                fallback,
+                localized_outlives_constraints,
+            );
+        }
+    }
+}
+
+/// Adds `LocalizedOutlivesConstraint`s between two connected points, according to the given edge
+/// direction.
+fn add_liveness_constraint(
+    region: RegionVid,
+    current_point: PointIndex,
+    next_point: PointIndex,
+    direction: ConstraintDirection,
+    localized_outlives_constraints: &mut LocalizedOutlivesConstraintSet,
+) {
+    match direction {
+        ConstraintDirection::Forward => {
+            // Covariant cases: loans flow in the regular direction, from the current point to the
+            // next point.
+            localized_outlives_constraints.push(LocalizedOutlivesConstraint {
+                source: region,
+                from: current_point,
+                target: region,
+                to: next_point,
+            });
+        }
+        ConstraintDirection::Backward => {
+            // Contravariant cases: loans flow in the inverse direction, from the next point to the
+            // current point.
+            localized_outlives_constraints.push(LocalizedOutlivesConstraint {
+                source: region,
+                from: next_point,
+                target: region,
+                to: current_point,
+            });
+        }
+        ConstraintDirection::Bidirectional => {
+            // For invariant cases, loans can flow in both directions: we add both edges.
+            localized_outlives_constraints.push(LocalizedOutlivesConstraint {
+                source: region,
+                from: current_point,
+                target: region,
+                to: next_point,
+            });
+            localized_outlives_constraints.push(LocalizedOutlivesConstraint {
+                source: region,
+                from: next_point,
+                target: region,
+                to: current_point,
+            });
+        }
+    }
+}
+
+/// Extracts variances for regions contained within types. Follows the same structure as
+/// `rustc_infer`'s `Generalizer`: we try to relate a type with itself to track and extract the
+/// variances of regions.
+struct VarianceExtractor<'a, 'tcx> {
+    tcx: TyCtxt<'tcx>,
+    ambient_variance: ty::Variance,
+    directions: &'a mut BTreeMap<RegionVid, ConstraintDirection>,
+    universal_regions: &'a UniversalRegions<'tcx>,
+}
+
+impl<'tcx> VarianceExtractor<'_, 'tcx> {
+    fn record_variance(&mut self, region: ty::Region<'tcx>, variance: ty::Variance) {
+        // We're only interested in the variance of vars and free regions.
+        //
+        // Note: even if we currently bail for two cases of unexpected region kinds here, missing
+        // variance data is not a soundness problem: the regions with missing variance will still be
+        // present in the constraint graph as they are live, and liveness edges construction has a
+        // fallback for this case.
+        //
+        // FIXME: that being said, we need to investigate these cases better to not ignore regions
+        // in general.
+        if region.is_bound() {
+            // We ignore these because they cannot be turned into the vids we need.
+            return;
+        }
+
+        if region.is_erased() {
+            // These cannot be turned into a vid either, and we also ignore them: the fact that they
+            // show up here looks like either an issue upstream or a combination with unexpectedly
+            // continuing compilation too far when we're in a tainted by errors situation.
+            //
+            // FIXME: investigate the `generic_const_exprs` test that triggers this issue,
+            // `ui/const-generics/generic_const_exprs/issue-97047-ice-2.rs`
+            return;
+        }
+
+        let direction = match variance {
+            ty::Covariant => ConstraintDirection::Forward,
+            ty::Contravariant => ConstraintDirection::Backward,
+            ty::Invariant => ConstraintDirection::Bidirectional,
+            ty::Bivariant => {
+                // We don't add edges for bivariant cases.
+                return;
+            }
+        };
+
+        let region = self.universal_regions.to_region_vid(region);
+        self.directions
+            .entry(region)
+            .and_modify(|entry| {
+                // If there's already a recorded direction for this region, we combine the two:
+                // - combining the same direction is idempotent
+                // - combining different directions is trivially bidirectional
+                if entry != &direction {
+                    *entry = ConstraintDirection::Bidirectional;
+                }
+            })
+            .or_insert(direction);
+    }
+}
+
+impl<'tcx> TypeRelation<TyCtxt<'tcx>> for VarianceExtractor<'_, 'tcx> {
+    fn cx(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+
+    fn relate_with_variance<T: Relate<TyCtxt<'tcx>>>(
+        &mut self,
+        variance: ty::Variance,
+        _info: ty::VarianceDiagInfo<TyCtxt<'tcx>>,
+        a: T,
+        b: T,
+    ) -> RelateResult<'tcx, T> {
+        let old_ambient_variance = self.ambient_variance;
+        self.ambient_variance = self.ambient_variance.xform(variance);
+        let r = self.relate(a, b)?;
+        self.ambient_variance = old_ambient_variance;
+        Ok(r)
+    }
+
+    fn tys(&mut self, a: Ty<'tcx>, b: Ty<'tcx>) -> RelateResult<'tcx, Ty<'tcx>> {
+        assert_eq!(a, b); // we are misusing TypeRelation here; both LHS and RHS ought to be ==
+        relate::structurally_relate_tys(self, a, b)
+    }
+
+    fn regions(
+        &mut self,
+        a: ty::Region<'tcx>,
+        b: ty::Region<'tcx>,
+    ) -> RelateResult<'tcx, ty::Region<'tcx>> {
+        assert_eq!(a, b); // we are misusing TypeRelation here; both LHS and RHS ought to be ==
+        self.record_variance(a, self.ambient_variance);
+        Ok(a)
+    }
+
+    fn consts(
+        &mut self,
+        a: ty::Const<'tcx>,
+        b: ty::Const<'tcx>,
+    ) -> RelateResult<'tcx, ty::Const<'tcx>> {
+        assert_eq!(a, b); // we are misusing TypeRelation here; both LHS and RHS ought to be ==
+        relate::structurally_relate_consts(self, a, b)
+    }
+
+    fn binders<T>(
+        &mut self,
+        a: ty::Binder<'tcx, T>,
+        _: ty::Binder<'tcx, T>,
+    ) -> RelateResult<'tcx, ty::Binder<'tcx, T>>
+    where
+        T: Relate<TyCtxt<'tcx>>,
+    {
+        self.relate(a.skip_binder(), a.skip_binder())?;
+        Ok(a)
+    }
+}
diff --git a/compiler/rustc_borrowck/src/polonius/mod.rs b/compiler/rustc_borrowck/src/polonius/mod.rs
index eee5e70efe3..a853ff266a1 100644
--- a/compiler/rustc_borrowck/src/polonius/mod.rs
+++ b/compiler/rustc_borrowck/src/polonius/mod.rs
@@ -34,45 +34,88 @@
 //!
 
 mod constraints;
-pub(crate) use constraints::*;
 mod dump;
-pub(crate) use dump::dump_polonius_mir;
 pub(crate) mod legacy;
+mod liveness_constraints;
 
+use std::collections::BTreeMap;
+
+use rustc_index::bit_set::SparseBitMatrix;
 use rustc_middle::mir::{Body, Location};
+use rustc_middle::ty::RegionVid;
 use rustc_mir_dataflow::points::PointIndex;
 
+pub(crate) use self::constraints::*;
+pub(crate) use self::dump::dump_polonius_mir;
+use self::liveness_constraints::create_liveness_constraints;
 use crate::RegionInferenceContext;
 use crate::constraints::OutlivesConstraint;
 use crate::region_infer::values::LivenessValues;
 use crate::type_check::Locations;
-use crate::universal_regions::UniversalRegions;
 
-/// Creates a constraint set for `-Zpolonius=next` by:
-/// - converting NLL typeck constraints to be localized
-/// - encoding liveness constraints
-pub(crate) fn create_localized_constraints<'tcx>(
-    regioncx: &mut RegionInferenceContext<'tcx>,
-    body: &Body<'tcx>,
-) -> LocalizedOutlivesConstraintSet {
-    let mut localized_outlives_constraints = LocalizedOutlivesConstraintSet::default();
-    convert_typeck_constraints(
-        body,
-        regioncx.liveness_constraints(),
-        regioncx.outlives_constraints(),
-        &mut localized_outlives_constraints,
-    );
-    create_liveness_constraints(
-        body,
-        regioncx.liveness_constraints(),
-        regioncx.universal_regions(),
-        &mut localized_outlives_constraints,
-    );
-
-    // FIXME: here, we can trace loan reachability in the constraint graph and record this as loan
-    // liveness for the next step in the chain, the NLL loan scope and active loans computations.
-
-    localized_outlives_constraints
+/// This struct holds the data needed to create the Polonius localized constraints.
+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>>,
+
+    /// 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>,
+}
+
+/// The direction a constraint can flow into. Used to create liveness constraints according to
+/// variance.
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
+enum ConstraintDirection {
+    /// For covariant cases, we add a forward edge `O at P1 -> O at P2`.
+    Forward,
+
+    /// For contravariant cases, we add a backward edge `O at P2 -> O at P1`
+    Backward,
+
+    /// For invariant cases, we add both the forward and backward edges `O at P1 <-> O at P2`.
+    Bidirectional,
+}
+
+impl PoloniusContext {
+    pub(crate) fn new() -> PoloniusContext {
+        Self { live_region_variances: BTreeMap::new(), live_regions: None }
+    }
+
+    /// Creates a constraint set for `-Zpolonius=next` by:
+    /// - converting NLL typeck constraints to be localized
+    /// - encoding liveness constraints
+    pub(crate) fn create_localized_constraints<'tcx>(
+        &self,
+        regioncx: &RegionInferenceContext<'tcx>,
+        body: &Body<'tcx>,
+    ) -> LocalizedOutlivesConstraintSet {
+        let mut localized_outlives_constraints = LocalizedOutlivesConstraintSet::default();
+        convert_typeck_constraints(
+            body,
+            regioncx.liveness_constraints(),
+            regioncx.outlives_constraints(),
+            &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_region_variances,
+            regioncx.universal_regions(),
+            &mut localized_outlives_constraints,
+        );
+
+        // FIXME: here, we can trace loan reachability in the constraint graph and record this as loan
+        // liveness for the next step in the chain, the NLL loan scope and active loans computations.
+
+        localized_outlives_constraints
+    }
 }
 
 /// Propagate loans throughout the subset graph at a given point (with some subtleties around the
@@ -109,72 +152,3 @@ fn convert_typeck_constraints<'tcx>(
         }
     }
 }
-
-/// Propagate loans throughout the CFG: for each statement in the MIR, create localized outlives
-/// constraints for loans that are propagated to the next statements.
-pub(crate) fn create_liveness_constraints<'tcx>(
-    body: &Body<'tcx>,
-    liveness: &LivenessValues,
-    universal_regions: &UniversalRegions<'tcx>,
-    localized_outlives_constraints: &mut LocalizedOutlivesConstraintSet,
-) {
-    for (block, bb) in body.basic_blocks.iter_enumerated() {
-        let statement_count = bb.statements.len();
-        for statement_index in 0..=statement_count {
-            let current_location = Location { block, statement_index };
-            let current_point = liveness.point_from_location(current_location);
-
-            if statement_index < statement_count {
-                // Intra-block edges, straight line constraints from each point to its successor
-                // within the same block.
-                let next_location = Location { block, statement_index: statement_index + 1 };
-                let next_point = liveness.point_from_location(next_location);
-                propagate_loans_between_points(
-                    current_point,
-                    next_point,
-                    liveness,
-                    universal_regions,
-                    localized_outlives_constraints,
-                );
-            } else {
-                // Inter-block edges, from the block's terminator to each successor block's entry
-                // point.
-                for successor_block in bb.terminator().successors() {
-                    let next_location = Location { block: successor_block, statement_index: 0 };
-                    let next_point = liveness.point_from_location(next_location);
-                    propagate_loans_between_points(
-                        current_point,
-                        next_point,
-                        liveness,
-                        universal_regions,
-                        localized_outlives_constraints,
-                    );
-                }
-            }
-        }
-    }
-}
-
-/// Propagate loans within a region between two points in the CFG, if that region is live at both
-/// the source and target points.
-fn propagate_loans_between_points(
-    current_point: PointIndex,
-    next_point: PointIndex,
-    _liveness: &LivenessValues,
-    universal_regions: &UniversalRegions<'_>,
-    localized_outlives_constraints: &mut LocalizedOutlivesConstraintSet,
-) {
-    // Universal regions are semantically live at all points.
-    // Note: we always have universal regions but they're not always (or often) involved in the
-    // subset graph. For now, we emit all their edges unconditionally, but some of these subgraphs
-    // will be disconnected from the rest of the graph and thus, unnecessary.
-    // FIXME: only emit the edges of universal regions that existential regions can reach.
-    for region in universal_regions.universal_regions_iter() {
-        localized_outlives_constraints.push(LocalizedOutlivesConstraint {
-            source: region,
-            from: current_point,
-            target: region,
-            to: next_point,
-        });
-    }
-}
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/liveness/mod.rs b/compiler/rustc_borrowck/src/type_check/liveness/mod.rs
index 683293bf828..3e9900cce5f 100644
--- a/compiler/rustc_borrowck/src/type_check/liveness/mod.rs
+++ b/compiler/rustc_borrowck/src/type_check/liveness/mod.rs
@@ -3,6 +3,7 @@ use rustc_data_structures::fx::FxHashSet;
 use rustc_middle::mir::visit::{TyContext, Visitor};
 use rustc_middle::mir::{Body, Local, Location, SourceInfo};
 use rustc_middle::span_bug;
+use rustc_middle::ty::relate::Relate;
 use rustc_middle::ty::visit::TypeVisitable;
 use rustc_middle::ty::{GenericArgsRef, Region, RegionVid, Ty, TyCtxt};
 use rustc_mir_dataflow::ResultsCursor;
@@ -13,6 +14,7 @@ use tracing::debug;
 
 use super::TypeChecker;
 use crate::constraints::OutlivesConstraintSet;
+use crate::polonius::PoloniusContext;
 use crate::region_infer::values::LivenessValues;
 use crate::universal_regions::UniversalRegions;
 
@@ -56,7 +58,13 @@ pub(super) fn generate<'a, 'tcx>(
 
     // Mark regions that should be live where they appear within rvalues or within a call: like
     // args, regions, and types.
-    record_regular_live_regions(typeck.tcx(), &mut typeck.constraints.liveness_constraints, body);
+    record_regular_live_regions(
+        typeck.tcx(),
+        &mut typeck.constraints.liveness_constraints,
+        &typeck.universal_regions,
+        &mut typeck.polonius_context,
+        body,
+    );
 }
 
 // The purpose of `compute_relevant_live_locals` is to define the subset of `Local`
@@ -130,9 +138,12 @@ fn regions_that_outlive_free_regions<'tcx>(
 fn record_regular_live_regions<'tcx>(
     tcx: TyCtxt<'tcx>,
     liveness_constraints: &mut LivenessValues,
+    universal_regions: &UniversalRegions<'tcx>,
+    polonius_context: &mut Option<PoloniusContext>,
     body: &Body<'tcx>,
 ) {
-    let mut visitor = LiveVariablesVisitor { tcx, liveness_constraints };
+    let mut visitor =
+        LiveVariablesVisitor { tcx, liveness_constraints, universal_regions, polonius_context };
     for (bb, data) in body.basic_blocks.iter_enumerated() {
         visitor.visit_basic_block_data(bb, data);
     }
@@ -142,6 +153,8 @@ fn record_regular_live_regions<'tcx>(
 struct LiveVariablesVisitor<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     liveness_constraints: &'a mut LivenessValues,
+    universal_regions: &'a UniversalRegions<'tcx>,
+    polonius_context: &'a mut Option<PoloniusContext>,
 }
 
 impl<'a, 'tcx> Visitor<'tcx> for LiveVariablesVisitor<'a, 'tcx> {
@@ -184,12 +197,17 @@ impl<'a, 'tcx> LiveVariablesVisitor<'a, 'tcx> {
     /// all regions appearing in the type of `value` must be live at `location`.
     fn record_regions_live_at<T>(&mut self, value: T, location: Location)
     where
-        T: TypeVisitable<TyCtxt<'tcx>>,
+        T: TypeVisitable<TyCtxt<'tcx>> + Relate<TyCtxt<'tcx>>,
     {
         debug!("record_regions_live_at(value={:?}, location={:?})", value, location);
         self.tcx.for_each_free_region(&value, |live_region| {
             let live_region_vid = live_region.as_var();
             self.liveness_constraints.add_location(live_region_vid, location);
         });
+
+        // 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);
+        }
     }
 }
diff --git a/compiler/rustc_borrowck/src/type_check/liveness/trace.rs b/compiler/rustc_borrowck/src/type_check/liveness/trace.rs
index f510d193dd9..2c658edc41c 100644
--- a/compiler/rustc_borrowck/src/type_check/liveness/trace.rs
+++ b/compiler/rustc_borrowck/src/type_check/liveness/trace.rs
@@ -5,6 +5,7 @@ use rustc_infer::infer::canonical::QueryRegionConstraints;
 use rustc_infer::infer::outlives::for_liveness;
 use rustc_middle::mir::{BasicBlock, Body, ConstraintCategory, Local, Location};
 use rustc_middle::traits::query::DropckOutlivesResult;
+use rustc_middle::ty::relate::Relate;
 use rustc_middle::ty::{Ty, TyCtxt, TypeVisitable, TypeVisitableExt};
 use rustc_mir_dataflow::ResultsCursor;
 use rustc_mir_dataflow::impls::MaybeInitializedPlaces;
@@ -532,11 +533,7 @@ impl<'tcx> LivenessContext<'_, '_, '_, 'tcx> {
 
     /// Stores the result that all regions in `value` are live for the
     /// points `live_at`.
-    fn add_use_live_facts_for(
-        &mut self,
-        value: impl TypeVisitable<TyCtxt<'tcx>>,
-        live_at: &IntervalSet<PointIndex>,
-    ) {
+    fn add_use_live_facts_for(&mut self, value: Ty<'tcx>, live_at: &IntervalSet<PointIndex>) {
         debug!("add_use_live_facts_for(value={:?})", value);
         Self::make_all_regions_live(self.elements, self.typeck, value, live_at);
     }
@@ -603,7 +600,7 @@ impl<'tcx> LivenessContext<'_, '_, '_, 'tcx> {
     fn make_all_regions_live(
         elements: &DenseLocationMap,
         typeck: &mut TypeChecker<'_, 'tcx>,
-        value: impl TypeVisitable<TyCtxt<'tcx>>,
+        value: impl TypeVisitable<TyCtxt<'tcx>> + Relate<TyCtxt<'tcx>>,
         live_at: &IntervalSet<PointIndex>,
     ) {
         debug!("make_all_regions_live(value={:?})", value);
@@ -621,6 +618,15 @@ impl<'tcx> LivenessContext<'_, '_, '_, 'tcx> {
                 typeck.constraints.liveness_constraints.add_points(live_region_vid, live_at);
             },
         });
+
+        // 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(
+                typeck.infcx.tcx,
+                typeck.universal_regions,
+                value,
+            );
+        }
     }
 
     fn compute_drop_data(typeck: &TypeChecker<'_, 'tcx>, dropped_ty: Ty<'tcx>) -> DropData<'tcx> {
diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs
index f918f005a9b..3968900d047 100644
--- a/compiler/rustc_borrowck/src/type_check/mod.rs
+++ b/compiler/rustc_borrowck/src/type_check/mod.rs
@@ -50,6 +50,7 @@ use crate::diagnostics::UniverseInfo;
 use crate::facts::AllFacts;
 use crate::location::LocationTable;
 use crate::member_constraints::MemberConstraintSet;
+use crate::polonius::PoloniusContext;
 use crate::region_infer::TypeTest;
 use crate::region_infer::values::{LivenessValues, PlaceholderIndex, PlaceholderIndices};
 use crate::renumber::RegionCtxt;
@@ -148,6 +149,12 @@ 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())
+    } else {
+        None
+    };
+
     let mut typeck = TypeChecker {
         infcx,
         last_span: body.span,
@@ -162,6 +169,7 @@ pub(crate) fn type_check<'a, 'tcx>(
         all_facts,
         borrow_set,
         constraints: &mut constraints,
+        polonius_context: &mut polonius_context,
     };
 
     typeck.check_user_type_annotations();
@@ -178,7 +186,18 @@ pub(crate) fn type_check<'a, 'tcx>(
     let opaque_type_values =
         opaque_types::take_opaques_and_register_member_constraints(&mut typeck);
 
-    MirTypeckResults { constraints, universal_region_relations, opaque_type_values }
+    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,
+        opaque_type_values,
+        polonius_context,
+    }
 }
 
 #[track_caller]
@@ -546,6 +565,8 @@ struct TypeChecker<'a, 'tcx> {
     all_facts: &'a mut Option<AllFacts>,
     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>,
 }
 
 /// Holder struct for passing results from MIR typeck to the rest of the non-lexical regions
@@ -554,6 +575,7 @@ pub(crate) struct MirTypeckResults<'tcx> {
     pub(crate) constraints: MirTypeckRegionConstraints<'tcx>,
     pub(crate) universal_region_relations: Frozen<UniversalRegionRelations<'tcx>>,
     pub(crate) opaque_type_values: FxIndexMap<OpaqueTypeKey<'tcx>, OpaqueHiddenType<'tcx>>,
+    pub(crate) polonius_context: Option<PoloniusContext>,
 }
 
 /// A collection of region constraints that must be satisfied for the
diff --git a/compiler/rustc_borrowck/src/universal_regions.rs b/compiler/rustc_borrowck/src/universal_regions.rs
index 1ac45cbea38..3dc4569c57b 100644
--- a/compiler/rustc_borrowck/src/universal_regions.rs
+++ b/compiler/rustc_borrowck/src/universal_regions.rs
@@ -337,7 +337,7 @@ impl<'tcx> UniversalRegions<'tcx> {
         self.indices.indices.iter().map(|(&r, &v)| (r, v))
     }
 
-    /// See `UniversalRegionIndices::to_region_vid`.
+    /// See [UniversalRegionIndices::to_region_vid].
     pub(crate) fn to_region_vid(&self, r: ty::Region<'tcx>) -> RegionVid {
         self.indices.to_region_vid(r)
     }
diff --git a/compiler/rustc_index/src/bit_set.rs b/compiler/rustc_index/src/bit_set.rs
index 664b77fd49e..38e2dbbde7d 100644
--- a/compiler/rustc_index/src/bit_set.rs
+++ b/compiler/rustc_index/src/bit_set.rs
@@ -179,7 +179,12 @@ impl<T: Idx> BitSet<T> {
     /// Insert `elem`. Returns whether the set has changed.
     #[inline]
     pub fn insert(&mut self, elem: T) -> bool {
-        assert!(elem.index() < self.domain_size);
+        assert!(
+            elem.index() < self.domain_size,
+            "inserting element at index {} but domain size is {}",
+            elem.index(),
+            self.domain_size,
+        );
         let (word_index, mask) = word_index_and_mask(elem);
         let word_ref = &mut self.words[word_index];
         let word = *word_ref;