Auto merge of #61278 - RalfJung:miri-tag-allocations, r=oli-obk

Miri: give machine the chance to tag all allocations

r? @oli-obk

The Miri side of this is at #61278.

src/librustc/mir/interpret/allocation.rs

@@ -111,9 +111,10 @@ pub trait AllocationExtra<Tag>: ::std::fmt::Debug + Clone {
 // For Tag=() and no extra state, we have is a trivial implementation.
 impl AllocationExtra<()> for () { }
 
-impl<Tag, Extra> Allocation<Tag, Extra> {
+// The constructors are all without extra; the extra gets added by a machine hook later.
+impl<Tag> Allocation<Tag> {
     /// Creates a read-only allocation initialized by the given bytes
-    pub fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, align: Align, extra: Extra) -> Self {
+    pub fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, align: Align) -> Self {
         let bytes = slice.into().into_owned();
         let undef_mask = UndefMask::new(Size::from_bytes(bytes.len() as u64), true);
         Self {
@@ -122,23 +123,23 @@ impl<Tag, Extra> Allocation<Tag, Extra> {
             undef_mask,
             align,
             mutability: Mutability::Immutable,
-            extra,
+            extra: (),
         }
     }
 
-    pub fn from_byte_aligned_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, extra: Extra) -> Self {
-        Allocation::from_bytes(slice, Align::from_bytes(1).unwrap(), extra)
+    pub fn from_byte_aligned_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>) -> Self {
+        Allocation::from_bytes(slice, Align::from_bytes(1).unwrap())
     }
 
-    pub fn undef(size: Size, align: Align, extra: Extra) -> Self {
+    pub fn undef(size: Size, align: Align) -> Self {
         assert_eq!(size.bytes() as usize as u64, size.bytes());
         Allocation {
             bytes: vec![0; size.bytes() as usize],
             relocations: Relocations::new(),
             undef_mask: UndefMask::new(size, false),
             align,
             mutability: Mutability::Mutable,
-            extra,
+            extra: (),
         }
     }
 }

src/librustc/mir/interpret/pointer.rs

@@ -116,13 +116,6 @@ impl<'tcx> Pointer<()> {
     {
         Pointer::new_with_tag(self.alloc_id, self.offset, tag)
     }
-
-    #[inline(always)]
-    pub fn with_default_tag<Tag>(self) -> Pointer<Tag>
-        where Tag: Default
-    {
-        self.with_tag(Tag::default())
-    }
 }
 
 impl<'tcx, Tag> Pointer<Tag> {

src/librustc/mir/interpret/value.rs

@@ -138,23 +138,22 @@ impl<'tcx> Scalar<()> {
                          "Scalar value {:#x} exceeds size of {} bytes", data, size);
     }
 
+    /// Tag this scalar with `new_tag` if it is a pointer, leave it unchanged otherwise.
+    ///
+    /// Used by `MemPlace::replace_tag`.
     #[inline]
     pub fn with_tag<Tag>(self, new_tag: Tag) -> Scalar<Tag> {
         match self {
             Scalar::Ptr(ptr) => Scalar::Ptr(ptr.with_tag(new_tag)),
             Scalar::Raw { data, size } => Scalar::Raw { data, size },
         }
     }
-
-    #[inline(always)]
-    pub fn with_default_tag<Tag>(self) -> Scalar<Tag>
-        where Tag: Default
-    {
-        self.with_tag(Tag::default())
-    }
 }
 
 impl<'tcx, Tag> Scalar<Tag> {
+    /// Erase the tag from the scalar, if any.
+    ///
+    /// Used by error reporting code to avoid having the error type depend on `Tag`.
     #[inline]
     pub fn erase_tag(self) -> Scalar {
         match self {
@@ -476,24 +475,10 @@ impl<Tag> fmt::Display for ScalarMaybeUndef<Tag> {
     }
 }
 
-impl<'tcx> ScalarMaybeUndef<()> {
-    #[inline]
-    pub fn with_tag<Tag>(self, new_tag: Tag) -> ScalarMaybeUndef<Tag> {
-        match self {
-            ScalarMaybeUndef::Scalar(s) => ScalarMaybeUndef::Scalar(s.with_tag(new_tag)),
-            ScalarMaybeUndef::Undef => ScalarMaybeUndef::Undef,
-        }
-    }
-
-    #[inline(always)]
-    pub fn with_default_tag<Tag>(self) -> ScalarMaybeUndef<Tag>
-        where Tag: Default
-    {
-        self.with_tag(Tag::default())
-    }
-}
-
 impl<'tcx, Tag> ScalarMaybeUndef<Tag> {
+    /// Erase the tag from the scalar, if any.
+    ///
+    /// Used by error reporting code to avoid having the error type depend on `Tag`.
     #[inline]
     pub fn erase_tag(self) -> ScalarMaybeUndef
     {

src/librustc/ty/context.rs

@@ -1149,7 +1149,7 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
     /// Allocates a byte or string literal for `mir::interpret`, read-only
     pub fn allocate_bytes(self, bytes: &[u8]) -> interpret::AllocId {
         // create an allocation that just contains these bytes
-        let alloc = interpret::Allocation::from_byte_aligned_bytes(bytes, ());
+        let alloc = interpret::Allocation::from_byte_aligned_bytes(bytes);
         let alloc = self.intern_const_alloc(alloc);
         self.alloc_map.lock().create_memory_alloc(alloc)
     }

src/librustc_mir/const_eval.rs

@@ -12,7 +12,7 @@ use rustc::hir::def_id::DefId;
 use rustc::mir::interpret::{ConstEvalErr, ErrorHandled};
 use rustc::mir;
 use rustc::ty::{self, TyCtxt, query::TyCtxtAt};
-use rustc::ty::layout::{self, LayoutOf, VariantIdx, Size};
+use rustc::ty::layout::{self, LayoutOf, VariantIdx};
 use rustc::ty::subst::Subst;
 use rustc::traits::Reveal;
 use rustc::util::common::ErrorReported;
@@ -117,7 +117,7 @@ fn op_to_const<'tcx>(
                 ),
                 Scalar::Raw { .. } => (
                     ecx.tcx.intern_const_alloc(Allocation::from_byte_aligned_bytes(
-                        b"" as &[u8], (),
+                        b"" as &[u8],
                     )),
                     0,
                 ),
@@ -397,27 +397,27 @@ impl<'a, 'mir, 'tcx> interpret::Machine<'a, 'mir, 'tcx>
     fn find_foreign_static(
         _def_id: DefId,
         _tcx: TyCtxtAt<'a, 'tcx, 'tcx>,
-        _memory_extra: &(),
     ) -> EvalResult<'tcx, Cow<'tcx, Allocation<Self::PointerTag>>> {
         err!(ReadForeignStatic)
     }
 
     #[inline(always)]
-    fn adjust_static_allocation<'b>(
-        alloc: &'b Allocation,
+    fn tag_allocation<'b>(
+        _id: AllocId,
+        alloc: Cow<'b, Allocation>,
+        _kind: Option<MemoryKind<!>>,
         _memory_extra: &(),
-    ) -> Cow<'b, Allocation<Self::PointerTag>> {
-        // We do not use a tag so we can just cheaply forward the reference
-        Cow::Borrowed(alloc)
+    ) -> (Cow<'b, Allocation<Self::PointerTag>>, Self::PointerTag) {
+        // We do not use a tag so we can just cheaply forward the allocation
+        (alloc, ())
     }
 
     #[inline(always)]
-    fn new_allocation(
-        _size: Size,
-        _extra: &Self::MemoryExtra,
-        _kind: MemoryKind<!>,
-    ) -> (Self::AllocExtra, Self::PointerTag) {
-        ((), ())
+    fn tag_static_base_pointer(
+        _id: AllocId,
+        _memory_extra: &(),
+    ) -> Self::PointerTag {
+        ()
     }
 
     fn box_alloc(

src/librustc_mir/hair/constant.rs

@@ -30,13 +30,13 @@ crate fn lit_to_const<'a, 'gcx, 'tcx>(
     let lit = match *lit {
         LitKind::Str(ref s, _) => {
             let s = s.as_str();
-            let allocation = Allocation::from_byte_aligned_bytes(s.as_bytes(), ());
+            let allocation = Allocation::from_byte_aligned_bytes(s.as_bytes());
             let allocation = tcx.intern_const_alloc(allocation);
             ConstValue::Slice { data: allocation, start: 0, end: s.len() }
         },
         LitKind::Err(ref s) => {
             let s = s.as_str();
-            let allocation = Allocation::from_byte_aligned_bytes(s.as_bytes(), ());
+            let allocation = Allocation::from_byte_aligned_bytes(s.as_bytes());
             let allocation = tcx.intern_const_alloc(allocation);
             return Ok(tcx.mk_const(ty::Const {
                 val: ConstValue::Slice{ data: allocation, start: 0, end: s.len() },

src/librustc_mir/interpret/cast.rs

@@ -86,7 +86,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> InterpretCx<'a, 'mir, 'tcx, M>
                             def_id,
                             substs,
                         ).ok_or_else(|| InterpError::TooGeneric.into());
-                        let fn_ptr = self.memory.create_fn_alloc(instance?).with_default_tag();
+                        let fn_ptr = self.memory.create_fn_alloc(instance?);
                         self.write_scalar(Scalar::Ptr(fn_ptr.into()), dest)?;
                     }
                     _ => bug!("reify fn pointer on {:?}", src.layout.ty),
@@ -115,7 +115,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> InterpretCx<'a, 'mir, 'tcx, M>
                             substs,
                             ty::ClosureKind::FnOnce,
                         );
-                        let fn_ptr = self.memory.create_fn_alloc(instance).with_default_tag();
+                        let fn_ptr = self.memory.create_fn_alloc(instance);
                         let val = Immediate::Scalar(Scalar::Ptr(fn_ptr.into()).into());
                         self.write_immediate(val, dest)?;
                     }

src/librustc_mir/interpret/eval_context.rs

@@ -15,7 +15,7 @@ use rustc::ty::query::TyCtxtAt;
 use rustc_data_structures::indexed_vec::IndexVec;
 use rustc::mir::interpret::{
     ErrorHandled,
-    GlobalId, Scalar, FrameInfo, AllocId,
+    GlobalId, Scalar, Pointer, FrameInfo, AllocId,
     EvalResult, InterpError,
     truncate, sign_extend,
 };
@@ -43,7 +43,10 @@ pub struct InterpretCx<'a, 'mir, 'tcx: 'a + 'mir, M: Machine<'a, 'mir, 'tcx>> {
     pub(crate) stack: Vec<Frame<'mir, 'tcx, M::PointerTag, M::FrameExtra>>,
 
     /// A cache for deduplicating vtables
-    pub(super) vtables: FxHashMap<(Ty<'tcx>, Option<ty::PolyExistentialTraitRef<'tcx>>), AllocId>,
+    pub(super) vtables: FxHashMap<
+        (Ty<'tcx>, Option<ty::PolyExistentialTraitRef<'tcx>>),
+        Pointer<M::PointerTag>
+    >,
 }
 
 /// A stack frame.
@@ -222,6 +225,11 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'a, 'mir, 'tcx>> InterpretCx<'a, 'mir, 'tc
         &mut self.memory
     }
 
+    #[inline(always)]
+    pub fn tag_static_base_pointer(&self, ptr: Pointer) -> Pointer<M::PointerTag> {
+        self.memory.tag_static_base_pointer(ptr)
+    }
+
     #[inline(always)]
     pub fn stack(&self) -> &[Frame<'mir, 'tcx, M::PointerTag, M::FrameExtra>] {
         &self.stack
@@ -360,11 +368,6 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'a, 'mir, 'tcx>> InterpretCx<'a, 'mir, 'tc
         }
     }
 
-    pub fn str_to_immediate(&mut self, s: &str) -> EvalResult<'tcx, Immediate<M::PointerTag>> {
-        let ptr = self.memory.allocate_static_bytes(s.as_bytes()).with_default_tag();
-        Ok(Immediate::new_slice(Scalar::Ptr(ptr), s.len() as u64, self))
-    }
-
     /// Returns the actual dynamic size and alignment of the place at the given type.
     /// Only the "meta" (metadata) part of the place matters.
     /// This can fail to provide an answer for extern types.

src/librustc_mir/interpret/intrinsics/type_name.rs

@@ -215,7 +215,7 @@ impl Write for AbsolutePathPrinter<'_, '_> {
 pub fn type_name<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, ty: Ty<'tcx>) -> &'tcx ty::Const<'tcx> {
     let path = AbsolutePathPrinter { tcx, path: String::new() }.print_type(ty).unwrap().path;
     let len = path.len();
-    let alloc = Allocation::from_byte_aligned_bytes(path.into_bytes(), ());
+    let alloc = Allocation::from_byte_aligned_bytes(path.into_bytes());
     let alloc = tcx.intern_const_alloc(alloc);
     tcx.mk_const(ty::Const {
         val: ConstValue::Slice {

src/librustc_mir/interpret/machine.rs

@@ -5,13 +5,13 @@
 use std::borrow::{Borrow, Cow};
 use std::hash::Hash;
 
-use rustc::hir::{self, def_id::DefId};
+use rustc::hir::def_id::DefId;
 use rustc::mir;
-use rustc::ty::{self, query::TyCtxtAt, layout::Size};
+use rustc::ty::{self, query::TyCtxtAt};
 
 use super::{
     Allocation, AllocId, EvalResult, Scalar, AllocationExtra,
-    InterpretCx, PlaceTy, MPlaceTy, OpTy, ImmTy, MemoryKind,
+    InterpretCx, PlaceTy, OpTy, ImmTy, MemoryKind,
 };
 
 /// Whether this kind of memory is allowed to leak
@@ -65,7 +65,7 @@ pub trait Machine<'a, 'mir, 'tcx>: Sized {
     /// Tag tracked alongside every pointer. This is used to implement "Stacked Borrows"
     /// <https://www.ralfj.de/blog/2018/08/07/stacked-borrows.html>.
     /// The `default()` is used for pointers to consts, statics, vtables and functions.
-    type PointerTag: ::std::fmt::Debug + Default + Copy + Eq + Hash + 'static;
+    type PointerTag: ::std::fmt::Debug + Copy + Eq + Hash + 'static;
 
     /// Extra data stored in every call frame.
     type FrameExtra;
@@ -90,7 +90,7 @@ pub trait Machine<'a, 'mir, 'tcx>: Sized {
     /// The memory kind to use for copied statics -- or None if statics should not be mutated
     /// and thus any such attempt will cause a `ModifiedStatic` error to be raised.
     /// Statics are copied under two circumstances: When they are mutated, and when
-    /// `static_with_default_tag` or `find_foreign_static` (see below) returns an owned allocation
+    /// `tag_allocation` or `find_foreign_static` (see below) returns an owned allocation
     /// that is added to the memory so that the work is not done twice.
     const STATIC_KIND: Option<Self::MemoryKinds>;
 
@@ -133,11 +133,12 @@ pub trait Machine<'a, 'mir, 'tcx>: Sized {
     /// This will only be called once per static and machine; the result is cached in
     /// the machine memory. (This relies on `AllocMap::get_or` being able to add the
     /// owned allocation to the map even when the map is shared.)
+    ///
+    /// This allocation will then be fed to `tag_allocation` to initialize the "extra" state.
     fn find_foreign_static(
         def_id: DefId,
         tcx: TyCtxtAt<'a, 'tcx, 'tcx>,
-        memory_extra: &Self::MemoryExtra,
-    ) -> EvalResult<'tcx, Cow<'tcx, Allocation<Self::PointerTag, Self::AllocExtra>>>;
+    ) -> EvalResult<'tcx, Cow<'tcx, Allocation>>;
 
     /// Called for all binary operations on integer(-like) types when one operand is a pointer
     /// value, and for the `Offset` operation that is inherently about pointers.
@@ -156,36 +157,38 @@ pub trait Machine<'a, 'mir, 'tcx>: Sized {
         dest: PlaceTy<'tcx, Self::PointerTag>,
     ) -> EvalResult<'tcx>;
 
-    /// Called to turn an allocation obtained from the `tcx` into one that has
-    /// the right type for this machine.
+    /// Called to initialize the "extra" state of an allocation and make the pointers
+    /// it contains (in relocations) tagged.  The way we construct allocations is
+    /// to always first construct it without extra and then add the extra.
+    /// This keeps uniform code paths for handling both allocations created by CTFE
+    /// for statics, and allocations ceated by Miri during evaluation.
+    ///
+    /// `kind` is the kind of the allocation being tagged; it can be `None` when
+    /// it's a static and `STATIC_KIND` is `None`.
     ///
     /// This should avoid copying if no work has to be done! If this returns an owned
     /// allocation (because a copy had to be done to add tags or metadata), machine memory will
     /// cache the result. (This relies on `AllocMap::get_or` being able to add the
     /// owned allocation to the map even when the map is shared.)
-    fn adjust_static_allocation<'b>(
-        alloc: &'b Allocation,
+    ///
+    /// For static allocations, the tag returned must be the same as the one returned by
+    /// `tag_static_base_pointer`.
+    fn tag_allocation<'b>(
+        id: AllocId,
+        alloc: Cow<'b, Allocation>,
+        kind: Option<MemoryKind<Self::MemoryKinds>>,
         memory_extra: &Self::MemoryExtra,
-    ) -> Cow<'b, Allocation<Self::PointerTag, Self::AllocExtra>>;
-
-    /// Computes the extra state and the tag for a new allocation.
-    fn new_allocation(
-        size: Size,
-        extra: &Self::MemoryExtra,
-        kind: MemoryKind<Self::MemoryKinds>,
-    ) -> (Self::AllocExtra, Self::PointerTag);
-
-    /// Executed when evaluating the `*` operator: Following a reference.
-    /// This has the chance to adjust the tag. It should not change anything else!
-    /// `mutability` can be `None` in case a raw ptr is being dereferenced.
-    #[inline]
-    fn tag_dereference(
-        _ecx: &InterpretCx<'a, 'mir, 'tcx, Self>,
-        place: MPlaceTy<'tcx, Self::PointerTag>,
-        _mutability: Option<hir::Mutability>,
-    ) -> EvalResult<'tcx, Scalar<Self::PointerTag>> {
-        Ok(place.ptr)
-    }
+    ) -> (Cow<'b, Allocation<Self::PointerTag, Self::AllocExtra>>, Self::PointerTag);
+
+    /// Return the "base" tag for the given static allocation: the one that is used for direct
+    /// accesses to this static/const/fn allocation.
+    ///
+    /// Be aware that requesting the `Allocation` for that `id` will lead to cycles
+    /// for cyclic statics!
+    fn tag_static_base_pointer(
+        id: AllocId,
+        memory_extra: &Self::MemoryExtra,
+    ) -> Self::PointerTag;
 
     /// Executes a retagging operation
     #[inline]