introduce reachability to the constraint graph

3 files changed, 112 insertions, 10 deletions

diff --git a/compiler/rustc_borrowck/src/nll.rs b/compiler/rustc_borrowck/src/nll.rs
index aa0bfd72147..590cbef05d5 100644
--- a/compiler/rustc_borrowck/src/nll.rs
+++ b/compiler/rustc_borrowck/src/nll.rs
@@ -103,7 +103,7 @@ pub(crate) fn compute_regions<'a, 'tcx>(
         constraints,
         universal_region_relations,
         opaque_type_values,
-        mut polonius_context,
+        polonius_context,
     } = type_check::type_check(
         infcx,
         body,
@@ -142,10 +142,10 @@ pub(crate) fn compute_regions<'a, 'tcx>(
         location_map,
     );
 
-    // If requested for `-Zpolonius=next`, convert NLL constraints to localized outlives
-    // constraints.
-    let localized_outlives_constraints = polonius_context.as_mut().map(|polonius_context| {
-        polonius_context.create_localized_constraints(infcx.tcx, &regioncx, body)
+    // If requested for `-Zpolonius=next`, convert NLL constraints to localized outlives constraints
+    // and use them to compute loan liveness.
+    let localized_outlives_constraints = polonius_context.as_ref().map(|polonius_context| {
+        polonius_context.compute_loan_liveness(infcx.tcx, &regioncx, body, borrow_set)
     });
 
     // If requested: dump NLL facts, and run legacy polonius analysis.
diff --git a/compiler/rustc_borrowck/src/polonius/loan_liveness.rs b/compiler/rustc_borrowck/src/polonius/loan_liveness.rs
new file mode 100644
index 00000000000..796a63da79d
--- /dev/null
+++ b/compiler/rustc_borrowck/src/polonius/loan_liveness.rs
@@ -0,0 +1,82 @@
+use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexSet};
+use rustc_middle::mir::Body;
+use rustc_middle::ty::{RegionVid, TyCtxt};
+use rustc_mir_dataflow::points::PointIndex;
+
+use super::{LiveLoans, LocalizedOutlivesConstraintSet};
+use crate::BorrowSet;
+use crate::region_infer::values::LivenessValues;
+
+/// With the full graph of constraints, we can compute loan reachability, and trace loan liveness
+/// throughout the CFG.
+pub(super) fn compute_loan_liveness<'tcx>(
+    _tcx: TyCtxt<'tcx>,
+    _body: &Body<'tcx>,
+    liveness: &LivenessValues,
+    borrow_set: &BorrowSet<'tcx>,
+    localized_outlives_constraints: &LocalizedOutlivesConstraintSet,
+) -> LiveLoans {
+    let mut live_loans = LiveLoans::new(borrow_set.len());
+    let graph = index_constraints(&localized_outlives_constraints);
+    let mut visited = FxHashSet::default();
+    let mut stack = Vec::new();
+
+    // Compute reachability per loan by traversing each loan's subgraph starting from where it is
+    // introduced.
+    for (loan_idx, loan) in borrow_set.iter_enumerated() {
+        visited.clear();
+        stack.clear();
+
+        let start_node = LocalizedNode {
+            region: loan.region,
+            point: liveness.point_from_location(loan.reserve_location),
+        };
+        stack.push(start_node);
+
+        while let Some(node) = stack.pop() {
+            if !visited.insert(node) {
+                continue;
+            }
+
+            // Record the loan as being live on entry to this point.
+            live_loans.insert(node.point, loan_idx);
+
+            for succ in outgoing_edges(&graph, node) {
+                stack.push(succ);
+            }
+        }
+    }
+
+    live_loans
+}
+
+/// The localized constraint graph is currently the per-node map of its physical edges. In the
+/// future, we'll add logical edges to model constraints that hold at all points in the CFG.
+type LocalizedConstraintGraph = FxHashMap<LocalizedNode, FxIndexSet<LocalizedNode>>;
+
+/// A node in the graph to be traversed, one of the two vertices of a localized outlives constraint.
+#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
+struct LocalizedNode {
+    region: RegionVid,
+    point: PointIndex,
+}
+
+/// Index the outlives constraints into a graph of edges per node.
+fn index_constraints(constraints: &LocalizedOutlivesConstraintSet) -> LocalizedConstraintGraph {
+    let mut edges = LocalizedConstraintGraph::default();
+    for constraint in &constraints.outlives {
+        let source = LocalizedNode { region: constraint.source, point: constraint.from };
+        let target = LocalizedNode { region: constraint.target, point: constraint.to };
+        edges.entry(source).or_default().insert(target);
+    }
+
+    edges
+}
+
+/// Returns the outgoing edges of a given node, not its transitive closure.
+fn outgoing_edges(
+    graph: &LocalizedConstraintGraph,
+    node: LocalizedNode,
+) -> impl Iterator<Item = LocalizedNode> + use<'_> {
+    graph.get(&node).into_iter().flat_map(|edges| edges.iter().copied())
+}
diff --git a/compiler/rustc_borrowck/src/polonius/mod.rs b/compiler/rustc_borrowck/src/polonius/mod.rs
index 7d0f9397021..313ed5d6649 100644
--- a/compiler/rustc_borrowck/src/polonius/mod.rs
+++ b/compiler/rustc_borrowck/src/polonius/mod.rs
@@ -37,6 +37,7 @@ mod constraints;
 mod dump;
 pub(crate) mod legacy;
 mod liveness_constraints;
+mod loan_liveness;
 mod typeck_constraints;
 
 use std::collections::BTreeMap;
@@ -49,8 +50,12 @@ use rustc_mir_dataflow::points::PointIndex;
 pub(crate) use self::constraints::*;
 pub(crate) use self::dump::dump_polonius_mir;
 use self::liveness_constraints::create_liveness_constraints;
+use self::loan_liveness::compute_loan_liveness;
 use self::typeck_constraints::convert_typeck_constraints;
-use crate::RegionInferenceContext;
+use crate::dataflow::BorrowIndex;
+use crate::{BorrowSet, RegionInferenceContext};
+
+pub(crate) type LiveLoans = SparseBitMatrix<PointIndex, BorrowIndex>;
 
 /// This struct holds the data needed to create the Polonius localized constraints.
 pub(crate) struct PoloniusContext {
@@ -82,14 +87,20 @@ impl PoloniusContext {
         Self { live_region_variances: BTreeMap::new(), live_regions: None }
     }
 
-    /// Creates a constraint set for `-Zpolonius=next` by:
+    /// Computes live loans using the set of loans model for `-Zpolonius=next`.
+    ///
+    /// First, creates a constraint graph combining regions and CFG points, by:
     /// - converting NLL typeck constraints to be localized
     /// - encoding liveness constraints
-    pub(crate) fn create_localized_constraints<'tcx>(
+    ///
+    /// Then, this graph is traversed, and combined with kills, reachability is recorded as loan
+    /// liveness, to be used by the loan scope and active loans computations.
+    pub(crate) fn compute_loan_liveness<'tcx>(
         &self,
         tcx: TyCtxt<'tcx>,
         regioncx: &RegionInferenceContext<'tcx>,
         body: &Body<'tcx>,
+        borrow_set: &BorrowSet<'tcx>,
     ) -> LocalizedOutlivesConstraintSet {
         let mut localized_outlives_constraints = LocalizedOutlivesConstraintSet::default();
         convert_typeck_constraints(
@@ -113,8 +124,17 @@ impl PoloniusContext {
             &mut localized_outlives_constraints,
         );
 
-        // FIXME: here, we can trace loan reachability in the constraint graph and record this as loan
-        // liveness for the next step in the chain, the NLL loan scope and active loans computations.
+        // Now that we have a complete graph, we can compute reachability to trace the liveness of
+        // loans for the next step in the chain, the NLL loan scope and active loans computations.
+        let _live_loans = compute_loan_liveness(
+            tcx,
+            body,
+            regioncx.liveness_constraints(),
+            borrow_set,
+            &localized_outlives_constraints,
+        );
+        // FIXME: record the live loans in the regioncx's liveness constraints, where the
+        // location-insensitive variant's data is stored.
 
         localized_outlives_constraints
     }