From 69e55f3da163cfb1f3a4e635a1d6d1e8030fb189 Mon Sep 17 00:00:00 2001 From: Aapo Alasuutari Date: Fri, 17 Jul 2026 14:08:50 +0300 Subject: [PATCH] Enable single Location to issue multiple borrows --- compiler/rustc_borrowck/src/borrow_set.rs | 267 ++++++++++-------- compiler/rustc_borrowck/src/dataflow.rs | 16 +- compiler/rustc_borrowck/src/lib.rs | 2 +- compiler/rustc_borrowck/src/path_utils.rs | 2 +- .../src/polonius/legacy/loan_invalidations.rs | 2 +- .../src/polonius/legacy/loan_kills.rs | 4 +- compiler/rustc_borrowck/src/type_check/mod.rs | 28 +- 7 files changed, 173 insertions(+), 148 deletions(-) diff --git a/compiler/rustc_borrowck/src/borrow_set.rs b/compiler/rustc_borrowck/src/borrow_set.rs index 6d68703642314..c92f94bf3c28c 100644 --- a/compiler/rustc_borrowck/src/borrow_set.rs +++ b/compiler/rustc_borrowck/src/borrow_set.rs @@ -1,63 +1,122 @@ +use std::collections::hash_map::Entry; use std::fmt; use std::ops::Index; -use rustc_data_structures::fx::{FxIndexMap, FxIndexSet}; +use rustc_data_structures::fx::{FxHashMap, FxIndexMap, FxIndexSet}; use rustc_hir::Mutability; +use rustc_index::IndexVec; use rustc_index::bit_set::DenseBitSet; use rustc_middle::mir::visit::{MutatingUseContext, NonUseContext, PlaceContext, Visitor}; use rustc_middle::mir::{self, Body, Local, Location, traversal}; +use rustc_middle::ty::data_structures::IndexSet; use rustc_middle::ty::{RegionVid, TyCtxt}; use rustc_middle::{bug, span_bug, ty}; use rustc_mir_dataflow::move_paths::MoveData; +use smallvec::{SmallVec, smallvec}; use tracing::debug; use crate::BorrowIndex; use crate::place_ext::PlaceExt; pub struct BorrowSet<'tcx> { + /// BorrowData storage. + borrows: IndexVec>, + /// The fundamental map relating bitvector indexes to the borrows - /// in the MIR. Each borrow is also uniquely identified in the MIR + /// in the MIR. Each borrow of a reference is uniquely identified in the MIR /// by the `Location` of the assignment statement in which it - /// appears on the right hand side. Thus the location is the map - /// key, and its position in the map corresponds to `BorrowIndex`. - pub(crate) location_map: FxIndexMap>, + /// appears on the right hand side, but for generic Reborrow there may be + /// multiple borrows per location. Thus the location is the map + /// key, and it identifies one or more `BorrowIndex` values. + /// + /// FIXME(reborrow): if the Reborrow experiment is rejected, this can be turned + /// back into a FxIndexMap or BorrowIndex>. See [PR]. + /// + /// [PR]: github.com/rust-lang/rust/pull/159449 + location_map: FxHashMap>, /// Locations which activate borrows. - /// NOTE: a given location may activate more than one borrow in the future - /// when more general two-phase borrow support is introduced, but for now we - /// only need to store one borrow index. - pub(crate) activation_map: FxIndexMap>, + activation_map: FxHashMap>, /// Map from local to all the borrows on that local. - pub(crate) local_map: FxIndexMap>, + local_map: FxIndexMap>, - pub(crate) locals_state_at_exit: LocalsStateAtExit, + locals_state_at_exit: LocalsStateAtExit, } -// These methods are public to support borrowck consumers. impl<'tcx> BorrowSet<'tcx> { - pub fn location_map(&self) -> &FxIndexMap> { - &self.location_map - } + // Public method to support Aquascope. + pub fn build( + tcx: TyCtxt<'tcx>, + body: &Body<'tcx>, + locals_are_invalidated_at_exit: bool, + move_data: &MoveData<'tcx>, + ) -> Self { + let mut visitor = GatherBorrows { + tcx, + body, + borrows: Default::default(), + location_map: Default::default(), + activation_map: Default::default(), + local_map: Default::default(), + pending_activations: Default::default(), + locals_state_at_exit: LocalsStateAtExit::build( + locals_are_invalidated_at_exit, + body, + move_data, + ), + }; + + for (block, block_data) in traversal::preorder(body) { + visitor.visit_basic_block_data(block, block_data); + } - pub fn activation_map(&self) -> &FxIndexMap> { - &self.activation_map + BorrowSet { + borrows: visitor.borrows, + location_map: visitor.location_map, + activation_map: visitor.activation_map, + local_map: visitor.local_map, + locals_state_at_exit: visitor.locals_state_at_exit, + } } - pub fn local_map(&self) -> &FxIndexMap> { - &self.local_map + // Public method to support Aquascope. + /// Iterate through all BorrowData in the BorrowSet. + pub fn iter(&self) -> impl Iterator> { + self.borrows.iter() } - pub fn locals_state_at_exit(&self) -> &LocalsStateAtExit { + // The following functions are not depended upon by outside consumers. + pub(crate) fn locals_state_at_exit(&self) -> &LocalsStateAtExit { &self.locals_state_at_exit } + + pub(crate) fn len(&self) -> usize { + self.borrows.len() + } + + pub(crate) fn iter_enumerated(&self) -> impl Iterator)> { + self.borrows.iter_enumerated() + } + + pub(crate) fn activations_at_location(&self, location: &Location) -> &[BorrowIndex] { + self.activation_map.get(&location).map_or(&[], |activations| &activations[..]) + } + + pub(crate) fn borrows_at_location(&self, location: &Location) -> Option<&[BorrowIndex]> { + self.location_map.get(location).map(|v| v.as_slice()) + } + + pub(crate) fn borrows_on_local(&self, local: Local) -> Option<&IndexSet> { + self.local_map.get(&local) + } } impl<'tcx> Index for BorrowSet<'tcx> { type Output = BorrowData<'tcx>; fn index(&self, index: BorrowIndex) -> &BorrowData<'tcx> { - &self.location_map[index.as_usize()] + &self.borrows[index] } } @@ -166,65 +225,12 @@ impl LocalsStateAtExit { } } -impl<'tcx> BorrowSet<'tcx> { - pub fn build( - tcx: TyCtxt<'tcx>, - body: &Body<'tcx>, - locals_are_invalidated_at_exit: bool, - move_data: &MoveData<'tcx>, - ) -> Self { - let mut visitor = GatherBorrows { - tcx, - body, - location_map: Default::default(), - activation_map: Default::default(), - local_map: Default::default(), - pending_activations: Default::default(), - locals_state_at_exit: LocalsStateAtExit::build( - locals_are_invalidated_at_exit, - body, - move_data, - ), - }; - - for (block, block_data) in traversal::preorder(body) { - visitor.visit_basic_block_data(block, block_data); - } - - BorrowSet { - location_map: visitor.location_map, - activation_map: visitor.activation_map, - local_map: visitor.local_map, - locals_state_at_exit: visitor.locals_state_at_exit, - } - } - - pub(crate) fn activations_at_location(&self, location: Location) -> &[BorrowIndex] { - self.activation_map.get(&location).map_or(&[], |activations| &activations[..]) - } - - pub(crate) fn len(&self) -> usize { - self.location_map.len() - } - - pub(crate) fn indices(&self) -> impl Iterator { - BorrowIndex::ZERO..BorrowIndex::from_usize(self.len()) - } - - pub(crate) fn iter_enumerated(&self) -> impl Iterator)> { - self.indices().zip(self.location_map.values()) - } - - pub(crate) fn get_index_of(&self, location: &Location) -> Option { - self.location_map.get_index_of(location).map(BorrowIndex::from) - } -} - struct GatherBorrows<'a, 'tcx> { tcx: TyCtxt<'tcx>, body: &'a Body<'tcx>, - location_map: FxIndexMap>, - activation_map: FxIndexMap>, + borrows: IndexVec>, + location_map: FxHashMap>, + activation_map: FxHashMap>, local_map: FxIndexMap>, /// When we encounter a 2-phase borrow statement, it will always @@ -240,6 +246,23 @@ struct GatherBorrows<'a, 'tcx> { locals_state_at_exit: LocalsStateAtExit, } +impl<'a, 'tcx> GatherBorrows<'a, 'tcx> { + fn insert_borrow(&mut self, location: Location, borrow: BorrowData<'tcx>) -> BorrowIndex { + let idx = self.borrows.push(borrow); + match self.location_map.entry(location) { + Entry::Occupied(entry) => { + bug!( + "Inserting a borrow {idx:?} at {location:?} attempted to override an existing list {entry:?}" + ); + } + Entry::Vacant(entry) => { + entry.insert(smallvec![idx]); + } + } + idx + } +} + impl<'a, 'tcx> Visitor<'tcx> for GatherBorrows<'a, 'tcx> { fn visit_assign( &mut self, @@ -265,10 +288,7 @@ impl<'a, 'tcx> Visitor<'tcx> for GatherBorrows<'a, 'tcx> { let idx = if !kind.is_two_phase_borrow() { debug!(" -> {:?}", location); - let (idx, _) = self - .location_map - .insert_full(location, borrow(TwoPhaseActivation::NotTwoPhase)); - BorrowIndex::from(idx) + self.insert_borrow(location, borrow(TwoPhaseActivation::NotTwoPhase)) } else { // When we encounter a 2-phase borrow statement, it will always // be assigning into a temporary TEMP: @@ -286,10 +306,7 @@ impl<'a, 'tcx> Visitor<'tcx> for GatherBorrows<'a, 'tcx> { // Consider the borrow not activated to start. When we find an activation, we'll update // this field. - let (idx, _) = self - .location_map - .insert_full(location, borrow(TwoPhaseActivation::NotActivated)); - let idx = BorrowIndex::from(idx); + let idx = self.insert_borrow(location, borrow(TwoPhaseActivation::NotActivated)); // Insert `temp` into the list of pending activations. From // now on, we'll be on the lookout for a use of it. Note that @@ -342,8 +359,7 @@ impl<'a, 'tcx> Visitor<'tcx> for GatherBorrows<'a, 'tcx> { borrowed_place, assigned_place: *assigned_place, }; - let (idx, _) = self.location_map.insert_full(location, borrow); - let idx = BorrowIndex::from(idx); + let idx = self.insert_borrow(location, borrow); self.local_map.entry(borrowed_place.local).or_default().insert(idx); } @@ -360,53 +376,56 @@ impl<'a, 'tcx> Visitor<'tcx> for GatherBorrows<'a, 'tcx> { // check whether we (earlier) saw a 2-phase borrow like // // TMP = &mut place - if let Some(&borrow_index) = self.pending_activations.get(&temp) { - let borrow_data = &mut self.location_map[borrow_index.as_usize()]; - - // Watch out: the use of TMP in the borrow itself - // doesn't count as an activation. =) - if borrow_data.reserve_location == location - && context == PlaceContext::MutatingUse(MutatingUseContext::Store) - { - return; - } + let Some(&borrow_index) = self.pending_activations.get(&temp) else { + return; + }; + let borrow_data = &mut self.borrows[borrow_index]; - if let TwoPhaseActivation::ActivatedAt(other_location) = borrow_data.activation_location - { - span_bug!( - self.body.source_info(location).span, - "found two uses for 2-phase borrow temporary {:?}: \ - {:?} and {:?}", - temp, - location, - other_location, - ); - } + // Watch out: the use of TMP in the borrow itself + // doesn't count as an activation. =) + if borrow_data.reserve_location == location + && context == PlaceContext::MutatingUse(MutatingUseContext::Store) + { + return; + } - // Otherwise, this is the unique later use that we expect. - // Double check: This borrow is indeed a two-phase borrow (that is, - // we are 'transitioning' from `NotActivated` to `ActivatedAt`) and - // we've not found any other activations (checked above). - assert_eq!( - borrow_data.activation_location, - TwoPhaseActivation::NotActivated, - "never found an activation for this borrow!", + if let TwoPhaseActivation::ActivatedAt(other_location) = borrow_data.activation_location { + span_bug!( + self.body.source_info(location).span, + "found two uses for 2-phase borrow temporary {:?}: \ + {:?} and {:?}", + temp, + location, + other_location, ); - self.activation_map.entry(location).or_default().push(borrow_index); - - borrow_data.activation_location = TwoPhaseActivation::ActivatedAt(location); } + + // Otherwise, this is the unique later use that we expect. + // Double check: This borrow is indeed a two-phase borrow (that is, + // we are 'transitioning' from `NotActivated` to `ActivatedAt`) and + // we've not found any other activations (checked above). + assert_eq!( + borrow_data.activation_location, + TwoPhaseActivation::NotActivated, + "never found an activation for this borrow!", + ); + self.activation_map.entry(location).or_default().push(borrow_index); + + borrow_data.activation_location = TwoPhaseActivation::ActivatedAt(location); } fn visit_rvalue(&mut self, rvalue: &mir::Rvalue<'tcx>, location: mir::Location) { if let &mir::Rvalue::Ref(region, kind, place) = rvalue { // double-check that we already registered a BorrowData for this - let borrow_data = &self.location_map[&location]; - assert_eq!(borrow_data.reserve_location, location); - assert_eq!(borrow_data.kind, kind); - assert_eq!(borrow_data.region, region.as_var()); - assert_eq!(borrow_data.borrowed_place, place); + let idxs = &self.location_map[&location]; + for idx in idxs { + let borrow_data = &self.borrows[*idx]; + assert_eq!(borrow_data.reserve_location, location); + assert_eq!(borrow_data.kind, kind); + assert_eq!(borrow_data.region, region.as_var()); + assert_eq!(borrow_data.borrowed_place, place); + } } self.super_rvalue(rvalue, location) diff --git a/compiler/rustc_borrowck/src/dataflow.rs b/compiler/rustc_borrowck/src/dataflow.rs index 0bc06ade95ea1..5bf692eaa7205 100644 --- a/compiler/rustc_borrowck/src/dataflow.rs +++ b/compiler/rustc_borrowck/src/dataflow.rs @@ -474,8 +474,7 @@ impl<'a, 'tcx> Borrows<'a, 'tcx> { let other_borrows_of_local = self .borrow_set - .local_map - .get(&place.local) + .borrows_on_local(place.local) .map(|bs| bs.iter().copied()) .into_flat_iter(); @@ -552,15 +551,18 @@ impl<'tcx> rustc_mir_dataflow::Analysis<'tcx> for Borrows<'_, 'tcx> { if place.ignore_borrow( self.tcx, self.body, - &self.borrow_set.locals_state_at_exit, + &self.borrow_set.locals_state_at_exit(), ) { return; } - let index = self.borrow_set.get_index_of(&location).unwrap_or_else(|| { - panic!("could not find BorrowIndex for location {location:?}"); - }); + let idxs = + self.borrow_set.borrows_at_location(&location).unwrap_or_else(|| { + panic!("could not find BorrowIndex for location {location:?}"); + }); - state.gen_(index); + for index in idxs { + state.gen_(*index); + } } // Make sure there are no remaining borrows for variables diff --git a/compiler/rustc_borrowck/src/lib.rs b/compiler/rustc_borrowck/src/lib.rs index 9d66863ef70e1..aa0e399aec1fd 100644 --- a/compiler/rustc_borrowck/src/lib.rs +++ b/compiler/rustc_borrowck/src/lib.rs @@ -1872,7 +1872,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, '_, 'tcx> { // Two-phase borrow support: For each activation that is newly // generated at this statement, check if it interferes with // another borrow. - for &borrow_index in self.borrow_set.activations_at_location(location) { + for &borrow_index in self.borrow_set.activations_at_location(&location) { let borrow = &self.borrow_set[borrow_index]; // only mutable borrows should be 2-phase diff --git a/compiler/rustc_borrowck/src/path_utils.rs b/compiler/rustc_borrowck/src/path_utils.rs index 2c94a32d369ce..51706a589ab28 100644 --- a/compiler/rustc_borrowck/src/path_utils.rs +++ b/compiler/rustc_borrowck/src/path_utils.rs @@ -26,7 +26,7 @@ pub(super) fn each_borrow_involving_path<'tcx, F, I, S>( // 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 }; + let Some(borrows_for_place_base) = borrow_set.borrows_on_local(place.local) else { return }; // check for loan restricting path P being used. Accounts for // borrows of P, P.a.b, etc. diff --git a/compiler/rustc_borrowck/src/polonius/legacy/loan_invalidations.rs b/compiler/rustc_borrowck/src/polonius/legacy/loan_invalidations.rs index 64ba106b68adc..26036aae114bf 100644 --- a/compiler/rustc_borrowck/src/polonius/legacy/loan_invalidations.rs +++ b/compiler/rustc_borrowck/src/polonius/legacy/loan_invalidations.rs @@ -423,7 +423,7 @@ impl<'a, 'tcx> LoanInvalidationsGenerator<'a, 'tcx> { // Two-phase borrow support: For each activation that is newly // generated at this statement, check if it interferes with // another borrow. - for &borrow_index in self.borrow_set.activations_at_location(location) { + for &borrow_index in self.borrow_set.activations_at_location(&location) { let borrow = &self.borrow_set[borrow_index]; // only mutable borrows should be 2-phase diff --git a/compiler/rustc_borrowck/src/polonius/legacy/loan_kills.rs b/compiler/rustc_borrowck/src/polonius/legacy/loan_kills.rs index 098c922bf7b91..bdf86b20dd301 100644 --- a/compiler/rustc_borrowck/src/polonius/legacy/loan_kills.rs +++ b/compiler/rustc_borrowck/src/polonius/legacy/loan_kills.rs @@ -115,7 +115,7 @@ impl<'tcx> LoanKillsGenerator<'_, 'tcx> { local, location ); - if let Some(borrow_indices) = self.borrow_set.local_map.get(&local) { + if let Some(borrow_indices) = self.borrow_set.borrows_on_local(local) { for &borrow_index in borrow_indices { let places_conflict = places_conflict::places_conflict( self.tcx, @@ -137,7 +137,7 @@ impl<'tcx> LoanKillsGenerator<'_, 'tcx> { /// Records the borrows on the specified local as `killed`. fn record_killed_borrows_for_local(&mut self, local: Local, location: Location) { - if let Some(borrow_indices) = self.borrow_set.local_map.get(&local) { + if let Some(borrow_indices) = self.borrow_set.borrows_on_local(local) { let location_index = self.location_table.mid_index(location); self.facts.loan_killed_at.reserve(borrow_indices.len()); for &borrow_index in borrow_indices { diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs index 00d91e3fc25ad..818ddfe7d75ef 100644 --- a/compiler/rustc_borrowck/src/type_check/mod.rs +++ b/compiler/rustc_borrowck/src/type_check/mod.rs @@ -2370,13 +2370,15 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> { // example). if let Some(polonius_facts) = polonius_facts { let _prof_timer = self.infcx.tcx.prof.generic_activity("polonius_fact_generation"); - if let Some(borrow_index) = borrow_set.get_index_of(&location) { + if let Some(idxs) = borrow_set.borrows_at_location(&location) { let region_vid = borrow_region.as_var(); - polonius_facts.loan_issued_at.push(( - region_vid.into(), - borrow_index, - location_table.mid_index(location), - )); + for borrow_index in idxs { + polonius_facts.loan_issued_at.push(( + region_vid.into(), + *borrow_index, + location_table.mid_index(location), + )); + } } } @@ -2513,13 +2515,15 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> { // linking the loan to the region. if let Some(polonius_facts) = polonius_facts { let _prof_timer = infcx.tcx.prof.generic_activity("polonius_fact_generation"); - if let Some(borrow_index) = borrow_set.get_index_of(&location) { + if let Some(borrows) = borrow_set.borrows_at_location(&location) { let region_vid = dest_region.as_var(); - polonius_facts.loan_issued_at.push(( - region_vid.into(), - borrow_index, - location_table.mid_index(location), - )); + for borrow_index in borrows { + polonius_facts.loan_issued_at.push(( + region_vid.into(), + *borrow_index, + location_table.mid_index(location), + )); + } } }