rustc_borrowck/

borrow_set.rs

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)
    }
}