Rollup merge of #151753 - aapoalas:aapoalas/reborrow-and-coerce-shared-first-version, r=oli-obk

Experiment: Reborrow traits

With this PR we now have basic functional Reborrow and CoerceShared traits. The current limitations are:

1. Reborrowable types can only have one lifetime parameter, so as to avoid having to figure out and store in rmeta the information of which lifetimes weaken during reborrowing.
2. Reborrowing of `&mut` wrappers is working (though I've not tested generic wrappers like `Option<&mut T>` yet), but CoerceShared of `&mut` wrappers currently causes an ICE.

The remaining tasks to complete before I'd consider this PR mergeable are:

- [x] Fix ICE on CoerceShared. Unfortunately this might require dipping into rmeta.
- [x] Expand the tests to give a more complete view of the current state of the experiment.

Reborrow traits experiment: #145612

Co-authored by @dingxiangfei2009

Author: JonathanBrouwer

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

@@ -1265,6 +1265,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,
@@ -1507,6 +1508,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

@@ -274,6 +274,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

@@ -1580,6 +1580,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),
@@ -2218,6 +2227,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             | Rvalue::ThreadLocalRef(..)
             | Rvalue::Repeat(..)
             | Rvalue::Ref(..)
+            | Rvalue::Reborrow(..)
             | Rvalue::RawPtr(..)
             | Rvalue::Cast(..)
             | Rvalue::BinaryOp(..)
@@ -2422,6 +2432,116 @@ 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>,
+    ) {
+        let Self { borrow_set, location_table, polonius_facts, constraints, infcx, body, .. } =
+            self;
+
+        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(reborrow): 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

@@ -518,6 +518,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

@@ -610,6 +610,10 @@ impl<'tcx> Visitor<'tcx> for Checker<'_, 'tcx> {
                 }
             }
 
+            Rvalue::Reborrow(..) => {
+                // FIXME(reborrow): 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

@@ -252,6 +252,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,19 @@ where
                 }
             }
 
+            mir::Rvalue::Reborrow(target, mutability, borrowed_place) => {
+                // A Reborrow allows mutation if it is Reborrow or if the CoerceShared target isn't
+                // Freeze.
+                if !borrowed_place.is_indirect()
+                    && (mutability.is_mut() || !target.is_freeze(self.ccx.tcx, self.ccx.typing_env))
+                {
+                    if Q::in_any_value_of_ty(self.ccx, *target) {
+                        self.state.qualif.insert(borrowed_place.local);
+                        self.state.borrow.insert(borrowed_place.local);
+                    }
+                }
+            }
+
             mir::Rvalue::Cast(..)
             | mir::Rvalue::Use(..)
             | mir::Rvalue::CopyForDeref(..)