Defer extra liveness calculation for Polonius Alpha

compiler/rustc_borrowck/src/lib.rs

@@ -311,7 +311,7 @@ struct CollectRegionConstraintsResult<'tcx> {
     deferred_closure_requirements: DeferredClosureRequirements<'tcx>,
     deferred_opaque_type_errors: Vec<DeferredOpaqueTypeError<'tcx>>,
     polonius_facts: Option<AllFacts<RustcFacts>>,
-    polonius_context: Option<PoloniusContext>,
+    polonius_context: Option<PoloniusContext<'tcx>>,
 }
 
 /// Start borrow checking by collecting the region constraints for
@@ -799,7 +799,7 @@ pub(crate) struct MirBorrowckCtxt<'a, 'diag, 'tcx> {
     /// Results of Polonius analysis.
     polonius_output: Option<&'a PoloniusOutput>,
     /// When using `-Zpolonius=next`: the data used to compute errors and diagnostics.
-    polonius_context: Option<&'a PoloniusContext>,
+    polonius_context: Option<&'a PoloniusContext<'tcx>>,
 }
 
 // Check that:

compiler/rustc_borrowck/src/nll.rs

@@ -46,7 +46,7 @@ pub(crate) struct NllOutput<'tcx> {
 
     /// When using `-Zpolonius=next`: the data used to compute errors and diagnostics, e.g.
     /// localized typeck and liveness constraints.
-    pub polonius_context: Option<PoloniusContext>,
+    pub polonius_context: Option<PoloniusContext<'tcx>>,
 }
 
 /// Rewrites the regions in the MIR to use NLL variables, also scraping out the set of universal
@@ -121,7 +121,7 @@ pub(crate) fn compute_regions<'tcx>(
     universal_region_relations: Frozen<UniversalRegionRelations<'tcx>>,
     constraints: MirTypeckRegionConstraints<'tcx>,
     mut polonius_facts: Option<AllFacts<RustcFacts>>,
-    mut polonius_context: Option<PoloniusContext>,
+    mut polonius_context: Option<PoloniusContext<'tcx>>,
 ) -> NllOutput<'tcx> {
     let polonius_output = root_cx.consumer.as_ref().map_or(false, |c| c.polonius_output())
         || infcx.tcx.sess.opts.unstable_opts.polonius.is_legacy_enabled();
@@ -144,20 +144,20 @@ pub(crate) fn compute_regions<'tcx>(
         &lowered_constraints,
     );
 
-    let num_points = location_map.num_points();
-
     // If requested for `-Zpolonius=next`, compute loan liveness information.
     // This is done prior to `RegionInferenceContext::new`, because we may add
     // additional liveness constraints.
     if let Some(polonius_context) = polonius_context.as_mut() {
         let _timer = infcx.tcx.prof.generic_activity("borrowck_polonius_loan_liveness");
         polonius_context.compute_loan_liveness(
+            infcx,
             &mut lowered_constraints.liveness_constraints,
             lowered_constraints.outlives_constraints.outlives().iter().copied(),
             &universal_region_relations.universal_regions,
             body,
+            move_data,
+            &location_map,
             borrow_set,
-            num_points,
         );
     }
 

compiler/rustc_borrowck/src/polonius/constraints.rs

@@ -1,8 +1,7 @@
 use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexSet};
-use rustc_index::interval::SparseIntervalMatrix;
 use rustc_middle::mir::{Body, Location};
 use rustc_middle::ty::RegionVid;
-use rustc_mir_dataflow::points::PointIndex;
+use rustc_mir_dataflow::points::{DenseLocationMap, PointIndex};
 
 use crate::BorrowSet;
 use crate::constraints::OutlivesConstraint;
@@ -49,11 +48,50 @@ pub(super) struct LocalizedConstraintGraph {
     logical_edges: FxHashMap<RegionVid, FxIndexSet<RegionVid>>,
 }
 
+/// For a given region, the relevant liveness and variance information.
+pub(super) struct RegionLiveness<'a> {
+    region: RegionVid,
+    direction: ConstraintDirection,
+    liveness: &'a LivenessValues,
+}
+
+impl<'a> RegionLiveness<'a> {
+    pub(super) fn new(
+        region: RegionVid,
+        live_region_variances: &LiveRegionVariances,
+        liveness: &'a LivenessValues,
+    ) -> Self {
+        // 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 direction = live_region_variances
+            .get(region)
+            .copied()
+            .flatten()
+            .unwrap_or(ConstraintDirection::Bidirectional);
+        Self { region, direction, liveness }
+    }
+
+    fn is_live_at(&self, point: PointIndex) -> bool {
+        self.liveness.points().contains(self.region, point)
+    }
+}
+
+/// The source of liveness information for a given region.
+pub(super) trait LivenessSource {
+    fn liveness_for_region(&mut self, region: RegionVid) -> RegionLiveness<'_>;
+    fn location_map(&self) -> &DenseLocationMap;
+}
+
 /// The visitor interface when traversing a `LocalizedConstraintGraph`.
 pub(super) trait LocalizedConstraintGraphVisitor {
     /// Callback called when traversing a given `loan` encounters a localized `node` it hasn't
     /// visited before.
-    fn on_node_traversed(&mut self, _loan: BorrowIndex, _node: LocalizedNode) {}
+    fn on_live_node_traversed(&mut self, _loan: BorrowIndex, _node: LocalizedNode) {}
 
     /// Callback called when discovering a new `successor` node for the `current_node`.
     fn on_successor_discovered(&mut self, _current_node: LocalizedNode, _successor: LocalizedNode) {
@@ -63,7 +101,7 @@ pub(super) trait LocalizedConstraintGraphVisitor {
 impl LocalizedConstraintGraph {
     /// Traverses the constraints and returns the indexed graph of edges per node.
     pub(super) fn new<'tcx>(
-        liveness: &LivenessValues,
+        location_map: &DenseLocationMap,
         outlives_constraints: impl Iterator<Item = OutlivesConstraint<'tcx>>,
     ) -> Self {
         let mut edges: FxHashMap<_, FxIndexSet<_>> = FxHashMap::default();
@@ -81,7 +119,7 @@ impl LocalizedConstraintGraph {
                 Locations::Single(location) => {
                     let node = LocalizedNode {
                         region: outlives_constraint.sup,
-                        point: liveness.point_from_location(location),
+                        point: location_map.point_from_location(location),
                     };
                     edges.entry(node).or_default().insert(outlives_constraint.sub);
                 }
@@ -96,14 +134,11 @@ impl LocalizedConstraintGraph {
     pub(super) fn traverse<'tcx>(
         &self,
         body: &Body<'tcx>,
-        liveness: &LivenessValues,
-        live_region_variances: &LiveRegionVariances,
         universal_regions: &UniversalRegions<'tcx>,
         borrow_set: &BorrowSet<'tcx>,
+        liveness_source: &mut impl LivenessSource,
         visitor: &mut impl LocalizedConstraintGraphVisitor,
     ) {
-        let live_regions = liveness.points();
-
         let mut visited = FxHashSet::default();
         let mut stack = Vec::new();
 
@@ -115,7 +150,7 @@ impl LocalizedConstraintGraph {
 
             let start_node = LocalizedNode {
                 region: loan.region,
-                point: liveness.point_from_location(loan.reserve_location),
+                point: liveness_source.location_map().point_from_location(loan.reserve_location),
             };
             stack.push(start_node);
 
@@ -124,9 +159,12 @@ impl LocalizedConstraintGraph {
                     continue;
                 }
 
+                let liveness = liveness_source.liveness_for_region(node.region);
                 // We've reached a node we haven't visited before.
-                let location = liveness.location_from_point(node.point);
-                visitor.on_node_traversed(loan_idx, node);
+                let location = liveness.liveness.location_map().to_location(node.point);
+                if liveness.is_live_at(node.point) {
+                    visitor.on_live_node_traversed(loan_idx, node);
+                }
 
                 // When we find a _new_ successor, we'd like to
                 // - visit it eventually,
@@ -161,28 +199,20 @@ impl LocalizedConstraintGraph {
                     // Intra-block edges, straight line constraints from each point to its successor
                     // within the same block.
                     let next_point = node.point + 1;
-                    if let Some(succ) = compute_forward_successor(
-                        node.region,
-                        next_point,
-                        live_regions,
-                        live_region_variances,
-                        is_universal_region,
-                    ) {
+                    if let Some(succ) =
+                        compute_forward_successor(&liveness, next_point, is_universal_region)
+                    {
                         successor_found(succ);
                     }
                 } else {
                     // Inter-block edges, from the block's terminator to each successor block's
                     // entry point.
                     for successor_block in body[location.block].terminator().successors() {
                         let next_location = Location { block: successor_block, statement_index: 0 };
-                        let next_point = liveness.point_from_location(next_location);
-                        if let Some(succ) = compute_forward_successor(
-                            node.region,
-                            next_point,
-                            live_regions,
-                            live_region_variances,
-                            is_universal_region,
-                        ) {
+                        let next_point = liveness.liveness.point_from_location(next_location);
+                        if let Some(succ) =
+                            compute_forward_successor(&liveness, next_point, is_universal_region)
+                        {
                             successor_found(succ);
                         }
                     }
@@ -194,13 +224,9 @@ impl LocalizedConstraintGraph {
                     if location.statement_index > 0 {
                         // Backward edges to the predecessor point in the same block.
                         let previous_point = PointIndex::from(node.point.as_usize() - 1);
-                        if let Some(succ) = compute_backward_successor(
-                            node.region,
-                            node.point,
-                            previous_point,
-                            live_regions,
-                            live_region_variances,
-                        ) {
+                        if let Some(succ) =
+                            compute_backward_successor(&liveness, node.point, previous_point)
+                        {
                             successor_found(succ);
                         }
                     } else {
@@ -212,14 +238,11 @@ impl LocalizedConstraintGraph {
                                 block: pred_block,
                                 statement_index: body[pred_block].statements.len(),
                             };
-                            let previous_point = liveness.point_from_location(previous_location);
-                            if let Some(succ) = compute_backward_successor(
-                                node.region,
-                                node.point,
-                                previous_point,
-                                live_regions,
-                                live_region_variances,
-                            ) {
+                            let previous_point =
+                                liveness.liveness.point_from_location(previous_location);
+                            if let Some(succ) =
+                                compute_backward_successor(&liveness, node.point, previous_point)
+                            {
                                 successor_found(succ);
                             }
                         }
@@ -239,40 +262,27 @@ impl LocalizedConstraintGraph {
 /// Returns the successor for the current region/point node when propagating a loan through forward
 /// edges, if applicable, according to liveness and variance.
 fn compute_forward_successor(
-    region: RegionVid,
+    liveness: &RegionLiveness<'_>,
     next_point: PointIndex,
-    live_regions: &SparseIntervalMatrix<RegionVid, PointIndex>,
-    live_region_variances: &LiveRegionVariances,
     is_universal_region: bool,
 ) -> Option<LocalizedNode> {
+    let region = liveness.region;
+
     // 1. Universal regions are semantically live at all points.
     if is_universal_region {
         let succ = LocalizedNode { region, point: next_point };
         return Some(succ);
     }
 
     // 2. Otherwise, gather the edges due to explicit region liveness, when applicable.
-    if !live_regions.contains(region, next_point) {
+    if !liveness.is_live_at(next_point) {
         return None;
     }
 
     // Here, `region` could be live at the current point, and is live at the next point: add a
     // constraint between them, according to variance.
 
-    // 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 direction = live_region_variances
-        .get(region)
-        .copied()
-        .flatten()
-        .unwrap_or(ConstraintDirection::Bidirectional);
-
-    match direction {
+    match liveness.direction {
         ConstraintDirection::Backward => {
             // Contravariant cases: loans flow in the inverse direction, but we're only interested
             // in forward successors and there are none here.
@@ -291,27 +301,19 @@ fn compute_forward_successor(
 /// Returns the successor for the current region/point node when propagating a loan through backward
 /// edges, if applicable, according to liveness and variance.
 fn compute_backward_successor(
-    region: RegionVid,
+    liveness: &RegionLiveness<'_>,
     current_point: PointIndex,
     previous_point: PointIndex,
-    live_regions: &SparseIntervalMatrix<RegionVid, PointIndex>,
-    live_region_variances: &LiveRegionVariances,
 ) -> Option<LocalizedNode> {
+    let region = liveness.region;
+
     // Liveness flows into the regions live at the next point. So, in a backwards view, we'll link
     // the region from the current point, if it's live there, to the previous point.
-    if !live_regions.contains(region, current_point) {
+    if !liveness.is_live_at(current_point) {
         return None;
     }
 
-    // FIXME: add the missing variance information and remove this fallback bidirectional edge. See
-    // the same comment in `compute_forward_successor`.
-    let direction = live_region_variances
-        .get(region)
-        .copied()
-        .flatten()
-        .unwrap_or(ConstraintDirection::Bidirectional);
-
-    match direction {
+    match liveness.direction {
         ConstraintDirection::Forward => {
             // Covariant cases: loans flow in the regular direction, but we're only interested in
             // backward successors and there are none here.

compiler/rustc_borrowck/src/polonius/dump.rs

@@ -5,24 +5,42 @@ use rustc_index::IndexVec;
 use rustc_middle::mir::pretty::{MirDumper, PassWhere, PrettyPrintMirOptions};
 use rustc_middle::mir::{Body, Location};
 use rustc_middle::ty::{RegionVid, TyCtxt};
-use rustc_mir_dataflow::points::PointIndex;
+use rustc_mir_dataflow::points::{DenseLocationMap, PointIndex};
 use rustc_session::config::MirIncludeSpans;
 
 use crate::borrow_set::BorrowSet;
 use crate::constraints::OutlivesConstraint;
-use crate::polonius::{LocalizedConstraintGraphVisitor, LocalizedNode, PoloniusContext};
+use crate::polonius::{
+    LiveRegionVariances, LivenessSource, LocalizedConstraintGraphVisitor, LocalizedNode,
+    PoloniusContext, RegionLiveness,
+};
 use crate::region_infer::values::LivenessValues;
 use crate::type_check::Locations;
 use crate::{BorrowckInferCtxt, ClosureRegionRequirements, RegionInferenceContext};
 
+/// A `LivenessSource` for already-existing liveness and variance data.
+struct CachedLivenessSource<'a> {
+    live_region_variances: &'a LiveRegionVariances,
+    liveness: &'a LivenessValues,
+}
+
+impl<'a> LivenessSource for CachedLivenessSource<'a> {
+    fn liveness_for_region(&mut self, region: RegionVid) -> RegionLiveness<'_> {
+        RegionLiveness::new(region, self.live_region_variances, self.liveness)
+    }
+    fn location_map(&self) -> &DenseLocationMap {
+        self.liveness.location_map()
+    }
+}
+
 /// `-Zdump-mir=polonius` dumps MIR annotated with NLL and polonius specific information.
 pub(crate) fn dump_polonius_mir<'tcx>(
     infcx: &BorrowckInferCtxt<'tcx>,
     body: &Body<'tcx>,
     regioncx: &RegionInferenceContext<'tcx>,
     closure_region_requirements: &Option<ClosureRegionRequirements<'tcx>>,
     borrow_set: &BorrowSet<'tcx>,
-    polonius_context: Option<&PoloniusContext>,
+    polonius_context: Option<&PoloniusContext<'tcx>>,
 ) {
     let tcx = infcx.tcx;
     if !tcx.sess.opts.unstable_opts.polonius.is_next_enabled() {
@@ -36,14 +54,17 @@ pub(crate) fn dump_polonius_mir<'tcx>(
 
     // If we have a polonius graph to dump along the rest of the MIR and NLL info, we extract its
     // constraints here.
+    let mut liveness_source = CachedLivenessSource {
+        live_region_variances: &polonius_context.live_region_variances,
+        liveness: regioncx.liveness_constraints(),
+    };
     let mut collector = LocalizedOutlivesConstraintCollector { constraints: Vec::new() };
     if let Some(graph) = &polonius_context.graph {
         graph.traverse(
             body,
-            regioncx.liveness_constraints(),
-            &polonius_context.live_region_variances,
             regioncx.universal_regions(),
             borrow_set,
+            &mut liveness_source,
             &mut collector,
         );
     }