move and maintain live loans in LivenessValues

lqd · lqd · commit 60d4eb2c1ba6 · 2023-12-01T14:04:54.000Z
Liveness data is pushed from multiple parts of NLL. Instead of changing
the call sites to maintain live loans, move the latter to `LivenessValues` where
this liveness data is pushed to, and maintain live loans there.

This fixes the differences in polonius scopes on some CFGs where a
variable was dead in tracing but as a MIR terminator its regions were marked
live from "constraint generation"
diff --git a/compiler/rustc_borrowck/src/nll.rs b/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.
diff --git a/compiler/rustc_borrowck/src/region_infer/mod.rs b/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)
     }
 }
 
diff --git a/compiler/rustc_borrowck/src/region_infer/values.rs b/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
diff --git a/compiler/rustc_borrowck/src/type_check/liveness/trace.rs b/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,
@@ -52,15 +51,12 @@ pub(super) fn trace<'mir, 'tcx>(
     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 mut live_loans = LiveLoans::new(borrow_set.len());
         let outlives_constraints = &borrowck_context.constraints.outlives_constraints;
-
-        let num_region_vars = typeck.infcx.num_region_vars();
-        let graph = outlives_constraints.graph(num_region_vars);
+        let graph = outlives_constraints.graph(typeck.infcx.num_region_vars());
         let region_graph =
             graph.region_graph(outlives_constraints, borrowck_context.universal_regions.fr_static);
 
@@ -73,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 {
@@ -86,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);
@@ -124,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> {
@@ -517,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`
@@ -575,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);
         }
     }
@@ -592,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,
@@ -617,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(
diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs
