WIP: implement borrowck of pinned borrows

Step 1: implement `PinSet`

Author: frank-king

compiler/rustc_borrowck/src/borrow_set.rs

@@ -1,6 +1,7 @@
 use std::fmt;
 use std::ops::Index;
 
+use rustc_abi::FieldIdx;
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
 use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::mir::visit::{MutatingUseContext, NonUseContext, PlaceContext, Visitor};
@@ -69,6 +70,13 @@ pub enum TwoPhaseActivation {
     ActivatedAt(Location),
 }
 
+/// Location where a place is pinned, if a borrowed place is pinned.
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
+pub(crate) enum Pinnedness<'tcx> {
+    Not,
+    Pinned { to: mir::Place<'tcx>, at: Location },
+}
+
 #[derive(Debug, Clone)]
 pub struct BorrowData<'tcx> {
     /// Location where the borrow reservation starts.
@@ -84,6 +92,8 @@ pub struct BorrowData<'tcx> {
     pub(crate) borrowed_place: mir::Place<'tcx>,
     /// Place to which the borrow was stored
     pub(crate) assigned_place: mir::Place<'tcx>,
+    /// Place and location to the pinning statement, if the borrowed place is pinned.
+    pub(crate) pinnedness: Pinnedness<'tcx>,
 }
 
 // These methods are public to support borrowck consumers.
@@ -260,6 +270,7 @@ impl<'a, 'tcx> Visitor<'tcx> for GatherBorrows<'a, 'tcx> {
                 activation_location,
                 borrowed_place,
                 assigned_place: *assigned_place,
+                pinnedness: Pinnedness::Not,
             };
 
             let idx = if !kind.is_two_phase_borrow() {
@@ -300,9 +311,40 @@ impl<'a, 'tcx> Visitor<'tcx> for GatherBorrows<'a, 'tcx> {
                 idx
             };
 
+            if let Some((pinned_place, successor)) =
+                succeeded_by_pin(self.tcx, self.body, location, *assigned_place)
+            {
+                self.location_map[&location].pinnedness =
+                    Pinnedness::Pinned { to: pinned_place, at: successor };
+            }
+
             self.local_map.entry(borrowed_place.local).or_default().insert(idx);
         }
 
+        // Whether the current statement is followed by a
+        // `pinned_place = Pin::<&[mut] T> { pointer: move assigned_place }` statement.
+        fn succeeded_by_pin<'tcx>(
+            tcx: TyCtxt<'tcx>,
+            body: &mir::Body<'tcx>,
+            location: mir::Location,
+            assigned_place: mir::Place<'tcx>,
+        ) -> Option<(mir::Place<'tcx>, mir::Location)> {
+            let successor = location.successor_within_block();
+            let stmt =
+                body.basic_blocks[successor.block].statements.get(successor.statement_index)?;
+            if let mir::StatementKind::Assign(box (pinned_place, ref rvalue)) = stmt.kind
+                && let mir::Rvalue::Aggregate(box agg_kind, operands) = rvalue
+                && let mir::AggregateKind::Adt(adt_did, _, args, _, _) = agg_kind
+                && tcx.adt_def(adt_did).is_pin()
+                && args.type_at(0).is_ref()
+                && let mir::Operand::Move(place) = operands[FieldIdx::ZERO]
+                && place == assigned_place
+            {
+                return Some((pinned_place, successor));
+            }
+            None
+        }
+
         self.super_assign(assigned_place, rvalue, location)
     }
 

compiler/rustc_borrowck/src/dataflow.rs

@@ -14,13 +14,16 @@ use rustc_mir_dataflow::impls::{
 use rustc_mir_dataflow::{Analysis, GenKill, JoinSemiLattice};
 use tracing::debug;
 
-use crate::{BorrowSet, PlaceConflictBias, PlaceExt, RegionInferenceContext, places_conflict};
+use crate::{
+    BorrowSet, PinSet, PlaceConflictBias, PlaceExt, RegionInferenceContext, places_conflict,
+};
 
 // This analysis is different to most others. Its results aren't computed with
 // `iterate_to_fixpoint`, but are instead composed from the results of three sub-analyses that are
 // computed individually with `iterate_to_fixpoint`.
 pub(crate) struct Borrowck<'a, 'tcx> {
     pub(crate) borrows: Borrows<'a, 'tcx>,
+    pub(crate) pins: Pins<'a, 'tcx>,
     pub(crate) uninits: MaybeUninitializedPlaces<'a, 'tcx>,
     pub(crate) ever_inits: EverInitializedPlaces<'a, 'tcx>,
 }
@@ -33,6 +36,7 @@ impl<'a, 'tcx> Analysis<'tcx> for Borrowck<'a, 'tcx> {
     fn bottom_value(&self, body: &mir::Body<'tcx>) -> Self::Domain {
         BorrowckDomain {
             borrows: self.borrows.bottom_value(body),
+            pins: self.pins.bottom_value(body),
             uninits: self.uninits.bottom_value(body),
             ever_inits: self.ever_inits.bottom_value(body),
         }
@@ -50,6 +54,7 @@ impl<'a, 'tcx> Analysis<'tcx> for Borrowck<'a, 'tcx> {
         loc: Location,
     ) {
         self.borrows.apply_early_statement_effect(&mut state.borrows, stmt, loc);
+        self.pins.apply_early_statement_effect(&mut state.pins, stmt, loc);
         self.uninits.apply_early_statement_effect(&mut state.uninits, stmt, loc);
         self.ever_inits.apply_early_statement_effect(&mut state.ever_inits, stmt, loc);
     }
@@ -61,6 +66,7 @@ impl<'a, 'tcx> Analysis<'tcx> for Borrowck<'a, 'tcx> {
         loc: Location,
     ) {
         self.borrows.apply_primary_statement_effect(&mut state.borrows, stmt, loc);
+        self.pins.apply_primary_statement_effect(&mut state.pins, stmt, loc);
         self.uninits.apply_primary_statement_effect(&mut state.uninits, stmt, loc);
         self.ever_inits.apply_primary_statement_effect(&mut state.ever_inits, stmt, loc);
     }
@@ -72,6 +78,7 @@ impl<'a, 'tcx> Analysis<'tcx> for Borrowck<'a, 'tcx> {
         loc: Location,
     ) {
         self.borrows.apply_early_terminator_effect(&mut state.borrows, term, loc);
+        self.pins.apply_early_terminator_effect(&mut state.pins, term, loc);
         self.uninits.apply_early_terminator_effect(&mut state.uninits, term, loc);
         self.ever_inits.apply_early_terminator_effect(&mut state.ever_inits, term, loc);
     }
@@ -83,6 +90,7 @@ impl<'a, 'tcx> Analysis<'tcx> for Borrowck<'a, 'tcx> {
         loc: Location,
     ) -> TerminatorEdges<'mir, 'tcx> {
         self.borrows.apply_primary_terminator_effect(&mut state.borrows, term, loc);
+        self.pins.apply_primary_terminator_effect(&mut state.pins, term, loc);
         self.uninits.apply_primary_terminator_effect(&mut state.uninits, term, loc);
         self.ever_inits.apply_primary_terminator_effect(&mut state.ever_inits, term, loc);
 
@@ -116,6 +124,8 @@ where
     fn fmt_with(&self, ctxt: &C, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.write_str("borrows: ")?;
         self.borrows.fmt_with(ctxt, f)?;
+        f.write_str(" pins: ")?;
+        self.pins.fmt_with(ctxt, f)?;
         f.write_str(" uninits: ")?;
         self.uninits.fmt_with(ctxt, f)?;
         f.write_str(" ever_inits: ")?;
@@ -134,6 +144,12 @@ where
             f.write_str("\n")?;
         }
 
+        if self.pins != old.pins {
+            f.write_str("pins: ")?;
+            self.pins.fmt_diff_with(&old.pins, ctxt, f)?;
+            f.write_str("\n")?;
+        }
+
         if self.uninits != old.uninits {
             f.write_str("uninits: ")?;
             self.uninits.fmt_diff_with(&old.uninits, ctxt, f)?;
@@ -154,6 +170,7 @@ where
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub(crate) struct BorrowckDomain {
     pub(crate) borrows: BorrowsDomain,
+    pub(crate) pins: PinsDomain,
     pub(crate) uninits: MaybeUninitializedPlacesDomain,
     pub(crate) ever_inits: EverInitializedPlacesDomain,
 }
@@ -164,6 +181,16 @@ rustc_index::newtype_index! {
     pub struct BorrowIndex {}
 }
 
+impl<C> DebugWithContext<C> for BorrowIndex {}
+
+rustc_index::newtype_index! {
+    #[orderable]
+    #[debug_format = "pi{}"]
+    pub struct PinIndex {}
+}
+
+impl<C> DebugWithContext<C> for PinIndex {}
+
 /// `Borrows` stores the data used in the analyses that track the flow
 /// of borrows.
 ///
@@ -178,6 +205,19 @@ pub struct Borrows<'a, 'tcx> {
     borrows_out_of_scope_at_location: FxIndexMap<Location, Vec<BorrowIndex>>,
 }
 
+/// `Pins` stores the data used in the analyses that track the flow
+/// of pins.
+///
+/// It uniquely identifies every pinned borrow by a
+/// `PinIndex`, and maps each such index to a `PinData`
+/// describing the pin. These indexes are used for representing the
+/// pins in compact bitvectors.
+pub(crate) struct Pins<'a, 'tcx> {
+    tcx: TyCtxt<'tcx>,
+    body: &'a Body<'tcx>,
+    pin_set: &'a PinSet<'tcx>,
+}
+
 struct OutOfScopePrecomputer<'a, 'tcx> {
     visited: DenseBitSet<mir::BasicBlock>,
     visit_stack: Vec<mir::BasicBlock>,
@@ -617,4 +657,153 @@ impl<'tcx> rustc_mir_dataflow::Analysis<'tcx> for Borrows<'_, 'tcx> {
     }
 }
 
-impl<C> DebugWithContext<C> for BorrowIndex {}
+impl<'a, 'tcx> Pins<'a, 'tcx> {
+    pub(crate) fn new(
+        tcx: TyCtxt<'tcx>,
+        body: &'a Body<'tcx>,
+        pin_set: &'a PinSet<'tcx>,
+    ) -> Self {
+        Pins { tcx, body, pin_set }
+    }
+
+    /// Kill any pins on `place`.
+    fn kill_pins_on_place(
+        &self,
+        state: &mut <Self as Analysis<'tcx>>::Domain,
+        place: Place<'tcx>,
+    ) {
+        debug!("kill_pins_on_place: place={:?}", place);
+
+        let other_pins_of_local = self
+            .pin_set
+            .local_map
+            .get(&place.local)
+            .into_iter()
+            .flat_map(|ps| ps.iter())
+            .copied();
+
+        // If the place is a local with no projections, all pins of this
+        // local must be killed. This is purely an optimization so we don't have to call
+        // `places_conflict` for every pin.
+        if place.projection.is_empty() {
+            state.kill_all(other_pins_of_local);
+            return;
+        }
+
+        // By passing `PlaceConflictBias::NoOverlap`, we conservatively assume that any given
+        // pair of array indices are not equal, so that when `places_conflict` returns true, we
+        // will be assured that two places being compared definitely denotes the same sets of
+        // locations.
+        let definitely_conflicting_pins = other_pins_of_local.filter(|&i| {
+            places_conflict(
+                self.tcx,
+                self.body,
+                self.pin_set[i].pinned_place,
+                place,
+                PlaceConflictBias::NoOverlap,
+            )
+        });
+
+        state.kill_all(definitely_conflicting_pins);
+    }
+}
+
+type PinsDomain = MixedBitSet<PinIndex>;
+
+/// Forward dataflow computation of the set of pins that are in scope at a particular location.
+/// - we gen the introduced pins
+/// - we kill pins on locals going out of scope
+/// - we kill pins when the pinned place is moved or overwritten
+impl<'tcx> rustc_mir_dataflow::Analysis<'tcx> for Pins<'_, 'tcx> {
+    type Domain = PinsDomain;
+
+    const NAME: &'static str = "pins";
+
+    fn bottom_value(&self, _: &mir::Body<'tcx>) -> Self::Domain {
+        // bottom = nothing is pinned yet
+        MixedBitSet::new_empty(self.pin_set.len())
+    }
+
+    fn initialize_start_block(&self, _: &mir::Body<'tcx>, _: &mut Self::Domain) {
+        // no pins have been created prior to function execution
+    }
+
+    fn apply_early_statement_effect(
+        &self,
+        _state: &mut Self::Domain,
+        _statement: &mir::Statement<'tcx>,
+        _location: Location,
+    ) {
+        // Pins don't go out of scope based on regions like borrows do
+    }
+
+    fn apply_primary_statement_effect(
+        &self,
+        state: &mut Self::Domain,
+        stmt: &mir::Statement<'tcx>,
+        location: Location,
+    ) {
+        match &stmt.kind {
+            mir::StatementKind::Assign(box (lhs, rhs)) => {
+                // Check if this is a Pin aggregate creation
+                if let mir::Rvalue::Aggregate(box agg_kind, _) = rhs
+                    && let mir::AggregateKind::Adt(adt_did, _, args, _, _) = agg_kind
+                    && self.tcx.adt_def(adt_did).is_pin()
+                    && args.type_at(0).is_ref()
+                {
+                    // Generate the pin
+                    if let Some(index) = self.pin_set.get_index_of(&location) {
+                        state.gen_(index);
+                    }
+                }
+
+                // Kill any pins on the place being assigned to
+                self.kill_pins_on_place(state, *lhs);
+            }
+
+            mir::StatementKind::StorageDead(local) => {
+                // Kill all pins on locals that are going out of scope
+                self.kill_pins_on_place(state, Place::from(*local));
+            }
+
+            mir::StatementKind::FakeRead(..)
+            | mir::StatementKind::SetDiscriminant { .. }
+            | mir::StatementKind::StorageLive(..)
+            | mir::StatementKind::Retag { .. }
+            | mir::StatementKind::PlaceMention(..)
+            | mir::StatementKind::AscribeUserType(..)
+            | mir::StatementKind::Coverage(..)
+            | mir::StatementKind::Intrinsic(..)
+            | mir::StatementKind::ConstEvalCounter
+            | mir::StatementKind::BackwardIncompatibleDropHint { .. }
+            | mir::StatementKind::Nop => {}
+        }
+    }
+
+    fn apply_early_terminator_effect(
+        &self,
+        _state: &mut Self::Domain,
+        _terminator: &mir::Terminator<'tcx>,
+        _location: Location,
+    ) {
+        // No early terminator effects for pins
+    }
+
+    fn apply_primary_terminator_effect<'mir>(
+        &self,
+        state: &mut Self::Domain,
+        terminator: &'mir mir::Terminator<'tcx>,
+        _location: Location,
+    ) -> TerminatorEdges<'mir, 'tcx> {
+        if let mir::TerminatorKind::InlineAsm { operands, .. } = &terminator.kind {
+            for op in operands {
+                if let mir::InlineAsmOperand::Out { place: Some(place), .. }
+                | mir::InlineAsmOperand::InOut { out_place: Some(place), .. } = *op
+                {
+                    self.kill_pins_on_place(state, place);
+                }
+            }
+        }
+        terminator.edges()
+    }
+}