Auto merge of #134670 - lqd:polonius-next-episode-4, r=<try>

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`

Author: bors

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| {

compiler/rustc_borrowck/src/polonius/liveness_constraints.rs

@@ -0,0 +1,316 @@
+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.
+        let Some(region) = self.universal_regions.try_to_region_vid(region) else {
+            return;
+        };
+
+        let direction = match variance {
+            ty::Variance::Covariant => ConstraintDirection::Forward,
+            ty::Variance::Contravariant => ConstraintDirection::Backward,
+            ty::Variance::Invariant => ConstraintDirection::Bidirectional,
+            ty::Variance::Bivariant => {
+                // We don't add edges for bivariant cases.
+                return;
+            }
+        };
+
+        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>>,
+    {
+        let result = self.relate(a.skip_binder(), a.skip_binder())?;
+        Ok(a.rebind(result))
+    }
+}