Auto merge of #56307 - RalfJung:stacked-borrows-2-phase, r=oli-obk

Stacked borrows: Support 2 phase borrows

Some infrastructure needed for Stacked Borrows to be able to handle 2-phase-borros.

r? @oli-obk

src/librustc/ich/impls_mir.rs

@@ -193,51 +193,18 @@ for mir::TerminatorKind<'gcx> {
 
 impl_stable_hash_for!(struct mir::Statement<'tcx> { source_info, kind });
 
-impl<'a, 'gcx> HashStable<StableHashingContext<'a>>
-for mir::StatementKind<'gcx> {
-    fn hash_stable<W: StableHasherResult>(&self,
-                                          hcx: &mut StableHashingContext<'a>,
-                                          hasher: &mut StableHasher<W>) {
-        mem::discriminant(self).hash_stable(hcx, hasher);
-
-        match *self {
-            mir::StatementKind::Assign(ref place, ref rvalue) => {
-                place.hash_stable(hcx, hasher);
-                rvalue.hash_stable(hcx, hasher);
-            }
-            mir::StatementKind::FakeRead(ref cause, ref place) => {
-                cause.hash_stable(hcx, hasher);
-                place.hash_stable(hcx, hasher);
-            }
-            mir::StatementKind::SetDiscriminant { ref place, variant_index } => {
-                place.hash_stable(hcx, hasher);
-                variant_index.hash_stable(hcx, hasher);
-            }
-            mir::StatementKind::StorageLive(ref place) |
-            mir::StatementKind::StorageDead(ref place) => {
-                place.hash_stable(hcx, hasher);
-            }
-            mir::StatementKind::EscapeToRaw(ref place) => {
-                place.hash_stable(hcx, hasher);
-            }
-            mir::StatementKind::Retag { fn_entry, ref place } => {
-                fn_entry.hash_stable(hcx, hasher);
-                place.hash_stable(hcx, hasher);
-            }
-            mir::StatementKind::AscribeUserType(ref place, ref variance, ref c_ty) => {
-                place.hash_stable(hcx, hasher);
-                variance.hash_stable(hcx, hasher);
-                c_ty.hash_stable(hcx, hasher);
-            }
-            mir::StatementKind::Nop => {}
-            mir::StatementKind::InlineAsm { ref asm, ref outputs, ref inputs } => {
-                asm.hash_stable(hcx, hasher);
-                outputs.hash_stable(hcx, hasher);
-                inputs.hash_stable(hcx, hasher);
-            }
-        }
-    }
-}
+impl_stable_hash_for!(impl<'gcx> for enum mir::StatementKind<'gcx> [ mir::StatementKind ] {
+    Assign(place, rvalue),
+    FakeRead(cause, place),
+    SetDiscriminant { place, variant_index },
+    StorageLive(place),
+    StorageDead(place),
+    EscapeToRaw(place),
+    Retag { fn_entry, two_phase, place },
+    AscribeUserType(place, variance, c_ty),
+    Nop,
+    InlineAsm { asm, outputs, inputs },
+});
 
 impl_stable_hash_for!(enum mir::FakeReadCause { ForMatchGuard, ForMatchedPlace, ForLet });
 

src/librustc/mir/interpret/value.rs

@@ -138,6 +138,14 @@ impl<'tcx, Tag> Scalar<Tag> {
         }
     }
 
+    #[inline]
+    pub fn with_tag(self, new_tag: Tag) -> Self {
+        match self {
+            Scalar::Ptr(ptr) => Scalar::Ptr(Pointer { tag: new_tag, ..ptr }),
+            Scalar::Bits { bits, size } => Scalar::Bits { bits, size },
+        }
+    }
+
     #[inline]
     pub fn ptr_null(cx: &impl HasDataLayout) -> Self {
         Scalar::Bits {

src/librustc/mir/mod.rs

@@ -1778,6 +1778,10 @@ pub enum StatementKind<'tcx> {
         /// `fn_entry` indicates whether this is the initial retag that happens in the
         /// function prolog.
         fn_entry: bool,
+        /// `two_phase` indicates whether this is just the reservation action of
+        /// a two-phase borrow.
+        two_phase: bool,
+        /// The place to retag
         place: Place<'tcx>,
     },
 
@@ -1841,8 +1845,12 @@ impl<'tcx> Debug for Statement<'tcx> {
         match self.kind {
             Assign(ref place, ref rv) => write!(fmt, "{:?} = {:?}", place, rv),
             FakeRead(ref cause, ref place) => write!(fmt, "FakeRead({:?}, {:?})", cause, place),
-            Retag { fn_entry, ref place } =>
-                write!(fmt, "Retag({}{:?})", if fn_entry { "[fn entry] " } else { "" }, place),
+            Retag { fn_entry, two_phase, ref place } =>
+                write!(fmt, "Retag({}{}{:?})",
+                    if fn_entry { "[fn entry] " } else { "" },
+                    if two_phase { "[2phase] " } else { "" },
+                    place,
+                ),
             EscapeToRaw(ref place) => write!(fmt, "EscapeToRaw({:?})", place),
             StorageLive(ref place) => write!(fmt, "StorageLive({:?})", place),
             StorageDead(ref place) => write!(fmt, "StorageDead({:?})", place),
@@ -3019,7 +3027,7 @@ EnumTypeFoldableImpl! {
         (StatementKind::StorageLive)(a),
         (StatementKind::StorageDead)(a),
         (StatementKind::InlineAsm) { asm, outputs, inputs },
-        (StatementKind::Retag) { fn_entry, place },
+        (StatementKind::Retag) { fn_entry, two_phase, place },
         (StatementKind::EscapeToRaw)(place),
         (StatementKind::AscribeUserType)(a, v, b),
         (StatementKind::Nop),

src/librustc/mir/visit.rs

@@ -154,9 +154,10 @@ macro_rules! make_mir_visitor {
 
             fn visit_retag(&mut self,
                            fn_entry: & $($mutability)* bool,
+                           two_phase: & $($mutability)* bool,
                            place: & $($mutability)* Place<'tcx>,
                            location: Location) {
-                self.super_retag(fn_entry, place, location);
+                self.super_retag(fn_entry, two_phase, place, location);
             }
 
             fn visit_place(&mut self,
@@ -417,8 +418,9 @@ macro_rules! make_mir_visitor {
                         }
                     }
                     StatementKind::Retag { ref $($mutability)* fn_entry,
+                                           ref $($mutability)* two_phase,
                                            ref $($mutability)* place } => {
-                        self.visit_retag(fn_entry, place, location);
+                        self.visit_retag(fn_entry, two_phase, place, location);
                     }
                     StatementKind::AscribeUserType(
                         ref $($mutability)* place,
@@ -724,6 +726,7 @@ macro_rules! make_mir_visitor {
 
             fn super_retag(&mut self,
                            _fn_entry: & $($mutability)* bool,
+                           _two_phase: & $($mutability)* bool,
                            place: & $($mutability)* Place<'tcx>,
                            location: Location) {
                 self.visit_place(

src/librustc_mir/interpret/machine.rs

@@ -204,6 +204,7 @@ pub trait Machine<'a, 'mir, 'tcx>: Sized {
     fn retag(
         _ecx: &mut EvalContext<'a, 'mir, 'tcx, Self>,
         _fn_entry: bool,
+        _two_phase: bool,
         _place: PlaceTy<'tcx, Self::PointerTag>,
     ) -> EvalResult<'tcx> {
         Ok(())

src/librustc_mir/interpret/place.rs

@@ -115,6 +115,16 @@ impl<Tag> MemPlace<Tag> {
         }
     }
 
+    #[inline]
+    pub fn with_tag(self, new_tag: Tag) -> Self
+    {
+        MemPlace {
+            ptr: self.ptr.with_tag(new_tag),
+            align: self.align,
+            meta: self.meta,
+        }
+    }
+
     #[inline(always)]
     pub fn from_scalar_ptr(ptr: Scalar<Tag>, align: Align) -> Self {
         MemPlace {
@@ -187,6 +197,16 @@ impl<'tcx, Tag> MPlaceTy<'tcx, Tag> {
         }
     }
 
+    #[inline]
+    pub fn with_tag(self, new_tag: Tag) -> Self
+    {
+        MPlaceTy {
+            mplace: self.mplace.with_tag(new_tag),
+            layout: self.layout,
+        }
+    }
+
+    #[inline]
     pub fn offset(
         self,
         offset: Size,

src/librustc_mir/interpret/step.rs

@@ -119,9 +119,9 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M>
             FakeRead(..) => {}
 
             // Stacked Borrows.
-            Retag { fn_entry, ref place } => {
+            Retag { fn_entry, two_phase, ref place } => {
                 let dest = self.eval_place(place)?;
-                M::retag(self, fn_entry, dest)?;
+                M::retag(self, fn_entry, two_phase, dest)?;
             }
             EscapeToRaw(ref op) => {
                 let op = self.eval_operand(op, None)?;

src/librustc_mir/shim.rs

@@ -229,7 +229,11 @@ fn build_drop_shim<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
             // Function arguments should be retagged
             mir.basic_blocks_mut()[START_BLOCK].statements.insert(0, Statement {
                 source_info,
-                kind: StatementKind::Retag { fn_entry: true, place: dropee_ptr.clone() },
+                kind: StatementKind::Retag {
+                    fn_entry: true,
+                    two_phase: false,
+                    place: dropee_ptr.clone(),
+                },
             });
             // We use raw ptr operations, better prepare the alias tracking for that
             mir.basic_blocks_mut()[START_BLOCK].statements.insert(1, Statement {

src/librustc_mir/transform/add_retag.rs

@@ -118,7 +118,7 @@ impl MirPass for AddRetag {
             basic_blocks[START_BLOCK].statements.splice(0..0,
                 places.into_iter().map(|place| Statement {
                     source_info,
-                    kind: StatementKind::Retag { fn_entry: true, place },
+                    kind: StatementKind::Retag { fn_entry: true, two_phase: false, place },
                 })
             );
         }
@@ -154,7 +154,7 @@ impl MirPass for AddRetag {
         for (source_info, dest_place, dest_block) in returns {
             basic_blocks[dest_block].statements.insert(0, Statement {
                 source_info,
-                kind: StatementKind::Retag { fn_entry: false, place: dest_place },
+                kind: StatementKind::Retag { fn_entry: false, two_phase: false, place: dest_place },
             });
         }
 
@@ -191,12 +191,21 @@ impl MirPass for AddRetag {
                     // Assignments of reference or ptr type are the ones where we may have
                     // to update tags.  This includes `x = &[mut] ...` and hence
                     // we also retag after taking a reference!
-                    StatementKind::Assign(ref place, _) if needs_retag(place) => {
+                    StatementKind::Assign(ref place, box ref rvalue) if needs_retag(place) => {
+                        let two_phase = match rvalue {
+                            Rvalue::Ref(_, borrow_kind, _) =>
+                                borrow_kind.allows_two_phase_borrow(),
+                            _ => false
+                        };
                         // Insert a retag after the assignment.
                         let source_info = block_data.statements[i].source_info;
                         block_data.statements.insert(i+1, Statement {
                             source_info,
-                            kind: StatementKind::Retag { fn_entry: false, place: place.clone() },
+                            kind: StatementKind::Retag {
+                                fn_entry: false,
+                                two_phase,
+                                place: place.clone(),
+                            },
                         });
                     }
                     // Do nothing for the rest

src/librustc_mir/transform/inline.rs

@@ -707,8 +707,14 @@ impl<'a, 'tcx> MutVisitor<'tcx> for Integrator<'a, 'tcx> {
         self.in_cleanup_block = false;
     }
 
-    fn visit_retag(&mut self, fn_entry: &mut bool, place: &mut Place<'tcx>, loc: Location) {
-        self.super_retag(fn_entry, place, loc);
+    fn visit_retag(
+        &mut self,
+        fn_entry: &mut bool,
+        two_phase: &mut bool,
+        place: &mut Place<'tcx>,
+        loc: Location,
+    ) {
+        self.super_retag(fn_entry, two_phase, place, loc);
 
         // We have to patch all inlined retags to be aware that they are no longer
         // happening on function entry.