Auto merge of rust-lang#118175 - lqd:unify-live-loans, r=matthewjasper

Centralize live loans maintenance to fix scope differences due to liveness

As found in the recent [polonius crater run](rust-lang#117593 (comment)), NLLs and the location-insensitive polonius computed different scopes on some specific CFG shapes, e.g. the following.

![image](https://github.com/rust-lang/rust/assets/247183/c3649f5e-3058-454e-854e-1a6b336bdd5e)

I had missed that liveness data was pushed from different sources than just the liveness computation: there are a few places that do this -- and some of them may be unneeded or at the very least untested, as no tests changed when I tried removing some of them.

Here, `_6` is e.g. dead on entry to `bb2[0]` during `liveness::trace`, but its regions will be marked as live later during "constraint generation" (which I plan to refactor away and put in the liveness module soon). This should cause the inflowing loans to be marked live, but they were only computed in `liveness::trace`.

Therefore, this PR moves live loan maintenance to `LivenessValues`, so that the various places pushing liveness data will all also update live loans at the same time -- except for promoteds which I don't believe need them, and their liveness handling is already interesting/peculiar.

All the regressions I saw in the initial crater run were related to this kind of shapes, and this change did fix all of them on the [next run](rust-lang#117593 (comment)).

r? `@matthewjasper`

(This will conflict with rust-lang#117880 but whichever lands first is fine by me, the end goal is the same for both)

Author: bors

compiler/rustc_borrowck/src/dataflow.rs

@@ -432,7 +432,8 @@ impl<'mir, 'tcx> Borrows<'mir, 'tcx> {
 
             assert_eq!(
                 borrows_out_of_scope_at_location, polonius_prec.loans_out_of_scope_at_location,
-                "the loans out of scope must be the same as the borrows out of scope"
+                "polonius loan scopes differ from NLL borrow scopes, for body {:?}",
+                body.span,
             );
 
             borrows_out_of_scope_at_location = polonius_prec.loans_out_of_scope_at_location;

compiler/rustc_borrowck/src/nll.rs

@@ -101,26 +101,22 @@ pub(crate) fn compute_regions<'cx, 'tcx>(
     let elements = &Rc::new(RegionValueElements::new(body));
 
     // Run the MIR type-checker.
-    let MirTypeckResults {
-        constraints,
-        universal_region_relations,
-        opaque_type_values,
-        live_loans,
-    } = type_check::type_check(
-        infcx,
-        param_env,
-        body,
-        promoted,
-        &universal_regions,
-        location_table,
-        borrow_set,
-        &mut all_facts,
-        flow_inits,
-        move_data,
-        elements,
-        upvars,
-        polonius_input,
-    );
+    let MirTypeckResults { constraints, universal_region_relations, opaque_type_values } =
+        type_check::type_check(
+            infcx,
+            param_env,
+            body,
+            promoted,
+            &universal_regions,
+            location_table,
+            borrow_set,
+            &mut all_facts,
+            flow_inits,
+            move_data,
+            elements,
+            upvars,
+            polonius_input,
+        );
 
     // Create the region inference context, taking ownership of the
     // region inference data that was contained in `infcx`, and the
@@ -161,7 +157,6 @@ pub(crate) fn compute_regions<'cx, 'tcx>(
         type_tests,
         liveness_constraints,
         elements,
-        live_loans,
     );
 
     // If requested: dump NLL facts, and run legacy polonius analysis.

compiler/rustc_borrowck/src/region_infer/mod.rs

@@ -7,7 +7,6 @@ use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
 use rustc_data_structures::graph::scc::Sccs;
 use rustc_errors::Diagnostic;
 use rustc_hir::def_id::CRATE_DEF_ID;
-use rustc_index::bit_set::SparseBitMatrix;
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_infer::infer::outlives::test_type_match;
 use rustc_infer::infer::region_constraints::{GenericKind, VarInfos, VerifyBound, VerifyIfEq};
@@ -31,8 +30,8 @@ use crate::{
     nll::PoloniusOutput,
     region_infer::reverse_sccs::ReverseSccGraph,
     region_infer::values::{
-        LivenessValues, PlaceholderIndices, PointIndex, RegionElement, RegionValueElements,
-        RegionValues, ToElementIndex,
+        LivenessValues, PlaceholderIndices, RegionElement, RegionValueElements, RegionValues,
+        ToElementIndex,
     },
     type_check::{free_region_relations::UniversalRegionRelations, Locations},
     universal_regions::UniversalRegions,
@@ -120,9 +119,6 @@ pub struct RegionInferenceContext<'tcx> {
     /// Information about how the universally quantified regions in
     /// scope on this function relate to one another.
     universal_region_relations: Frozen<UniversalRegionRelations<'tcx>>,
-
-    /// The set of loans that are live at a given point in the CFG, when using `-Zpolonius=next`.
-    live_loans: SparseBitMatrix<PointIndex, BorrowIndex>,
 }
 
 /// Each time that `apply_member_constraint` is successful, it appends
@@ -335,7 +331,6 @@ impl<'tcx> RegionInferenceContext<'tcx> {
         type_tests: Vec<TypeTest<'tcx>>,
         liveness_constraints: LivenessValues,
         elements: &Rc<RegionValueElements>,
-        live_loans: SparseBitMatrix<PointIndex, BorrowIndex>,
     ) -> Self {
         debug!("universal_regions: {:#?}", universal_regions);
         debug!("outlives constraints: {:#?}", outlives_constraints);
@@ -389,7 +384,6 @@ impl<'tcx> RegionInferenceContext<'tcx> {
             type_tests,
             universal_regions,
             universal_region_relations,
-            live_loans,
         };
 
         result.init_free_and_bound_regions();
@@ -2325,7 +2319,7 @@ impl<'tcx> RegionInferenceContext<'tcx> {
     /// 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);
-        self.live_loans.contains(point, loan_idx)
+        self.liveness_constraints.is_loan_live_at(loan_idx, point)
     }
 }
 

compiler/rustc_borrowck/src/region_infer/values.rs

@@ -11,6 +11,8 @@ use rustc_middle::ty::{self, RegionVid};
 use std::fmt::Debug;
 use std::rc::Rc;
 
+use crate::dataflow::BorrowIndex;
+
 /// Maps between a `Location` and a `PointIndex` (and vice versa).
 pub(crate) struct RegionValueElements {
     /// For each basic block, how many points are contained within?
@@ -120,14 +122,45 @@ pub(crate) enum RegionElement {
 /// Records the CFG locations where each region is live. When we initially compute liveness, we use
 /// an interval matrix storing liveness ranges for each region-vid.
 pub(crate) struct LivenessValues {
+    /// The map from locations to points.
     elements: Rc<RegionValueElements>,
+
+    /// For each region: the points where it is live.
     points: SparseIntervalMatrix<RegionVid, PointIndex>,
+
+    /// When using `-Zpolonius=next`, for each point: the loans flowing into the live regions at
+    /// that point.
+    pub(crate) loans: Option<LiveLoans>,
+}
+
+/// Data used to compute the loans that are live at a given point in the CFG, when using
+/// `-Zpolonius=next`.
+pub(crate) struct LiveLoans {
+    /// The set of loans that flow into a given region. When individual regions are marked as live
+    /// in the CFG, these inflowing loans are recorded as live.
+    pub(crate) inflowing_loans: SparseBitMatrix<RegionVid, BorrowIndex>,
+
+    /// The set of loans that are live at a given point in the CFG.
+    pub(crate) live_loans: SparseBitMatrix<PointIndex, BorrowIndex>,
+}
+
+impl LiveLoans {
+    pub(crate) fn new(num_loans: usize) -> Self {
+        LiveLoans {
+            live_loans: SparseBitMatrix::new(num_loans),
+            inflowing_loans: SparseBitMatrix::new(num_loans),
+        }
+    }
 }
 
 impl LivenessValues {
     /// Create an empty map of regions to locations where they're live.
     pub(crate) fn new(elements: Rc<RegionValueElements>) -> Self {
-        Self { points: SparseIntervalMatrix::new(elements.num_points), elements }
+        LivenessValues {
+            points: SparseIntervalMatrix::new(elements.num_points),
+            elements,
+            loans: None,
+        }
     }
 
     /// Iterate through each region that has a value in this set.
@@ -140,12 +173,30 @@ impl LivenessValues {
         debug!("LivenessValues::add_location(region={:?}, location={:?})", region, location);
         let point = self.elements.point_from_location(location);
         self.points.insert(region, point);
+
+        // When available, record the loans flowing into this region as live at the given point.
+        if let Some(loans) = self.loans.as_mut() {
+            if let Some(inflowing) = loans.inflowing_loans.row(region) {
+                loans.live_loans.union_row(point, inflowing);
+            }
+        }
     }
 
     /// Records `region` as being live at all the given `points`.
     pub(crate) fn add_points(&mut self, region: RegionVid, points: &IntervalSet<PointIndex>) {
         debug!("LivenessValues::add_points(region={:?}, points={:?})", region, points);
         self.points.union_row(region, points);
+
+        // When available, record the loans flowing into this region as live at the given points.
+        if let Some(loans) = self.loans.as_mut() {
+            if let Some(inflowing) = loans.inflowing_loans.row(region) {
+                if !inflowing.is_empty() {
+                    for point in points.iter() {
+                        loans.live_loans.union_row(point, inflowing);
+                    }
+                }
+            }
+        }
     }
 
     /// Records `region` as being live at all the control-flow points.
@@ -185,6 +236,15 @@ impl LivenessValues {
     pub(crate) fn point_from_location(&self, location: Location) -> PointIndex {
         self.elements.point_from_location(location)
     }
+
+    /// When using `-Zpolonius=next`, returns whether the `loan_idx` is live at the given `point`.
+    pub(crate) fn is_loan_live_at(&self, loan_idx: BorrowIndex, point: PointIndex) -> bool {
+        self.loans
+            .as_ref()
+            .expect("Accessing live loans requires `-Zpolonius=next`")
+            .live_loans
+            .contains(point, loan_idx)
+    }
 }
 
 /// Maps from `ty::PlaceholderRegion` values that are used in the rest of

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

@@ -1,12 +1,12 @@
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
 use rustc_data_structures::graph::WithSuccessors;
-use rustc_index::bit_set::{HybridBitSet, SparseBitMatrix};
+use rustc_index::bit_set::HybridBitSet;
 use rustc_index::interval::IntervalSet;
 use rustc_infer::infer::canonical::QueryRegionConstraints;
 use rustc_infer::infer::outlives::for_liveness;
 use rustc_middle::mir::{BasicBlock, Body, ConstraintCategory, Local, Location};
 use rustc_middle::traits::query::DropckOutlivesResult;
-use rustc_middle::ty::{RegionVid, Ty, TyCtxt, TypeVisitable, TypeVisitableExt};
+use rustc_middle::ty::{Ty, TyCtxt, TypeVisitable, TypeVisitableExt};
 use rustc_span::DUMMY_SP;
 use rustc_trait_selection::traits::query::type_op::outlives::DropckOutlives;
 use rustc_trait_selection::traits::query::type_op::{TypeOp, TypeOpOutput};
@@ -16,9 +16,8 @@ use rustc_mir_dataflow::impls::MaybeInitializedPlaces;
 use rustc_mir_dataflow::move_paths::{HasMoveData, MoveData, MovePathIndex};
 use rustc_mir_dataflow::ResultsCursor;
 
-use crate::dataflow::BorrowIndex;
 use crate::{
-    region_infer::values::{self, PointIndex, RegionValueElements},
+    region_infer::values::{self, LiveLoans, PointIndex, RegionValueElements},
     type_check::liveness::local_use_map::LocalUseMap,
     type_check::liveness::polonius,
     type_check::NormalizeLocation,
@@ -49,22 +48,17 @@ pub(super) fn trace<'mir, 'tcx>(
     boring_locals: Vec<Local>,
     polonius_drop_used: Option<Vec<(Local, Location)>>,
 ) {
-    debug!("trace()");
-
     let local_use_map = &LocalUseMap::build(&relevant_live_locals, elements, body);
 
     // When using `-Zpolonius=next`, compute the set of loans that can reach a given region.
-    let num_loans = typeck.borrowck_context.borrow_set.len();
-    let mut inflowing_loans = SparseBitMatrix::new(num_loans);
     if typeck.tcx().sess.opts.unstable_opts.polonius.is_next_enabled() {
-        let borrowck_context = &typeck.borrowck_context;
+        let borrowck_context = &mut typeck.borrowck_context;
         let borrow_set = &borrowck_context.borrow_set;
-        let constraint_set = &borrowck_context.constraints.outlives_constraints;
-
-        let num_region_vars = typeck.infcx.num_region_vars();
-        let graph = constraint_set.graph(num_region_vars);
+        let mut live_loans = LiveLoans::new(borrow_set.len());
+        let outlives_constraints = &borrowck_context.constraints.outlives_constraints;
+        let graph = outlives_constraints.graph(typeck.infcx.num_region_vars());
         let region_graph =
-            graph.region_graph(constraint_set, borrowck_context.universal_regions.fr_static);
+            graph.region_graph(outlives_constraints, borrowck_context.universal_regions.fr_static);
 
         // Traverse each issuing region's constraints, and record the loan as flowing into the
         // outlived region.
@@ -75,9 +69,13 @@ pub(super) fn trace<'mir, 'tcx>(
                     continue;
                 }
 
-                inflowing_loans.insert(succ, loan);
+                live_loans.inflowing_loans.insert(succ, loan);
             }
         }
+
+        // Store the inflowing loans in the liveness constraints: they will be used to compute live
+        // loans when liveness data is recorded there.
+        borrowck_context.constraints.liveness_constraints.loans = Some(live_loans);
     };
 
     let cx = LivenessContext {
@@ -88,7 +86,6 @@ pub(super) fn trace<'mir, 'tcx>(
         local_use_map,
         move_data,
         drop_data: FxIndexMap::default(),
-        inflowing_loans,
     };
 
     let mut results = LivenessResults::new(cx);
@@ -126,9 +123,6 @@ struct LivenessContext<'me, 'typeck, 'flow, 'tcx> {
     /// Index indicating where each variable is assigned, used, or
     /// dropped.
     local_use_map: &'me LocalUseMap,
-
-    /// Set of loans that flow into a given region, when using `-Zpolonius=next`.
-    inflowing_loans: SparseBitMatrix<RegionVid, BorrowIndex>,
 }
 
 struct DropData<'tcx> {
@@ -519,14 +513,7 @@ impl<'tcx> LivenessContext<'_, '_, '_, 'tcx> {
         live_at: &IntervalSet<PointIndex>,
     ) {
         debug!("add_use_live_facts_for(value={:?})", value);
-
-        Self::make_all_regions_live(
-            self.elements,
-            self.typeck,
-            value,
-            live_at,
-            &self.inflowing_loans,
-        );
+        Self::make_all_regions_live(self.elements, self.typeck, value, live_at);
     }
 
     /// Some variable with type `live_ty` is "drop live" at `location`
@@ -577,14 +564,7 @@ impl<'tcx> LivenessContext<'_, '_, '_, 'tcx> {
         // All things in the `outlives` array may be touched by
         // the destructor and must be live at this point.
         for &kind in &drop_data.dropck_result.kinds {
-            Self::make_all_regions_live(
-                self.elements,
-                self.typeck,
-                kind,
-                live_at,
-                &self.inflowing_loans,
-            );
-
+            Self::make_all_regions_live(self.elements, self.typeck, kind, live_at);
             polonius::add_drop_of_var_derefs_origin(self.typeck, dropped_local, &kind);
         }
     }
@@ -594,20 +574,13 @@ impl<'tcx> LivenessContext<'_, '_, '_, 'tcx> {
         typeck: &mut TypeChecker<'_, 'tcx>,
         value: impl TypeVisitable<TyCtxt<'tcx>>,
         live_at: &IntervalSet<PointIndex>,
-        inflowing_loans: &SparseBitMatrix<RegionVid, BorrowIndex>,
     ) {
         debug!("make_all_regions_live(value={:?})", value);
         debug!(
             "make_all_regions_live: live_at={}",
             values::pretty_print_points(elements, live_at.iter()),
         );
 
-        // When using `-Zpolonius=next`, we want to record the loans that flow into this value's
-        // regions as being live at the given `live_at` points: this will be used to compute the
-        // location where a loan goes out of scope.
-        let num_loans = typeck.borrowck_context.borrow_set.len();
-        let value_loans = &mut HybridBitSet::new_empty(num_loans);
-
         value.visit_with(&mut for_liveness::FreeRegionsVisitor {
             tcx: typeck.tcx(),
             param_env: typeck.param_env,
@@ -619,21 +592,8 @@ impl<'tcx> LivenessContext<'_, '_, '_, 'tcx> {
                     .constraints
                     .liveness_constraints
                     .add_points(live_region_vid, live_at);
-
-                // There can only be inflowing loans for this region when we are using
-                // `-Zpolonius=next`.
-                if let Some(inflowing) = inflowing_loans.row(live_region_vid) {
-                    value_loans.union(inflowing);
-                }
             },
         });
-
-        // Record the loans reaching the value.
-        if !value_loans.is_empty() {
-            for point in live_at.iter() {
-                typeck.borrowck_context.live_loans.union_row(point, value_loans);
-            }
-        }
     }
 
     fn compute_drop_data(