compute loan scopes with liveness in -Zpolonius=next

Author: lqd

compiler/rustc_borrowck/src/dataflow.rs

@@ -238,15 +238,231 @@ pub fn calculate_borrows_out_of_scope_at_location<'tcx>(
     prec.borrows_out_of_scope_at_location
 }
 
+struct StackEntry {
+    bb: mir::BasicBlock,
+    lo: usize,
+    hi: usize,
+}
+
+struct PoloniusOutOfScopePrecomputer<'a, 'tcx> {
+    visited: BitSet<mir::BasicBlock>,
+    visit_stack: Vec<StackEntry>,
+    body: &'a Body<'tcx>,
+    regioncx: &'a RegionInferenceContext<'tcx>,
+    loans_out_of_scope_at_location: FxIndexMap<Location, Vec<BorrowIndex>>,
+    placeholders: Vec<RegionVid>,
+}
+
+impl<'a, 'tcx> PoloniusOutOfScopePrecomputer<'a, 'tcx> {
+    fn new(body: &'a Body<'tcx>, regioncx: &'a RegionInferenceContext<'tcx>) -> Self {
+        // Compute the placeholder list once, as it will be used for all the loan scopes we'll
+        // compute.
+        // FIXME: they're surely already available somewhere.
+        let placeholders = regioncx
+            .regions()
+            .filter(|&r| {
+                use rustc_infer::infer::*;
+                let origin = regioncx.var_infos[r].origin;
+                let is_placeholder = matches!(
+                    origin,
+                    RegionVariableOrigin::Nll(NllRegionVariableOrigin::Placeholder(_))
+                );
+                is_placeholder
+            })
+            .collect();
+
+        Self {
+            visited: BitSet::new_empty(body.basic_blocks.len()),
+            visit_stack: vec![],
+            body,
+            regioncx,
+            loans_out_of_scope_at_location: FxIndexMap::default(),
+            placeholders,
+        }
+    }
+}
+
+impl<'tcx> PoloniusOutOfScopePrecomputer<'_, 'tcx> {
+    /// Loans are in scope while they are live: whether they are contained within any live region.
+    /// In the location-insensitive analysis, a loan will be contained in a region if the issuing
+    /// region can reach it in the subset graph. So this is a reachability problem.
+    fn precompute_loans_out_of_scope(
+        &mut self,
+        loan_idx: BorrowIndex,
+        issuing_region: RegionVid,
+        issued_location: Location,
+    ) {
+        // Let's precompute the reachability set of the issuing region, via reachability on the
+        // condensation graph. We can also early return when reaching regions that outlive free
+        // regions via member constraints. (The `OutOfScopePrecomputer` wouldn't be called on a
+        // region that outlives free regions via outlives constraints.)
+
+        let liveness = &self.regioncx.liveness_constraints;
+        let sccs = &self.regioncx.constraint_sccs;
+
+        let mut reachability = BitSet::new_empty(sccs.num_sccs());
+
+        let issuing_region_scc = sccs.scc(issuing_region);
+        let mut stack = vec![issuing_region_scc];
+        reachability.insert(issuing_region_scc);
+
+        let member_constraints = &self.regioncx.member_constraints;
+
+        while let Some(scc) = stack.pop() {
+            // Handle successors of this SCC:
+            //
+            // 1. Push outlives successors to the worklist stack
+            for &succ_scc in sccs.successors(scc) {
+                if reachability.insert(succ_scc) {
+                    stack.push(succ_scc);
+                }
+            }
+
+            // 2. Deal with member constraints
+            //
+            // The issuing region can flow into the choice regions here, and they are either:
+            // - placeholders or free regions themselves,
+            // - or also transitively outlive a free region.
+            //
+            // That is to say, if there are member constraints here, the loan escapes the
+            // function and cannot go out of scope. We can early return.
+            if member_constraints.indices(scc).next().is_some() {
+                return;
+            }
+        }
+
+        // We visit one BB at a time. The complication is that we may start in the
+        // middle of the first BB visited (the one containing `location`), in which
+        // case we may have to later on process the first part of that BB if there
+        // is a path back to its start.
+
+        // For visited BBs, we record the index of the first statement processed.
+        // (In fully processed BBs this index is 0.) Note also that we add BBs to
+        // `visited` once they are added to `stack`, before they are actually
+        // processed, because this avoids the need to look them up again on
+        // completion.
+        self.visited.insert(issued_location.block);
+
+        let mut first_lo = issued_location.statement_index;
+        let first_hi = self.body[issued_location.block].statements.len();
+
+        self.visit_stack.push(StackEntry { bb: issued_location.block, lo: first_lo, hi: first_hi });
+
+        while let Some(StackEntry { bb, lo, hi }) = self.visit_stack.pop() {
+            // If we process the first part of the first basic block (i.e. we encounter that block
+            // for the second time), we no longer have to visit its successors again.
+            let mut finished_early = bb == issued_location.block && hi != first_hi;
+            for i in lo..=hi {
+                let location = Location { block: bb, statement_index: i };
+
+                // The loan is out of scope at point `location` if it's not contained within any
+                // live regions.
+                let mut issuing_region_can_reach_live_regions = false;
+
+                // Check reachability of all live regions:
+                // - the local regions that are live at this point,
+                // - the placeholders, which are live at all points and don't need liveness to be
+                //   computed, and are thus absent from the liveness values.
+                //
+                // As mentioned above, we don't need to check for free regions, if the issuing
+                // region outlived a free region via outlives constraints, we wouldn't need to
+                // compute its loan's scope.
+                for live_region in
+                    liveness.live_regions_at(location).chain(self.placeholders.iter().copied())
+                {
+                    let live_region_scc = sccs.scc(live_region);
+
+                    // If a single live region is reachable from the issuing region, then the loan
+                    // is still live at this point. We can stop checking other live regions at this
+                    // location, and go to the next location.
+                    if reachability.contains(live_region_scc) {
+                        issuing_region_can_reach_live_regions = true;
+                        break;
+                    }
+                }
+
+                // If no live region is reachable from the issuing region, then the loan is
+                // killed at this point, and goes out of scope.
+                if !issuing_region_can_reach_live_regions {
+                    debug!("loan {:?} gets killed at {:?}", loan_idx, location);
+                    self.loans_out_of_scope_at_location.entry(location).or_default().push(loan_idx);
+                    finished_early = true;
+                    break;
+                }
+            }
+
+            if !finished_early {
+                // Add successor BBs to the work list, if necessary.
+                let bb_data = &self.body[bb];
+                debug_assert!(hi == bb_data.statements.len());
+                for succ_bb in bb_data.terminator().successors() {
+                    if !self.visited.insert(succ_bb) {
+                        if succ_bb == issued_location.block && first_lo > 0 {
+                            // `succ_bb` has been seen before. If it wasn't fully processed, add its
+                            // first part to `visit_stack` for processing.
+                            self.visit_stack.push(StackEntry {
+                                bb: succ_bb,
+                                lo: 0,
+                                hi: first_lo - 1,
+                            });
+
+                            // And update this entry with 0, to represent the whole BB being
+                            // processed.
+                            first_lo = 0;
+                        }
+                    } else {
+                        // `succ_bb` hasn't been seen before. Add it to `visit_stack` for
+                        // processing.
+                        self.visit_stack.push(StackEntry {
+                            bb: succ_bb,
+                            lo: 0,
+                            hi: self.body[succ_bb].statements.len(),
+                        });
+                    }
+                }
+            }
+        }
+
+        assert!(self.visit_stack.is_empty(), "stack should be empty");
+
+        self.visited.clear();
+    }
+}
+
 impl<'a, 'tcx> Borrows<'a, 'tcx> {
     pub fn new(
         tcx: TyCtxt<'tcx>,
         body: &'a Body<'tcx>,
         regioncx: &'a RegionInferenceContext<'tcx>,
         borrow_set: &'a BorrowSet<'tcx>,
     ) -> Self {
-        let borrows_out_of_scope_at_location =
+        let mut borrows_out_of_scope_at_location =
             calculate_borrows_out_of_scope_at_location(body, regioncx, borrow_set);
+
+        // The in-tree polonius analysis computes loans going out of scope using the set-of-loans
+        // model, and makes sure they're identical to the existing computation of the set-of-points
+        // model.
+        if tcx.sess.opts.unstable_opts.polonius.is_next_enabled() {
+            let mut polonius_prec = PoloniusOutOfScopePrecomputer::new(body, regioncx);
+            for (loan_idx, loan_data) in borrow_set.iter_enumerated() {
+                let issuing_region = loan_data.region;
+                let issued_location = loan_data.reserve_location;
+
+                polonius_prec.precompute_loans_out_of_scope(
+                    loan_idx,
+                    issuing_region,
+                    issued_location,
+                );
+            }
+
+            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"
+            );
+
+            borrows_out_of_scope_at_location = polonius_prec.loans_out_of_scope_at_location;
+        }
+
         Borrows { tcx, body, borrow_set, borrows_out_of_scope_at_location }
     }