Reborrow traits

Author: aapoalas

compiler/rustc_borrowck/src/borrow_set.rs

@@ -2,11 +2,12 @@ use std::fmt;
 use std::ops::Index;
 
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
+use rustc_hir::Mutability;
 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::span_bug;
 use rustc_middle::ty::{RegionVid, TyCtxt};
+use rustc_middle::{bug, span_bug, ty};
 use rustc_mir_dataflow::move_paths::MoveData;
 use tracing::debug;
 
@@ -300,6 +301,50 @@ impl<'a, 'tcx> Visitor<'tcx> for GatherBorrows<'a, 'tcx> {
                 idx
             };
 
+            self.local_map.entry(borrowed_place.local).or_default().insert(idx);
+        } else if let &mir::Rvalue::Reborrow(target, mutability, borrowed_place) = rvalue {
+            let borrowed_place_ty = borrowed_place.ty(self.body, self.tcx).ty;
+            let &ty::Adt(reborrowed_adt, _reborrowed_args) = borrowed_place_ty.kind() else {
+                unreachable!()
+            };
+            let &ty::Adt(target_adt, assigned_args) = target.kind() else { unreachable!() };
+            let Some(ty::GenericArgKind::Lifetime(region)) = assigned_args.get(0).map(|r| r.kind())
+            else {
+                bug!(
+                    "hir-typeck passed but {} does not have a lifetime argument",
+                    if mutability == Mutability::Mut { "Reborrow" } else { "CoerceShared" }
+                );
+            };
+            let region = region.as_var();
+            let kind = if mutability == Mutability::Mut {
+                // Reborrow
+                if target_adt.did() != reborrowed_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() == reborrowed_adt.did() {
+                    bug!(
+                        "hir-typeck passed but CoerceShared involves matching types at {location:?}"
+                    )
+                }
+                mir::BorrowKind::Shared
+            };
+            let borrow = BorrowData {
+                kind,
+                region,
+                reserve_location: location,
+                activation_location: TwoPhaseActivation::NotTwoPhase,
+                borrowed_place,
+                assigned_place: *assigned_place,
+            };
+            let (idx, _) = self.location_map.insert_full(location, borrow);
+            let idx = BorrowIndex::from(idx);
+
             self.local_map.entry(borrowed_place.local).or_default().insert(idx);
         }
 

compiler/rustc_borrowck/src/dataflow.rs

@@ -549,7 +549,7 @@ impl<'tcx> rustc_mir_dataflow::Analysis<'tcx> for Borrows<'_, 'tcx> {
     ) {
         match &stmt.kind {
             mir::StatementKind::Assign(box (lhs, rhs)) => {
-                if let mir::Rvalue::Ref(_, _, place) = rhs {
+                if let mir::Rvalue::Ref(_, _, place) | mir::Rvalue::Reborrow(_, _, place) = rhs {
                     if place.ignore_borrow(
                         self.tcx,
                         self.body,

compiler/rustc_borrowck/src/lib.rs

@@ -1266,6 +1266,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, '_, 'tcx> {
         let mut error_reported = false;
 
         let borrows_in_scope = self.borrows_in_scope(location, state);
+        debug!(?borrows_in_scope, ?location);
 
         each_borrow_involving_path(
             self,
@@ -1508,6 +1509,36 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, '_, 'tcx> {
                 );
             }
 
+            &Rvalue::Reborrow(_target, mutability, place) => {
+                let access_kind = (
+                    Deep,
+                    if mutability == Mutability::Mut {
+                        Write(WriteKind::MutableBorrow(BorrowKind::Mut {
+                            kind: MutBorrowKind::Default,
+                        }))
+                    } else {
+                        Read(ReadKind::Borrow(BorrowKind::Shared))
+                    },
+                );
+
+                self.access_place(
+                    location,
+                    (place, span),
+                    access_kind,
+                    LocalMutationIsAllowed::Yes,
+                    state,
+                );
+
+                let action = InitializationRequiringAction::Borrow;
+
+                self.check_if_path_or_subpath_is_moved(
+                    location,
+                    action,
+                    (place.as_ref(), span),
+                    state,
+                );
+            }
+
             &Rvalue::RawPtr(kind, place) => {
                 let access_kind = match kind {
                     RawPtrKind::Mut => (

compiler/rustc_borrowck/src/polonius/legacy/loan_invalidations.rs

@@ -275,6 +275,21 @@ impl<'a, 'tcx> LoanInvalidationsGenerator<'a, 'tcx> {
                 self.access_place(location, place, access_kind, LocalMutationIsAllowed::No);
             }
 
+            &Rvalue::Reborrow(_target, mutability, place) => {
+                let access_kind = (
+                    Deep,
+                    if mutability == Mutability::Mut {
+                        Reservation(WriteKind::MutableBorrow(BorrowKind::Mut {
+                            kind: MutBorrowKind::TwoPhaseBorrow,
+                        }))
+                    } else {
+                        Read(ReadKind::Borrow(BorrowKind::Shared))
+                    },
+                );
+
+                self.access_place(location, place, access_kind, LocalMutationIsAllowed::No);
+            }
+
             &Rvalue::RawPtr(kind, place) => {
                 let access_kind = match kind {
                     RawPtrKind::Mut => (

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -1581,6 +1581,15 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                 self.add_reborrow_constraint(location, *region, borrowed_place);
             }
 
+            Rvalue::Reborrow(target, mutability, borrowed_place) => {
+                self.add_generic_reborrow_constraint(
+                    *mutability,
+                    location,
+                    borrowed_place,
+                    *target,
+                );
+            }
+
             Rvalue::BinaryOp(
                 BinOp::Eq | BinOp::Ne | BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge,
                 box (left, right),
@@ -2217,6 +2226,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             | Rvalue::ThreadLocalRef(_)
             | Rvalue::Repeat(..)
             | Rvalue::Ref(..)
+            | Rvalue::Reborrow(..)
             | Rvalue::RawPtr(..)
             | Rvalue::Cast(..)
             | Rvalue::BinaryOp(..)
@@ -2421,6 +2431,122 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
         }
     }
 
+    fn add_generic_reborrow_constraint(
+        &mut self,
+        mutability: Mutability,
+        location: Location,
+        borrowed_place: &Place<'tcx>,
+        dest_ty: Ty<'tcx>,
+    ) {
+        // These constraints are only meaningful during borrowck:
+        let Self { borrow_set, location_table, polonius_facts, constraints, infcx, body, .. } =
+            self;
+
+        // If we are reborrowing the referent of another reference, we
+        // need to add outlives relationships. In a case like `&mut
+        // *p`, where the `p` has type `&'b mut Foo`, for example, we
+        // need to ensure that `'b: 'a`.
+
+        debug!(
+            "add_generic_reborrow_constraint({:?}, {:?}, {:?}, {:?})",
+            mutability, location, borrowed_place, dest_ty
+        );
+
+        let tcx = infcx.tcx;
+        let def = body.source.def_id().expect_local();
+        let upvars = tcx.closure_captures(def);
+        let field =
+            path_utils::is_upvar_field_projection(tcx, upvars, borrowed_place.as_ref(), body);
+        let category = if let Some(field) = field {
+            ConstraintCategory::ClosureUpvar(field)
+        } else {
+            ConstraintCategory::Boring
+        };
+
+        let borrowed_ty = borrowed_place.ty(self.body, tcx).ty;
+
+        let ty::Adt(dest_adt, dest_args) = dest_ty.kind() else { bug!() };
+        let [dest_arg, ..] = ***dest_args else { bug!() };
+        let ty::GenericArgKind::Lifetime(dest_region) = dest_arg.kind() else { bug!() };
+        constraints.liveness_constraints.add_location(dest_region.as_var(), location);
+
+        // In Polonius mode, we also push a `loan_issued_at` fact
+        // 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) {
+                let region_vid = dest_region.as_var();
+                polonius_facts.loan_issued_at.push((
+                    region_vid.into(),
+                    borrow_index,
+                    location_table.mid_index(location),
+                ));
+            }
+        }
+
+        if mutability.is_not() {
+            // FIXME(@aapoalas): for CoerceShared we need to relate the types manually, field by
+            // field. We cannot just attempt to relate `T` and `<T as CoerceShared>::Target` by
+            // calling relate_types as they are (generally) two unrelated user-defined ADTs, such as
+            // `CustomMut<'a>` and `CustomRef<'a>`, or `CustomMut<'a, T>` and `CustomRef<'a, T>`.
+            // Field-by-field relate_types is expected to work based on the wf-checks that the
+            // CoerceShared trait performs.
+            let ty::Adt(borrowed_adt, borrowed_args) = borrowed_ty.kind() else { unreachable!() };
+            let borrowed_fields = borrowed_adt.all_fields().collect::<Vec<_>>();
+            for dest_field in dest_adt.all_fields() {
+                let Some(borrowed_field) =
+                    borrowed_fields.iter().find(|f| f.name == dest_field.name)
+                else {
+                    continue;
+                };
+                let dest_ty = dest_field.ty(tcx, dest_args);
+                let borrowed_ty = borrowed_field.ty(tcx, borrowed_args);
+                if let (
+                    ty::Ref(borrow_region, _, Mutability::Mut),
+                    ty::Ref(ref_region, _, Mutability::Not),
+                ) = (borrowed_ty.kind(), dest_ty.kind())
+                {
+                    self.relate_types(
+                        borrowed_ty.peel_refs(),
+                        ty::Variance::Covariant,
+                        dest_ty.peel_refs(),
+                        location.to_locations(),
+                        category,
+                    )
+                    .unwrap();
+                    self.constraints.outlives_constraints.push(OutlivesConstraint {
+                        sup: ref_region.as_var(),
+                        sub: borrow_region.as_var(),
+                        locations: location.to_locations(),
+                        span: location.to_locations().span(self.body),
+                        category,
+                        variance_info: ty::VarianceDiagInfo::default(),
+                        from_closure: false,
+                    });
+                } else {
+                    self.relate_types(
+                        borrowed_ty,
+                        ty::Variance::Covariant,
+                        dest_ty,
+                        location.to_locations(),
+                        category,
+                    )
+                    .unwrap();
+                }
+            }
+        } else {
+            // Exclusive reborrow
+            self.relate_types(
+                borrowed_ty,
+                ty::Variance::Covariant,
+                dest_ty,
+                location.to_locations(),
+                category,
+            )
+            .unwrap();
+        }
+    }
+
     fn prove_aggregate_predicates(
         &mut self,
         aggregate_kind: &AggregateKind<'tcx>,

compiler/rustc_codegen_cranelift/src/base.rs

@@ -629,6 +629,11 @@ fn codegen_stmt<'tcx>(fx: &mut FunctionCx<'_, '_, 'tcx>, cur_block: Block, stmt:
                     let ref_ = place.place_ref(fx, lval.layout());
                     lval.write_cvalue(fx, ref_);
                 }
+                Rvalue::Reborrow(_, _, place) => {
+                    let cplace = codegen_place(fx, place);
+                    let val = cplace.to_cvalue(fx);
+                    lval.write_cvalue(fx, val)
+                }
                 Rvalue::ThreadLocalRef(def_id) => {
                     let val = crate::constant::codegen_tls_ref(fx, def_id, lval.layout());
                     lval.write_cvalue(fx, val);

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -510,6 +510,14 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 self.codegen_place_to_pointer(bx, place, mk_ref)
             }
 
+            // Note: Exclusive reborrowing is always equal to a memcpy, as the types do not change.
+            // Generic shared reborrowing is not (necessarily) a simple memcpy, but currently the
+            // coherence check places such restrictions on the CoerceShared trait as to guarantee
+            // that it is.
+            mir::Rvalue::Reborrow(_, _, place) => {
+                self.codegen_operand(bx, &mir::Operand::Copy(place))
+            }
+
             mir::Rvalue::RawPtr(kind, place) => {
                 let mk_ptr = move |tcx: TyCtxt<'tcx>, ty: Ty<'tcx>| {
                     Ty::new_ptr(tcx, ty, kind.to_mutbl_lossy())

compiler/rustc_const_eval/src/check_consts/check.rs

@@ -609,6 +609,10 @@ impl<'tcx> Visitor<'tcx> for Checker<'_, 'tcx> {
                 }
             }
 
+            Rvalue::Reborrow(..) => {
+                // FIXME(@aapoalas): figure out if this is relevant at all.
+            }
+
             Rvalue::RawPtr(RawPtrKind::FakeForPtrMetadata, place) => {
                 // These are only inserted for slice length, so the place must already be indirect.
                 // This implies we do not have to worry about whether the borrow escapes.

compiler/rustc_const_eval/src/check_consts/qualifs.rs

@@ -255,6 +255,8 @@ where
             in_place::<Q, _>(cx, in_local, place.as_ref())
         }
 
+        Rvalue::Reborrow(_, _, place) => in_place::<Q, _>(cx, in_local, place.as_ref()),
+
         Rvalue::WrapUnsafeBinder(op, _) => in_operand::<Q, _>(cx, in_local, op),
 
         Rvalue::Aggregate(kind, operands) => {

compiler/rustc_const_eval/src/check_consts/resolver.rs

@@ -191,6 +191,16 @@ where
                 }
             }
 
+            mir::Rvalue::Reborrow(_target, mutability, borrowed_place) => {
+                if !borrowed_place.is_indirect() && mutability.is_mut() {
+                    let place_ty = borrowed_place.ty(self.ccx.body, self.ccx.tcx).ty;
+                    if Q::in_any_value_of_ty(self.ccx, place_ty) {
+                        self.state.qualif.insert(borrowed_place.local);
+                        self.state.borrow.insert(borrowed_place.local);
+                    }
+                }
+            }
+
             mir::Rvalue::Cast(..)
             | mir::Rvalue::Use(..)
             | mir::Rvalue::CopyForDeref(..)