localize typeck constraints

it's still partially a skeleton, but works well enough for almost all tests to pass

compiler/rustc_borrowck/src/nll.rs

@@ -142,9 +142,9 @@ pub(crate) fn compute_regions<'a, 'tcx>(
 
     // If requested for `-Zpolonius=next`, convert NLL constraints to localized outlives
     // constraints.
-    let localized_outlives_constraints = polonius_context
-        .as_mut()
-        .map(|polonius_context| polonius_context.create_localized_constraints(&mut regioncx, body));
+    let localized_outlives_constraints = polonius_context.as_mut().map(|polonius_context| {
+        polonius_context.create_localized_constraints(infcx.tcx, &regioncx, body)
+    });
 
     // If requested: dump NLL facts, and run legacy polonius analysis.
     let polonius_output = all_facts.as_ref().and_then(|all_facts| {

compiler/rustc_borrowck/src/polonius/mod.rs

@@ -43,7 +43,7 @@ use std::collections::BTreeMap;
 
 use rustc_index::bit_set::SparseBitMatrix;
 use rustc_middle::mir::Body;
-use rustc_middle::ty::RegionVid;
+use rustc_middle::ty::{RegionVid, TyCtxt};
 use rustc_mir_dataflow::points::PointIndex;
 
 pub(crate) use self::constraints::*;
@@ -87,14 +87,17 @@ impl PoloniusContext {
     /// - encoding liveness constraints
     pub(crate) fn create_localized_constraints<'tcx>(
         &self,
+        tcx: TyCtxt<'tcx>,
         regioncx: &RegionInferenceContext<'tcx>,
         body: &Body<'tcx>,
     ) -> LocalizedOutlivesConstraintSet {
         let mut localized_outlives_constraints = LocalizedOutlivesConstraintSet::default();
         convert_typeck_constraints(
+            tcx,
             body,
             regioncx.liveness_constraints(),
             regioncx.outlives_constraints(),
+            regioncx.universal_regions(),
             &mut localized_outlives_constraints,
         );
 

compiler/rustc_borrowck/src/polonius/typeck_constraints.rs

@@ -1,16 +1,22 @@
-use rustc_middle::mir::{Body, Location};
+use rustc_data_structures::fx::FxHashSet;
+use rustc_middle::mir::{Body, Location, Statement, StatementKind, Terminator, TerminatorKind};
+use rustc_middle::ty::{TyCtxt, TypeVisitable};
+use rustc_mir_dataflow::points::PointIndex;
 
 use super::{LocalizedOutlivesConstraint, LocalizedOutlivesConstraintSet};
 use crate::constraints::OutlivesConstraint;
 use crate::region_infer::values::LivenessValues;
 use crate::type_check::Locations;
+use crate::universal_regions::UniversalRegions;
 
 /// Propagate loans throughout the subset graph at a given point (with some subtleties around the
 /// location where effects start to be visible).
 pub(super) fn convert_typeck_constraints<'tcx>(
+    tcx: TyCtxt<'tcx>,
     body: &Body<'tcx>,
     liveness: &LivenessValues,
     outlives_constraints: impl Iterator<Item = OutlivesConstraint<'tcx>>,
+    universal_regions: &UniversalRegions<'tcx>,
     localized_outlives_constraints: &mut LocalizedOutlivesConstraintSet,
 ) {
     for outlives_constraint in outlives_constraints {
@@ -35,7 +41,191 @@ pub(super) fn convert_typeck_constraints<'tcx>(
                 }
             }
 
-            _ => {}
+            Locations::Single(location) => {
+                // This constraint is marked as holding at one location, we localize it to that
+                // location or its successor, depending on the corresponding MIR
+                // statement/terminator. Unfortunately, they all show up from typeck as coming "on
+                // entry", so for now we modify them to take effects that should apply "on exit"
+                // into account.
+                //
+                // FIXME: this approach is subtle, complicated, and hard to test, so we should track
+                // this information better in MIR typeck instead, for example with a new `Locations`
+                // variant that contains which node is crossing over between entry and exit.
+                let point = liveness.point_from_location(location);
+                let (from, to) = if let Some(stmt) =
+                    body[location.block].statements.get(location.statement_index)
+                {
+                    localize_statement_constraint(
+                        tcx,
+                        body,
+                        stmt,
+                        liveness,
+                        &outlives_constraint,
+                        location,
+                        point,
+                        universal_regions,
+                    )
+                } else {
+                    assert_eq!(location.statement_index, body[location.block].statements.len());
+                    let terminator = body[location.block].terminator();
+                    localize_terminator_constraint(
+                        tcx,
+                        body,
+                        terminator,
+                        liveness,
+                        &outlives_constraint,
+                        point,
+                        universal_regions,
+                    )
+                };
+                localized_outlives_constraints.push(LocalizedOutlivesConstraint {
+                    source: outlives_constraint.sup,
+                    from,
+                    target: outlives_constraint.sub,
+                    to,
+                });
+            }
+        }
+    }
+}
+
+/// For a given outlives constraint arising from a MIR statement, computes the CFG `from`-`to`
+/// intra-block nodes to localize the constraint.
+fn localize_statement_constraint<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    body: &Body<'tcx>,
+    stmt: &Statement<'tcx>,
+    liveness: &LivenessValues,
+    outlives_constraint: &OutlivesConstraint<'tcx>,
+    current_location: Location,
+    current_point: PointIndex,
+    universal_regions: &UniversalRegions<'tcx>,
+) -> (PointIndex, PointIndex) {
+    match &stmt.kind {
+        StatementKind::Assign(box (lhs, rhs)) => {
+            // To create localized outlives constraints without midpoints, we rely on the property
+            // that no input regions from the RHS of the assignment will flow into themselves: they
+            // should not appear in the output regions in the LHS. We believe this to be true by
+            // construction of the MIR, via temporaries, and assert it here.
+            //
+            // We think we don't need midpoints because:
+            // - every LHS Place has a unique set of regions that don't appear elsewhere
+            // - this implies that for them to be part of the RHS, the same Place must be read and
+            //   written
+            // - and that should be impossible in MIR
+            //
+            // When we have a more complete implementation in the future, tested with crater, etc,
+            // we can relax this to a debug assert instead, or remove it.
+            assert!(
+                {
+                    let mut lhs_regions = FxHashSet::default();
+                    tcx.for_each_free_region(lhs, |region| {
+                        let region = universal_regions.to_region_vid(region);
+                        lhs_regions.insert(region);
+                    });
+
+                    let mut rhs_regions = FxHashSet::default();
+                    tcx.for_each_free_region(rhs, |region| {
+                        let region = universal_regions.to_region_vid(region);
+                        rhs_regions.insert(region);
+                    });
+
+                    // The intersection between LHS and RHS regions should be empty.
+                    lhs_regions.is_disjoint(&rhs_regions)
+                },
+                "there should be no common regions between the LHS and RHS of an assignment"
+            );
+
+            // As mentioned earlier, we should be tracking these better upstream but: we want to
+            // relate the types on entry to the type of the place on exit. That is, outlives
+            // constraints on the RHS are on entry, and outlives constraints to/from the LHS are on
+            // exit (i.e. on entry to the successor location).
+            let lhs_ty = body.local_decls[lhs.local].ty;
+            let successor_location = Location {
+                block: current_location.block,
+                statement_index: current_location.statement_index + 1,
+            };
+            let successor_point = liveness.point_from_location(successor_location);
+            compute_constraint_direction(
+                tcx,
+                outlives_constraint,
+                &lhs_ty,
+                current_point,
+                successor_point,
+                universal_regions,
+            )
+        }
+        _ => {
+            // For the other cases, we localize an outlives constraint to where it arises.
+            (current_point, current_point)
         }
     }
 }
+
+/// For a given outlives constraint arising from a MIR terminator, computes the CFG `from`-`to`
+/// inter-block nodes to localize the constraint.
+fn localize_terminator_constraint<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    body: &Body<'tcx>,
+    terminator: &Terminator<'tcx>,
+    liveness: &LivenessValues,
+    outlives_constraint: &OutlivesConstraint<'tcx>,
+    current_point: PointIndex,
+    universal_regions: &UniversalRegions<'tcx>,
+) -> (PointIndex, PointIndex) {
+    // FIXME: check if other terminators need the same handling as `Call`s, in particular
+    // Assert/Yield/Drop. A handful of tests are failing with Drop related issues, as well as some
+    // coroutine tests, and that may be why.
+    match &terminator.kind {
+        // FIXME: also handle diverging calls.
+        TerminatorKind::Call { destination, target: Some(target), .. } => {
+            // Calls are similar to assignments, and thus follow the same pattern. If there is a
+            // target for the call we also relate what flows into the destination here to entry to
+            // that successor.
+            let destination_ty = destination.ty(&body.local_decls, tcx);
+            let successor_location = Location { block: *target, statement_index: 0 };
+            let successor_point = liveness.point_from_location(successor_location);
+            compute_constraint_direction(
+                tcx,
+                outlives_constraint,
+                &destination_ty,
+                current_point,
+                successor_point,
+                universal_regions,
+            )
+        }
+        _ => {
+            // Typeck constraints guide loans between regions at the current point, so we do that in
+            // the general case, and liveness will take care of making them flow to the terminator's
+            // successors.
+            (current_point, current_point)
+        }
+    }
+}
+
+/// For a given constraint, returns the `from`-`to` edge according to whether the constraint flows
+/// to or from a free region in the given `value`, some kind of result for an effectful operation,
+/// like the LHS of an assignment.
+fn compute_constraint_direction<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    outlives_constraint: &OutlivesConstraint<'tcx>,
+    value: &impl TypeVisitable<TyCtxt<'tcx>>,
+    current_point: PointIndex,
+    successor_point: PointIndex,
+    universal_regions: &UniversalRegions<'tcx>,
+) -> (PointIndex, PointIndex) {
+    let mut to = current_point;
+    let mut from = current_point;
+    tcx.for_each_free_region(value, |region| {
+        let region = universal_regions.to_region_vid(region);
+        if region == outlives_constraint.sub {
+            // This constraint flows into the result, its effects start becoming visible on exit.
+            to = successor_point;
+        } else if region == outlives_constraint.sup {
+            // This constraint flows from the result, its effects start becoming visible on exit.
+            from = successor_point;
+        }
+    });
+
+    (from, to)
+}