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.tcx,
             &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,10 @@
+use std::rc::Rc;
+
 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;
@@ -43,12 +45,25 @@ pub(super) struct LocalizedConstraintGraph {
     /// when traversing from the node to the successor region.
     edges: FxHashMap<LocalizedNode, FxIndexSet<RegionVid>>,
 
+    location_map: Rc<DenseLocationMap>,
+
     /// The logical edges representing the outlives constraints that hold at all points in the CFG,
     /// which we don't localize to avoid creating a lot of unnecessary edges in the graph. Some CFGs
     /// can be big, and we don't need to create such a physical edge for every point in the CFG.
     logical_edges: FxHashMap<RegionVid, FxIndexSet<RegionVid>>,
 }
 
+pub(super) trait LocalizedConstraintGraphTraversal {
+    type Visitor<'a>: LocalizedConstraintGraphVisitor
+    where
+        Self: 'a;
+
+    fn mk_visitor(
+        &mut self,
+        region: RegionVid,
+    ) -> (&LivenessValues, &LiveRegionVariances, Self::Visitor<'_>);
+}
+
 /// 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
@@ -63,7 +78,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: Rc<DenseLocationMap>,
         outlives_constraints: impl Iterator<Item = OutlivesConstraint<'tcx>>,
     ) -> Self {
         let mut edges: FxHashMap<_, FxIndexSet<_>> = FxHashMap::default();
@@ -81,29 +96,26 @@ 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);
                 }
             }
         }
 
-        LocalizedConstraintGraph { edges, logical_edges }
+        LocalizedConstraintGraph { edges, logical_edges, location_map }
     }
 
     /// Traverses the localized constraint graph per-loan, and notifies the `visitor` of discovered
     /// nodes and successors.
     pub(super) fn traverse<'tcx>(
         &self,
         body: &Body<'tcx>,
-        liveness: &LivenessValues,
-        live_region_variances: &LiveRegionVariances,
         universal_regions: &UniversalRegions<'tcx>,
         borrow_set: &BorrowSet<'tcx>,
-        visitor: &mut impl LocalizedConstraintGraphVisitor,
+        traversal: &mut impl LocalizedConstraintGraphTraversal,
     ) {
-        let live_regions = liveness.points();
-
+        let location_map = &self.location_map;
         let mut visited = FxHashSet::default();
         let mut stack = Vec::new();
 
@@ -115,7 +127,7 @@ impl LocalizedConstraintGraph {
 
             let start_node = LocalizedNode {
                 region: loan.region,
-                point: liveness.point_from_location(loan.reserve_location),
+                point: location_map.point_from_location(loan.reserve_location),
             };
             stack.push(start_node);
 
@@ -124,8 +136,11 @@ impl LocalizedConstraintGraph {
                     continue;
                 }
 
+                let (liveness, live_region_variances, mut visitor) =
+                    traversal.mk_visitor(node.region);
+
                 // We've reached a node we haven't visited before.
-                let location = liveness.location_from_point(node.point);
+                let location = location_map.to_location(node.point);
                 visitor.on_node_traversed(loan_idx, node);
 
                 // When we find a _new_ successor, we'd like to
@@ -164,7 +179,7 @@ impl LocalizedConstraintGraph {
                     if let Some(succ) = compute_forward_successor(
                         node.region,
                         next_point,
-                        live_regions,
+                        liveness.points(),
                         live_region_variances,
                         is_universal_region,
                     ) {
@@ -175,11 +190,11 @@ impl LocalizedConstraintGraph {
                     // 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);
+                        let next_point = location_map.point_from_location(next_location);
                         if let Some(succ) = compute_forward_successor(
                             node.region,
                             next_point,
-                            live_regions,
+                            liveness.points(),
                             live_region_variances,
                             is_universal_region,
                         ) {
@@ -198,7 +213,7 @@ impl LocalizedConstraintGraph {
                             node.region,
                             node.point,
                             previous_point,
-                            live_regions,
+                            liveness.points(),
                             live_region_variances,
                         ) {
                             successor_found(succ);
@@ -212,12 +227,13 @@ impl LocalizedConstraintGraph {
                                 block: pred_block,
                                 statement_index: body[pred_block].statements.len(),
                             };
-                            let previous_point = liveness.point_from_location(previous_location);
+                            let previous_point =
+                                location_map.point_from_location(previous_location);
                             if let Some(succ) = compute_backward_successor(
                                 node.region,
                                 node.point,
                                 previous_point,
-                                live_regions,
+                                liveness.points(),
                                 live_region_variances,
                             ) {
                                 successor_found(succ);

compiler/rustc_borrowck/src/polonius/dump.rs

@@ -10,7 +10,10 @@ use rustc_session::config::MirIncludeSpans;
 
 use crate::borrow_set::BorrowSet;
 use crate::constraints::OutlivesConstraint;
-use crate::polonius::{LocalizedConstraintGraphVisitor, LocalizedNode, PoloniusContext};
+use crate::polonius::{
+    LiveRegionVariances, LocalizedConstraintGraphTraversal, LocalizedConstraintGraphVisitor,
+    LocalizedNode, PoloniusContext,
+};
 use crate::region_infer::values::LivenessValues;
 use crate::type_check::Locations;
 use crate::{BorrowckInferCtxt, ClosureRegionRequirements, RegionInferenceContext};
@@ -22,7 +25,7 @@ pub(crate) fn dump_polonius_mir<'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,16 +39,13 @@ 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 collector = LocalizedOutlivesConstraintCollector { constraints: Vec::new() };
+    let mut collector = LocalizedOutlivesConstraintCollectorTraversal {
+        liveness: regioncx.liveness_constraints(),
+        live_region_variances: &polonius_context.live_region_variances,
+        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 collector,
-        );
+        graph.traverse(body, regioncx.universal_regions(), borrow_set, &mut collector);
     }
 
     let extra_data = &|pass_where, out: &mut dyn io::Write| {
@@ -84,12 +84,38 @@ struct LocalizedOutlivesConstraint {
     to: PointIndex,
 }
 
-/// Visitor to record constraints encountered when traversing the localized constraint graph.
-struct LocalizedOutlivesConstraintCollector {
+struct LocalizedOutlivesConstraintCollectorTraversal<'a> {
+    liveness: &'a LivenessValues,
+    live_region_variances: &'a LiveRegionVariances,
     constraints: Vec<LocalizedOutlivesConstraint>,
 }
 
-impl LocalizedConstraintGraphVisitor for LocalizedOutlivesConstraintCollector {
+impl<'outer> LocalizedConstraintGraphTraversal
+    for LocalizedOutlivesConstraintCollectorTraversal<'outer>
+{
+    type Visitor<'a>
+        = LocalizedOutlivesConstraintCollector<'a>
+    where
+        Self: 'a;
+
+    fn mk_visitor(
+        &mut self,
+        _region: RegionVid,
+    ) -> (&LivenessValues, &LiveRegionVariances, Self::Visitor<'_>) {
+        (
+            self.liveness,
+            self.live_region_variances,
+            LocalizedOutlivesConstraintCollector { constraints: &mut self.constraints },
+        )
+    }
+}
+
+/// Visitor to record constraints encountered when traversing the localized constraint graph.
+struct LocalizedOutlivesConstraintCollector<'a> {
+    constraints: &'a mut Vec<LocalizedOutlivesConstraint>,
+}
+
+impl LocalizedConstraintGraphVisitor for LocalizedOutlivesConstraintCollector<'_> {
     fn on_successor_discovered(&mut self, current_node: LocalizedNode, successor: LocalizedNode) {
         self.constraints.push(LocalizedOutlivesConstraint {
             source: current_node.region,
@@ -246,8 +272,8 @@ fn emit_polonius_mir<'tcx>(
 
                 for constraint in localized_outlives_constraints {
                     let LocalizedOutlivesConstraint { source, from, target, to } = constraint;
-                    let from = liveness.location_from_point(*from);
-                    let to = liveness.location_from_point(*to);
+                    let from = liveness.location_map().to_location(*from);
+                    let to = liveness.location_map().to_location(*to);
                     writeln!(out, "| {source:?} at {from:?} -> {target:?} at {to:?}")?;
                 }
                 writeln!(out, "|")?;
@@ -438,7 +464,7 @@ fn emit_mermaid_constraint_graph<'tcx>(
     };
     let region_name = |region: RegionVid| format!("'{}", region.index());
     let node_name = |region: RegionVid, point: PointIndex| {
-        let location = liveness.location_from_point(point);
+        let location = liveness.location_map().to_location(point);
         format!("{}_{}", region_name(region), location_name(location))
     };
 

compiler/rustc_borrowck/src/polonius/liveness.rs

@@ -0,0 +1,66 @@
+use rustc_data_structures::fx::FxHashMap;
+use rustc_middle::mir::Local;
+use rustc_middle::ty::{GenericArg, RegionVid, Ty};
+
+use crate::BorrowckInferCtxt;
+use crate::universal_regions::UniversalRegions;
+
+#[derive(Default)]
+pub(crate) struct DeferredLocals<'tcx> {
+    /// For each region, the local whose liveness is deferred.
+    ///
+    /// Importantly, because of MIR renumbering, this will always be a 1:1 relationship.
+    by_region: FxHashMap<RegionVid, Local>,
+
+    /// For each deferred local, gets the regions contained within that local at use and drop.
+    drop_args_by_local: FxHashMap<Local, Vec<GenericArg<'tcx>>>,
+}
+
+impl<'tcx> DeferredLocals<'tcx> {
+    pub(crate) fn defer_local(
+        &mut self,
+        infcx: &BorrowckInferCtxt<'tcx>,
+        universal_regions: &UniversalRegions<'tcx>,
+        local: Local,
+        local_ty: Ty<'tcx>,
+        dropck_kinds: &[GenericArg<'tcx>],
+    ) {
+        let tcx = infcx.tcx;
+
+        // We already have drop data for this local, because we need to register
+        // region constraints eagerly. So, we'll store this so we don't need to
+        // recompute.
+        self.drop_args_by_local.insert(local, dropck_kinds.to_vec());
+
+        // Then, we want to map all the regions contained within this local to
+        // the local itself. Later, when asked for liveness of a given region,
+        // we can trace liveness for the local containing it.
+        let by_region = &mut self.by_region;
+        tcx.for_each_free_region(&local_ty, |region| {
+            // See note in [`VarianceExtractor::record_variance`].
+            if region.is_bound() || region.is_erased() {
+                return;
+            }
+            let vid = universal_regions.to_region_vid(region);
+            // Because of MIR renumbering, we should always have a 1:1 mapping
+            // between a region and a local.
+            let previous = by_region.insert(vid, local);
+            debug_assert!(
+                previous.is_none(),
+                "{vid:?} is in the type of both {previous:?} and {local:?}, but \
+                MIR renumbering should ensure that this is impossible.",
+            );
+        });
+    }
+
+    /// For a given region, return the local whose liveness is deferred, and
+    /// the regions within that local at use and drop.
+    pub(crate) fn use_deferred_local(
+        &mut self,
+        region: RegionVid,
+    ) -> Option<(Local, Vec<GenericArg<'tcx>>)> {
+        let local = self.by_region.remove(&region)?;
+        let drop_args = self.drop_args_by_local.remove(&local)?;
+        Some((local, drop_args))
+    }
+}