introduce polonius context

This context struct will hold data to help creating localized
constraints:
- the live regions, with the shape matching a CFG walk, indexed per
  point
- the variance of these live regions, represented as the direction we'll
  add the appropriate

We also add this structure to the mir typeck to record liveness data,
and make it responsible for localized constraint creation.

3 files changed, 104 insertions, 47 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/mod.rs b/compiler/rustc_borrowck/src/polonius/mod.rs
index eee5e70efe3..16487de09e2 100644
--- a/compiler/rustc_borrowck/src/polonius/mod.rs
+++ b/compiler/rustc_borrowck/src/polonius/mod.rs
@@ -34,45 +34,84 @@
 //!
 
 mod constraints;
-pub(crate) use constraints::*;
 mod dump;
-pub(crate) use dump::dump_polonius_mir;
 pub(crate) mod legacy;
 
+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 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: 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(num_regions: usize) -> PoloniusContext {
+        Self {
+            live_region_variances: BTreeMap::new(),
+            live_regions: SparseBitMatrix::new(num_regions),
+        }
+    }
+
+    /// 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,
+        );
+        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
+    }
 }
 
 /// Propagate loans throughout the subset graph at a given point (with some subtleties around the
diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs
index f918f005a9b..ea13e9b0fcf 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,13 @@ 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))
+    } else {
+        None
+    };
+
     let mut typeck = TypeChecker {
         infcx,
         last_span: body.span,
@@ -162,6 +170,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 +187,12 @@ 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 }
+    MirTypeckResults {
+        constraints,
+        universal_region_relations,
+        opaque_type_values,
+        polonius_context,
+    }
 }
 
 #[track_caller]
@@ -546,6 +560,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 +570,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