| use std::collections::hash_map::Entry; | |
| use std::fmt; | |
| use std::ops::Index; | |
| 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, PlaceElem, traversal}; | |
| use rustc_middle::ty; | |
| use rustc_middle::ty::data_structures::IndexSet; | |
| use rustc_middle::ty::{RegionVid, TyCtxt}; | |
| use rustc_mir_dataflow::move_paths::MoveData; | |
| use rustc_span::{bug, span_bug}; | |
| use smallvec::{SmallVec, smallvec}; | |
| use tracing::debug; | |
| use crate::BorrowIndex; | |
| use crate::place_ext::PlaceExt; | |
| pub struct BorrowSet<'tcx> { | |
| /// BorrowData storage. | |
| borrows: IndexVec<BorrowIndex, BorrowData<'tcx>>, | |
| /// The fundamental map relating bitvector indexes to the borrows | |
| /// 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, 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<Location, BorrowData<'tcx> or BorrowIndex>. See [PR]. | |
| /// | |
| /// [PR]: github.com/rust-lang/rust/pull/159449 | |
| location_map: FxHashMap<Location, SmallVec<[BorrowIndex; 1]>>, | |
| /// Locations which activate borrows. | |
| activation_map: FxHashMap<Location, SmallVec<[BorrowIndex; 1]>>, | |
| /// Map from local to all the borrows on that local. | |
| local_map: FxIndexMap<mir::Local, FxIndexSet<BorrowIndex>>, | |
| locals_state_at_exit: LocalsStateAtExit, | |
| } | |
| impl<'tcx> BorrowSet<'tcx> { | |
| // 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); | |
| } | |
| 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, | |
| } | |
| } | |
| // Public method to support Aquascope and Creusot. | |
| /// Iterate through all BorrowData in the BorrowSet. | |
| pub fn iter(&self) -> impl Iterator<Item = &BorrowData<'tcx>> { | |
| self.borrows.iter() | |
| } | |
| // Public method to support Creusot. | |
| pub fn locals_state_at_exit(&self) -> &LocalsStateAtExit { | |
| &self.locals_state_at_exit | |
| } | |
| // Public method to support Creusot. | |
| pub fn len(&self) -> usize { | |
| self.borrows.len() | |
| } | |
| pub fn iter_enumerated(&self) -> impl Iterator<Item = (BorrowIndex, &BorrowData<'tcx>)> { | |
| self.borrows.iter_enumerated() | |
| } | |
| // Public method to support Creusot. | |
| pub fn activations_at_location(&self, location: &Location) -> &[BorrowIndex] { | |
| self.activation_map.get(&location).map_or(&[], |activations| &activations[..]) | |
| } | |
| // Public method to support Creusot. | |
| pub fn borrows_at_location(&self, location: &Location) -> Option<&[BorrowIndex]> { | |
| self.location_map.get(location).map(|v| v.as_slice()) | |
| } | |
| // Public method to support Creusot. | |
| pub fn borrows_on_local(&self, local: Local) -> Option<&IndexSet<BorrowIndex>> { | |
| self.local_map.get(&local) | |
| } | |
| } | |
| impl<'tcx> Index<BorrowIndex> for BorrowSet<'tcx> { | |
| type Output = BorrowData<'tcx>; | |
| fn index(&self, index: BorrowIndex) -> &BorrowData<'tcx> { | |
| &self.borrows[index] | |
| } | |
| } | |
| /// Location where a two-phase borrow is activated, if a borrow | |
| /// is in fact a two-phase borrow. | |
| #[derive(Copy, Clone, PartialEq, Eq, Debug)] | |
| pub enum TwoPhaseActivation { | |
| NotTwoPhase, | |
| NotActivated, | |
| ActivatedAt(Location), | |
| } | |
| #[derive(Debug, Clone)] | |
| pub struct BorrowData<'tcx> { | |
| /// Location where the borrow reservation starts. | |
| /// In many cases, this will be equal to the activation location but not always. | |
| pub(crate) reserve_location: Location, | |
| /// Location where the borrow is activated. | |
| pub(crate) activation_location: TwoPhaseActivation, | |
| /// What kind of borrow this is | |
| pub(crate) kind: mir::BorrowKind, | |
| /// The region for which this borrow is live | |
| pub(crate) region: RegionVid, | |
| /// Place from which we are borrowing | |
| pub(crate) borrowed_place: mir::Place<'tcx>, | |
| /// Place to which the borrow was stored | |
| pub(crate) assigned_place: mir::Place<'tcx>, | |
| } | |
| // These methods are public to support borrowck consumers. | |
| impl<'tcx> BorrowData<'tcx> { | |
| pub fn reserve_location(&self) -> Location { | |
| self.reserve_location | |
| } | |
| pub fn activation_location(&self) -> TwoPhaseActivation { | |
| self.activation_location | |
| } | |
| pub fn kind(&self) -> mir::BorrowKind { | |
| self.kind | |
| } | |
| pub fn region(&self) -> RegionVid { | |
| self.region | |
| } | |
| pub fn borrowed_place(&self) -> mir::Place<'tcx> { | |
| self.borrowed_place | |
| } | |
| pub fn assigned_place(&self) -> mir::Place<'tcx> { | |
| self.assigned_place | |
| } | |
| } | |
| impl<'tcx> fmt::Display for BorrowData<'tcx> { | |
| fn fmt(&self, w: &mut fmt::Formatter<'_>) -> fmt::Result { | |
| let kind = match self.kind { | |
| mir::BorrowKind::Shared => "", | |
| mir::BorrowKind::Fake(mir::FakeBorrowKind::Deep) => "fake ", | |
| mir::BorrowKind::Fake(mir::FakeBorrowKind::Shallow) => "fake shallow ", | |
| mir::BorrowKind::Mut { kind: mir::MutBorrowKind::ClosureCapture } => "uniq ", | |
| // FIXME: differentiate `TwoPhaseBorrow` | |
| mir::BorrowKind::Mut { | |
| kind: mir::MutBorrowKind::Default | mir::MutBorrowKind::TwoPhaseBorrow, | |
| } => "mut ", | |
| }; | |
| write!(w, "&{:?} {}{:?}", self.region, kind, self.borrowed_place) | |
| } | |
| } | |
| pub enum LocalsStateAtExit { | |
| AllAreInvalidated, | |
| SomeAreInvalidated { has_storage_dead_or_moved: DenseBitSet<Local> }, | |
| } | |
| impl LocalsStateAtExit { | |
| fn build<'tcx>( | |
| locals_are_invalidated_at_exit: bool, | |
| body: &Body<'tcx>, | |
| move_data: &MoveData<'tcx>, | |
| ) -> Self { | |
| struct HasStorageDead(DenseBitSet<Local>); | |
| impl<'tcx> Visitor<'tcx> for HasStorageDead { | |
| fn visit_local(&mut self, local: Local, ctx: PlaceContext, _: Location) { | |
| if ctx == PlaceContext::NonUse(NonUseContext::StorageDead) { | |
| self.0.insert(local); | |
| } | |
| } | |
| } | |
| if locals_are_invalidated_at_exit { | |
| LocalsStateAtExit::AllAreInvalidated | |
| } else { | |
| let mut has_storage_dead = | |
| HasStorageDead(DenseBitSet::new_empty(body.local_decls.len())); | |
| has_storage_dead.visit_body(body); | |
| let mut has_storage_dead_or_moved = has_storage_dead.0; | |
| for move_out in &move_data.move_outs { | |
| has_storage_dead_or_moved.insert(move_data.base_local(move_out.path)); | |
| } | |
| LocalsStateAtExit::SomeAreInvalidated { has_storage_dead_or_moved } | |
| } | |
| } | |
| } | |
| struct GatherBorrows<'a, 'tcx> { | |
| tcx: TyCtxt<'tcx>, | |
| body: &'a Body<'tcx>, | |
| borrows: IndexVec<BorrowIndex, BorrowData<'tcx>>, | |
| location_map: FxHashMap<Location, SmallVec<[BorrowIndex; 1]>>, | |
| activation_map: FxHashMap<Location, SmallVec<[BorrowIndex; 1]>>, | |
| local_map: FxIndexMap<mir::Local, FxIndexSet<BorrowIndex>>, | |
| /// When we encounter a 2-phase borrow statement, it will always | |
| /// be assigning into a temporary TEMP: | |
| /// | |
| /// TEMP = &foo | |
| /// | |
| /// We add TEMP into this map with `b`, where `b` is the index of | |
| /// the borrow. When we find a later use of this activation, we | |
| /// remove from the map (and add to the "tombstone" set below). | |
| pending_activations: FxIndexMap<mir::Local, BorrowIndex>, | |
| 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 | |
| } | |
| fn insert_borrows( | |
| &mut self, | |
| location: Location, | |
| borrows: SmallVec<[BorrowData<'tcx>; 1]>, | |
| ) -> SmallVec<[BorrowIndex; 1]> { | |
| let mut idxs = SmallVec::<[BorrowIndex; 1]>::with_capacity(borrows.len()); | |
| // FIXME(reborrow): why doesn't SmallVec offer reserve? | |
| for borrow in borrows { | |
| idxs.push(self.borrows.push(borrow)); | |
| } | |
| match self.location_map.entry(location) { | |
| Entry::Occupied(entry) => { | |
| bug!( | |
| "Inserting borrows {idxs:?} at {location:?} attempted to override an existing list {entry:?}" | |
| ); | |
| } | |
| Entry::Vacant(entry) => { | |
| entry.insert(idxs.clone()); | |
| } | |
| } | |
| idxs | |
| } | |
| fn gather_reborrows( | |
| &mut self, | |
| v: &mut SmallVec<[BorrowData<'tcx>; 1]>, | |
| kind: mir::BorrowKind, | |
| location: Location, | |
| target_adt: ty::AdtDef<'tcx>, | |
| target_args: &'tcx ty::List<ty::GenericArg<'tcx>>, | |
| target_place: mir::Place<'tcx>, | |
| source_adt: ty::AdtDef<'tcx>, | |
| source_args: &'tcx ty::List<ty::GenericArg<'tcx>>, | |
| source_place: mir::Place<'tcx>, | |
| ) { | |
| let mut did_reborrow = false; | |
| for (source_idx, source_field) in source_adt.all_fields().enumerate() { | |
| let source_field_ty = source_field.ty(self.tcx, source_args).skip_norm_wip(); | |
| match source_field_ty.kind() { | |
| ty::Ref(source_region, _, source_mutability) if source_mutability.is_mut() => { | |
| if source_region.is_static() { | |
| bug!( | |
| "Cannot implement Reborrow on a type containing a &'static mut T field" | |
| ); | |
| } | |
| let Some((target_idx, target_field)) = target_adt | |
| .all_fields() | |
| .enumerate() | |
| .find(|(_, f)| f.name == source_field.name) | |
| else { | |
| // Reborrow dropped this field. | |
| continue; | |
| }; | |
| let ty::Ref(target_region, _, _) = | |
| target_field.ty(self.tcx, target_args).skip_norm_wip().kind() | |
| else { | |
| bug!( | |
| "Reborrow source field type is &mut T but target field is not a reference" | |
| ); | |
| }; | |
| did_reborrow = true; | |
| let source_field_deref_place = source_place.project_deeper( | |
| &[PlaceElem::Field(source_idx.into(), source_field_ty), PlaceElem::Deref], | |
| self.tcx, | |
| ); | |
| let target_field_place = target_place.project_to_field( | |
| target_idx.into(), | |
| &self.body.local_decls, | |
| self.tcx, | |
| ); | |
| v.push(BorrowData { | |
| kind, | |
| region: target_region.as_var(), | |
| reserve_location: location, | |
| activation_location: TwoPhaseActivation::NotTwoPhase, | |
| borrowed_place: source_field_deref_place, | |
| assigned_place: target_field_place, | |
| }); | |
| } | |
| ty::Adt(source_field_adt, source_field_args) | |
| if source_field_args.get(0).is_some_and(|f| f.as_region().is_some()) | |
| && !self.tcx.type_is_copy_modulo_regions( | |
| self.body.typing_env(self.tcx), | |
| self.tcx.erase_and_anonymize_regions(source_field_ty), | |
| ) => | |
| { | |
| let Some((target_idx, target_field)) = target_adt | |
| .all_fields() | |
| .enumerate() | |
| .find(|(_, f)| f.name == source_field.name) | |
| else { | |
| // Reborrow dropped this field. | |
| continue; | |
| }; | |
| let ty::Adt(target_field_adt, target_field_args) = | |
| target_field.ty(self.tcx, target_args).skip_norm_wip().kind() | |
| else { | |
| bug!("Reborrow source field type is a !Copy ADT but target field is not"); | |
| }; | |
| did_reborrow = true; | |
| let source_field_place = source_place.project_to_field( | |
| source_idx.into(), | |
| &self.body.local_decls, | |
| self.tcx, | |
| ); | |
| let target_field_place = target_place.project_to_field( | |
| target_idx.into(), | |
| &self.body.local_decls, | |
| self.tcx, | |
| ); | |
| self.gather_reborrows( | |
| v, | |
| kind, | |
| location, | |
| *target_field_adt, | |
| target_field_args, | |
| target_field_place, | |
| *source_field_adt, | |
| source_field_args, | |
| source_field_place, | |
| ); | |
| } | |
| _ => continue, | |
| } | |
| } | |
| if !did_reborrow { | |
| // Key point: if source contained no reference, a phantom dereference must be performed | |
| // to avoid capturing the local variable's place. | |
| let source_phantom_deref_place = | |
| source_place.project_deeper(&[PlaceElem::PhantomDeref], self.tcx); | |
| if target_args.regions().count() != 1 { | |
| bug!( | |
| "ADT containing no '&mut T' or 'T: Reborrow' fields must only have one lifetime to implement Reborrow" | |
| ); | |
| } | |
| let target_region = target_args.regions().next().unwrap(); | |
| v.push(BorrowData { | |
| kind, | |
| region: target_region.as_var(), | |
| reserve_location: location, | |
| activation_location: TwoPhaseActivation::NotTwoPhase, | |
| borrowed_place: source_phantom_deref_place, | |
| assigned_place: target_place, | |
| }); | |
| } | |
| } | |
| } | |
| impl<'a, 'tcx> Visitor<'tcx> for GatherBorrows<'a, 'tcx> { | |
| fn visit_assign( | |
| &mut self, | |
| assigned_place: &mir::Place<'tcx>, | |
| rvalue: &mir::Rvalue<'tcx>, | |
| location: mir::Location, | |
| ) { | |
| if let &mir::Rvalue::Ref(region, kind, borrowed_place) = rvalue { | |
| if borrowed_place.ignore_borrow(self.tcx, self.body, &self.locals_state_at_exit) { | |
| debug!("ignoring_borrow of {:?}", borrowed_place); | |
| return; | |
| } | |
| let region = region.as_var(); | |
| let borrow = |activation_location| BorrowData { | |
| kind, | |
| region, | |
| reserve_location: location, | |
| activation_location, | |
| borrowed_place, | |
| assigned_place: *assigned_place, | |
| }; | |
| let idx = if !kind.is_two_phase_borrow() { | |
| debug!(" -> {:?}", location); | |
| 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: | |
| // | |
| // TEMP = &foo | |
| // | |
| // so extract `temp`. | |
| let Some(temp) = assigned_place.as_local() else { | |
| span_bug!( | |
| self.body.source_info(location).span, | |
| "expected 2-phase borrow to assign to a local, not `{:?}`", | |
| assigned_place, | |
| ); | |
| }; | |
| // Consider the borrow not activated to start. When we find an activation, we'll update | |
| // this field. | |
| 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 | |
| // we are guaranteed that this use will come after the | |
| // assignment. | |
| let prev = self.pending_activations.insert(temp, idx); | |
| assert_eq!(prev, None, "temporary associated with multiple two phase borrows"); | |
| idx | |
| }; | |
| self.local_map.entry(borrowed_place.local).or_default().insert(idx); | |
| } else if let &mir::Rvalue::Reborrow(target, mutability, source_place) = rvalue { | |
| let source_ty = source_place.ty(self.body, self.tcx).ty; | |
| let &ty::Adt(source_adt, source_args) = source_ty.kind() else { unreachable!() }; | |
| let &ty::Adt(target_adt, target_args) = target.kind() else { unreachable!() }; | |
| let kind = if mutability == Mutability::Mut { | |
| // Reborrow | |
| if target_adt.did() != source_adt.did() { | |
| bug!( | |
| "hir-typeck passed but Reborrow involves mismatching types at {location:?}" | |
| ) | |
| } | |
| mir::BorrowKind::Mut { kind: mir::MutBorrowKind::Default } | |
| } else { | |
| // CoerceShared | |
| if target_adt.did() == source_adt.did() { | |
| bug!( | |
| "hir-typeck passed but CoerceShared involves matching types at {location:?}" | |
| ) | |
| } | |
| mir::BorrowKind::Shared | |
| }; | |
| let mut reborrows = smallvec![]; | |
| self.gather_reborrows( | |
| &mut reborrows, | |
| kind, | |
| location, | |
| target_adt, | |
| target_args, | |
| *assigned_place, | |
| source_adt, | |
| source_args, | |
| source_place, | |
| ); | |
| let idxs = self.insert_borrows(location, reborrows); | |
| let locals = self.local_map.entry(source_place.local).or_default(); | |
| for idx in idxs { | |
| locals.insert(idx); | |
| } | |
| } | |
| self.super_assign(assigned_place, rvalue, location) | |
| } | |
| fn visit_local(&mut self, temp: Local, context: PlaceContext, location: Location) { | |
| if !context.is_use() { | |
| return; | |
| } | |
| // We found a use of some temporary TMP | |
| // check whether we (earlier) saw a 2-phase borrow like | |
| // | |
| // TMP = &mut place | |
| let Some(&borrow_index) = self.pending_activations.get(&temp) else { | |
| return; | |
| }; | |
| let borrow_data = &mut self.borrows[borrow_index]; | |
| // 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; | |
| } | |
| 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, | |
| ); | |
| } | |
| // 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 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) | |
| } | |
| } |