Rollup of 7 pull requests

compiler/rustc_borrowck/src/nll.rs

@@ -156,7 +156,7 @@ pub(crate) fn compute_regions<'tcx>(
             &universal_region_relations.universal_regions,
             body,
             move_data,
-            Rc::clone(&location_map),
+            &location_map,
             borrow_set,
         );
     }

compiler/rustc_borrowck/src/polonius/constraints.rs

@@ -1,5 +1,3 @@
-use std::rc::Rc;
-
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_index::IndexVec;
 use rustc_middle::mir::{Body, Location};
@@ -41,8 +39,6 @@ pub(super) struct LocalizedNode {
 /// The localized constraint graph indexes the physical and logical edges to lazily compute a given
 /// node's successors during traversal.
 pub(super) struct LocalizedConstraintGraph {
-    location_map: Rc<DenseLocationMap>,
-
     /// The actual, physical, edges we have recorded for a given node. We localize them on-demand
     /// when traversing from the node to the successor region.
     edges: FxHashMap<LocalizedNode, SmallVec<[RegionVid; 4]>>,
@@ -67,7 +63,7 @@ pub(super) trait LocalizedConstraintGraphVisitor {
 impl LocalizedConstraintGraph {
     /// Traverses the constraints and returns the indexed graph of edges per node.
     pub(super) fn new<'tcx>(
-        location_map: Rc<DenseLocationMap>,
+        location_map: &DenseLocationMap,
         outlives_constraints: impl Iterator<Item = OutlivesConstraint<'tcx>>,
     ) -> Self {
         let mut edges: FxHashMap<_, SmallVec<[RegionVid; 4]>> = FxHashMap::default();
@@ -96,7 +92,7 @@ impl LocalizedConstraintGraph {
             }
         }
 
-        LocalizedConstraintGraph { location_map, edges, logical_edges }
+        LocalizedConstraintGraph { edges, logical_edges }
     }
 
     /// Traverses the localized constraint graph per-loan, and notifies the `visitor` of discovered
@@ -105,6 +101,7 @@ impl LocalizedConstraintGraph {
         &self,
         body: &Body<'tcx>,
         borrow_set: &BorrowSet<'tcx>,
+        location_map: &DenseLocationMap,
         liveness_source: &mut impl LivenessSource,
         visitor: &mut impl LocalizedConstraintGraphVisitor,
     ) {
@@ -119,15 +116,15 @@ impl LocalizedConstraintGraph {
 
             let start_node = LocalizedNode {
                 region: loan.region,
-                point: self.location_map.point_from_location(loan.reserve_location),
+                point: location_map.point_from_location(loan.reserve_location),
             };
             visited.insert(start_node);
             stack.push(start_node);
 
             while let Some(node) = stack.pop() {
                 let liveness = liveness_source.liveness_for_region(node.region);
                 // We've reached a node we haven't visited before.
-                let location = self.location_map.to_location(node.point);
+                let location = location_map.to_location(node.point);
                 visitor.on_node_traversed(loan_idx, node, liveness.is_live_at(node.point));
 
                 // When we find a _new_ successor, we'd like to
@@ -183,7 +180,7 @@ impl LocalizedConstraintGraph {
                         for successor_block in body[location.block].terminator().successors() {
                             let next_location =
                                 Location { block: successor_block, statement_index: 0 };
-                            let next_point = self.location_map.point_from_location(next_location);
+                            let next_point = location_map.point_from_location(next_location);
                             if liveness.is_live_at(next_point) {
                                 successor_found(LocalizedNode {
                                     region: node.region,
@@ -226,7 +223,7 @@ impl LocalizedConstraintGraph {
                                 statement_index: body[pred_block].statements.len(),
                             };
                             let previous_point =
-                                self.location_map.point_from_location(previous_location);
+                                location_map.point_from_location(previous_location);
                             successor_found(LocalizedNode {
                                 region: node.region,
                                 point: previous_point,

compiler/rustc_borrowck/src/polonius/dump.rs

@@ -5,7 +5,7 @@ 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;
@@ -71,7 +71,8 @@ pub(crate) fn dump_polonius_mir<'tcx>(
     };
     let mut collector = MirDumpCollector::default();
     if let Some(graph) = &polonius_context.graph {
-        graph.traverse(body, borrow_set, &mut liveness_source, &mut collector);
+        let location_map = regioncx.liveness_constraints().location_map();
+        graph.traverse(body, borrow_set, location_map, &mut liveness_source, &mut collector);
     }
 
     let extra_data = &|pass_where, out: &mut dyn io::Write| {
@@ -172,7 +173,7 @@ fn emit_polonius_dump<'tcx>(
                     "POLONIUS_CONSTRAINTS" => {
                         edge_count = emit_mermaid_constraint_graph(
                             borrow_set,
-                            regioncx.liveness_constraints(),
+                            regioncx.liveness_constraints().location_map(),
                             &collector.constraints,
                             out,
                         )?;
@@ -271,7 +272,7 @@ fn emit_polonius_mir<'tcx>(
         out,
     )?;
 
-    let liveness = regioncx.liveness_constraints();
+    let location_map = regioncx.liveness_constraints().location_map();
 
     // Add localized outlives constraints
     match pass_where {
@@ -281,8 +282,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 = location_map.to_location(*from);
+                    let to = location_map.to_location(*to);
                     writeln!(out, "| {source:?} at {from:?} -> {target:?} at {to:?}")?;
                 }
                 writeln!(out, "|")?;
@@ -462,12 +463,12 @@ fn emit_mermaid_nll_sccs<'tcx>(
 /// region, and loan introductions.
 fn emit_mermaid_constraint_graph<'tcx>(
     borrow_set: &BorrowSet<'tcx>,
-    liveness: &LivenessValues,
+    location_map: &DenseLocationMap,
     localized_outlives_constraints: &[LocalizedOutlivesConstraint],
     out: &mut dyn io::Write,
 ) -> io::Result<usize> {
     let node_label = |region: RegionVid, point: PointIndex| {
-        let location = liveness.location_from_point(point);
+        let location = location_map.to_location(point);
         node_name(region, location)
     };
 
@@ -533,6 +534,7 @@ fn emit_loan_reachability(
     reachability: &FxIndexMap<BorrowIndex, Vec<LocalizedNode>>,
     out: &mut dyn io::Write,
 ) -> io::Result<()> {
+    let location_map = liveness.location_map();
     for (loan, _) in borrow_set.iter_enumerated() {
         let Some(reachability) = reachability.get(&loan) else {
             continue;
@@ -553,7 +555,7 @@ fn emit_loan_reachability(
         for (idx, node) in reachability.iter().enumerate() {
             writeln!(out, "<li>")?;
 
-            let location = liveness.location_from_point(node.point);
+            let location = location_map.to_location(node.point);
             let kind = if idx == 0 { "starts in" } else { "reaches" };
             writeln!(
                 out,

compiler/rustc_borrowck/src/polonius/liveness.rs

@@ -57,9 +57,11 @@ impl<'a> RegionLiveness<'a> {
 }
 
 /// The data needed to compute region liveness on-demand while traversing the localized outlives
-/// constraint graph to compute loan liveness.
+/// constraint graph to compute loan liveness. Each region comes from a local's type, so this holds
+/// mostly the locals and drop-liveness specifics about their type, to be used as inputs to the
+/// `LivenessComputation` giving us the region liveness.
 #[derive(Default)]
-pub(crate) struct DeferredLocals<'tcx> {
+pub(crate) struct DeferredRegionLiveness<'tcx> {
     /// For each region, the local whose liveness is deferred.
     ///
     /// Importantly, because of MIR renumbering, this will always be a 1:1 relationship.
@@ -69,7 +71,9 @@ pub(crate) struct DeferredLocals<'tcx> {
     drop_args_by_local: IndexVec<Local, Option<Vec<GenericArg<'tcx>>>>,
 }
 
-impl<'tcx> DeferredLocals<'tcx> {
+impl<'tcx> DeferredRegionLiveness<'tcx> {
+    /// We want to defer computing the liveness of the region in the local's type, so we store the
+    /// necessary inputs to do that later.
     pub(crate) fn defer_local(
         &mut self,
         tcx: TyCtxt<'tcx>,
@@ -86,7 +90,6 @@ impl<'tcx> DeferredLocals<'tcx> {
         // 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() {
@@ -95,7 +98,7 @@ impl<'tcx> DeferredLocals<'tcx> {
             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);
+            let previous = self.by_region.insert(vid, local);
             debug_assert!(
                 previous.is_none(),
                 "{vid:?} is in the type of both {previous:?} and {local:?}, but \
@@ -106,7 +109,7 @@ impl<'tcx> DeferredLocals<'tcx> {
 
     /// For a given region, compute the liveness for the local containing it, if it is deferred.
     #[inline]
-    pub(crate) fn compute_deferred_local(
+    pub(crate) fn ensure_deferred_liveness(
         &mut self,
         region: RegionVid,
         universal_regions: &UniversalRegions<'tcx>,

compiler/rustc_borrowck/src/polonius/mod.rs

@@ -39,8 +39,6 @@ pub(crate) mod legacy;
 mod liveness;
 mod liveness_constraints;
 
-use std::rc::Rc;
-
 use rustc_data_structures::fx::FxHashSet;
 use rustc_index::IndexVec;
 use rustc_index::bit_set::DenseBitSet;
@@ -54,7 +52,7 @@ pub(crate) use self::dump::dump_polonius_mir;
 pub(crate) use self::liveness_constraints::record_live_region_variance;
 use crate::constraints::OutlivesConstraint;
 use crate::dataflow::BorrowIndex;
-pub(crate) use crate::polonius::liveness::DeferredLocals;
+pub(crate) use crate::polonius::liveness::DeferredRegionLiveness;
 use crate::polonius::liveness::{LivenessSource, RegionLiveness};
 use crate::region_infer::values::LivenessValues;
 use crate::type_check::liveness::{LivenessComputation, LocalUseMap};
@@ -105,7 +103,7 @@ pub(crate) struct PoloniusContext<'tcx> {
     /// diagnostics, to focus on the locals we consider relevant and match NLL diagnostics.
     pub(crate) boring_nll_locals: FxHashSet<Local>,
 
-    pub(crate) deferred_locals_for_liveness: DeferredLocals<'tcx>,
+    pub(crate) deferred_liveness: DeferredRegionLiveness<'tcx>,
 
     pub(crate) local_use_map: Option<LocalUseMap>,
 }
@@ -143,7 +141,7 @@ impl<'tcx> PoloniusContext<'tcx> {
         universal_regions: &UniversalRegions<'tcx>,
         body: &Body<'tcx>,
         move_data: &MoveData<'tcx>,
-        location_map: Rc<DenseLocationMap>,
+        location_map: &DenseLocationMap,
         borrow_set: &BorrowSet<'tcx>,
     ) {
         // We don't need to prepare the graph (index NLL constraints, etc.) if we have no loans to
@@ -152,27 +150,25 @@ impl<'tcx> PoloniusContext<'tcx> {
             // From the outlives constraints, liveness, and variances, we can compute reachability
             // on the lazy localized constraint graph to trace the liveness of loans, for the next
             // step in the chain (the NLL loan scope and active loans computations).
-            let graph =
-                LocalizedConstraintGraph::new(Rc::clone(&location_map), outlives_constraints);
+            let graph = LocalizedConstraintGraph::new(location_map, outlives_constraints);
 
             let local_use_map = self
                 .local_use_map
                 .as_ref()
                 .expect("local use map should be computed before loan liveness");
-            let deferred_locals_for_liveness =
-                std::mem::take(&mut self.deferred_locals_for_liveness);
+            let deferred_liveness = std::mem::take(&mut self.deferred_liveness);
             let mut live_loans = LiveLoans::new(location_map.num_points(), borrow_set.len());
             let comp =
-                LivenessComputation::new(infcx, body, &location_map, move_data, &local_use_map);
+                LivenessComputation::new(infcx, body, location_map, move_data, &local_use_map);
             let mut liveness_source = DeferredLivenessSource {
                 liveness,
                 live_region_variances: &mut self.live_region_variances,
                 universal_regions,
-                deferred_locals_for_liveness,
+                deferred_liveness,
                 comp,
             };
             let mut visitor = LoanLivenessVisitor { live_loans: &mut live_loans };
-            graph.traverse(body, borrow_set, &mut liveness_source, &mut visitor);
+            graph.traverse(body, borrow_set, location_map, &mut liveness_source, &mut visitor);
             liveness.record_live_loans(live_loans);
 
             // The graph can be traversed again during MIR dumping, so we store it here.
@@ -188,14 +184,14 @@ struct DeferredLivenessSource<'a, 'tcx> {
     liveness: &'a mut LivenessValues,
     live_region_variances: &'a mut LiveRegionVariances,
     universal_regions: &'a UniversalRegions<'tcx>,
-    deferred_locals_for_liveness: DeferredLocals<'tcx>,
+    deferred_liveness: DeferredRegionLiveness<'tcx>,
     comp: LivenessComputation<'a, 'tcx>,
 }
 
 impl<'a> LivenessSource for DeferredLivenessSource<'a, '_> {
     #[inline]
     fn liveness_for_region(&mut self, region: RegionVid) -> RegionLiveness<'_> {
-        self.deferred_locals_for_liveness.compute_deferred_local(
+        self.deferred_liveness.ensure_deferred_liveness(
             region,
             self.universal_regions,
             &mut self.liveness,

compiler/rustc_borrowck/src/region_infer/region_context.rs

@@ -694,7 +694,7 @@ impl<'tcx> RegionInferenceContextInner<'tcx> {
     /// region is contained within the type of a variable that is live at this point.
     /// Note: for now, the sets of live loans is only available when using `-Zpolonius=next`.
     pub(crate) fn is_loan_live_at(&self, loan_idx: BorrowIndex, location: Location) -> bool {
-        let point = self.liveness_constraints.point_from_location(location);
+        let point = self.liveness_constraints.location_map().point_from_location(location);
         self.liveness_constraints.is_loan_live_at(loan_idx, point)
     }
 }

compiler/rustc_borrowck/src/region_infer/values.rs

@@ -191,16 +191,6 @@ impl LivenessValues {
         )
     }
 
-    #[inline]
-    pub(crate) fn point_from_location(&self, location: Location) -> PointIndex {
-        self.location_map.point_from_location(location)
-    }
-
-    #[inline]
-    pub(crate) fn location_from_point(&self, point: PointIndex) -> Location {
-        self.location_map.to_location(point)
-    }
-
     /// When using `-Zpolonius=next`, records the given live loans for the loan scopes and active
     /// loans dataflow computations.
     pub(crate) fn record_live_loans(&mut self, live_loans: LiveLoans) {
@@ -214,6 +204,11 @@ impl LivenessValues {
             .expect("Accessing live loans requires `-Zpolonius=next`")
             .contains(point, loan_idx)
     }
+
+    #[inline]
+    pub(crate) fn location_map(&self) -> &DenseLocationMap {
+        &self.location_map
+    }
 }
 
 /// Maps from `ty::PlaceholderRegion` values that are used in the rest of

compiler/rustc_borrowck/src/type_check/liveness/mod.rs

@@ -69,24 +69,42 @@ pub(super) fn generate<'tcx>(
     // for now.
     // FIXME: this NLL optimization idea, to reduce work to relevant locals only, still makes sense
     // for polonius, and should be investigated to improve liveness performance.
-    let deferred_locals = 'deferred: {
+    let deferred_locals = if typeck.polonius_context.is_none() || typeck.borrow_set.len() == 0 {
         // If we aren't going to be using the additional liveness information,
         // don't even bother computing the larger relevant set.
         // Similarly, since this liveness information is ultimately used for *loan*
         // liveness, we don't need to compute it when there are no loans.
-        if typeck.polonius_context.is_none() || typeck.borrow_set.len() == 0 {
-            break 'deferred FxIndexSet::default();
-        }
-
-        let free_regions = typeck.universal_regions.universal_regions_iter().collect();
-        let (polonius_relevant, _) =
-            compute_relevant_live_locals(typeck.tcx(), &free_regions, typeck.body);
+        FxIndexSet::default()
+    } else {
+        // As described above, we can defer computing liveness for the regions that are
+        // NLL-boring-and-polonius-relevant. Let's find these.
+        //
+        // Note that this is basically a simplified version of `compute_relevant_live_locals`
+        // avoiding the allocations that are unnecessary in our more limited use-case.
+        let tcx = typeck.tcx();
+        let universal_regions = typeck.universal_regions;
+        let boring_nll_locals: FxHashSet<_> = boring_locals.iter().copied().collect();
+        let deferred_polonius_locals = typeck
+            .body
+            .local_decls
+            .iter_enumerated()
+            .filter_map(|(local, decl)| {
+                // The polonius-relevant locals are the ones whose types do not only contain
+                // universal regions.
+                if boring_nll_locals.contains(&local)
+                    && !tcx.all_free_regions_meet(&decl.ty, |r| {
+                        universal_regions.is_universal_region(r.as_var())
+                    })
+                {
+                    Some(local)
+                } else {
+                    None
+                }
+            })
+            .collect();
 
-        let boring: FxHashSet<_> = boring_locals.iter().copied().collect();
-        let deferred =
-            polonius_relevant.into_iter().filter(|local| boring.contains(local)).collect();
-        typeck.polonius_context.as_mut().unwrap().boring_nll_locals = boring;
-        deferred
+        typeck.polonius_context.as_mut().unwrap().boring_nll_locals = boring_nll_locals;
+        deferred_polonius_locals
     };
 
     trace::trace(