Only consider places with the same local in each_borrow_involving_path.

Author: cjgillot

compiler/rustc_borrowck/src/invalidation.rs

@@ -351,15 +351,14 @@ impl<'cx, 'tcx> InvalidationGenerator<'cx, 'tcx> {
         let tcx = self.tcx;
         let body = self.body;
         let borrow_set = self.borrow_set;
-        let indices = self.borrow_set.indices();
         each_borrow_involving_path(
             self,
             tcx,
             body,
             location,
             (sd, place),
             borrow_set,
-            indices,
+            |_| true,
             |this, borrow_index, borrow| {
                 match (rw, borrow.kind) {
                     // Obviously an activation is compatible with its own

compiler/rustc_borrowck/src/lib.rs

@@ -23,7 +23,7 @@ use rustc_errors::{Diagnostic, DiagnosticBuilder, DiagnosticMessage, Subdiagnost
 use rustc_fluent_macro::fluent_messages;
 use rustc_hir as hir;
 use rustc_hir::def_id::LocalDefId;
-use rustc_index::bit_set::ChunkedBitSet;
+use rustc_index::bit_set::{BitSet, ChunkedBitSet};
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_infer::infer::{
     DefiningAnchor, InferCtxt, NllRegionVariableOrigin, RegionVariableOrigin, TyCtxtInferExt,
@@ -41,7 +41,6 @@ use rustc_session::lint::builtin::UNUSED_MUT;
 use rustc_span::{Span, Symbol};
 use rustc_target::abi::FieldIdx;
 
-use either::Either;
 use smallvec::SmallVec;
 use std::cell::OnceCell;
 use std::cell::RefCell;
@@ -1049,12 +1048,16 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
         let borrow_set = self.borrow_set.clone();
 
         // Use polonius output if it has been enabled.
-        let polonius_output = self.polonius_output.clone();
-        let borrows_in_scope = if let Some(polonius) = &polonius_output {
+        let mut polonius_output;
+        let borrows_in_scope = if let Some(polonius) = &self.polonius_output {
             let location = self.location_table.start_index(location);
-            Either::Left(polonius.errors_at(location).iter().copied())
+            polonius_output = BitSet::new_empty(borrow_set.len());
+            for &idx in polonius.errors_at(location) {
+                polonius_output.insert(idx);
+            }
+            &polonius_output
         } else {
-            Either::Right(flow_state.borrows.iter())
+            &flow_state.borrows
         };
 
         each_borrow_involving_path(
@@ -1064,7 +1067,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
             location,
             (sd, place_span.0),
             &borrow_set,
-            borrows_in_scope,
+            |borrow_index| borrows_in_scope.contains(borrow_index),
             |this, borrow_index, borrow| match (rw, borrow.kind) {
                 // Obviously an activation is compatible with its own
                 // reservation (or even prior activating uses of same

compiler/rustc_borrowck/src/path_utils.rs

@@ -33,20 +33,24 @@ pub(super) fn each_borrow_involving_path<'tcx, F, I, S>(
     _location: Location,
     access_place: (AccessDepth, Place<'tcx>),
     borrow_set: &BorrowSet<'tcx>,
-    candidates: I,
+    is_candidate: I,
     mut op: F,
 ) where
     F: FnMut(&mut S, BorrowIndex, &BorrowData<'tcx>) -> Control,
-    I: Iterator<Item = BorrowIndex>,
+    I: Fn(BorrowIndex) -> bool,
 {
     let (access, place) = access_place;
 
-    // FIXME: analogous code in check_loans first maps `place` to
-    // its base_path.
+    // The number of candidates can be large, but borrows for different locals cannot conflict with
+    // each other, so we restrict the working set a priori.
+    let Some(borrows_for_place_base) = borrow_set.local_map.get(&place.local) else { return };
 
     // check for loan restricting path P being used. Accounts for
     // borrows of P, P.a.b, etc.
-    for i in candidates {
+    for &i in borrows_for_place_base {
+        if !is_candidate(i) {
+            continue;
+        }
         let borrowed = &borrow_set[i];
 
         if places_conflict::borrow_conflicts_with_place(