Auto merge of #56094 - RalfJung:memory-data-revived, r=oli-obk

miri: Memory data revived, Hooks for stack frame push/pop

r? @oli-obk

Author: bors

src/librustc/mir/interpret/allocation.rs

@@ -1055,7 +1055,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().allocate(alloc)
     }

src/librustc/ty/context.rs

@@ -65,6 +65,7 @@ pub fn mk_borrowck_eval_cx<'a, 'mir, 'tcx>(
         return_place: None,
         return_to_block: StackPopCleanup::Goto(None), // never pop
         stmt: 0,
+        extra: (),
     });
     Ok(ecx)
 }
@@ -353,9 +354,12 @@ impl<'a, 'mir, 'tcx> interpret::Machine<'a, 'mir, 'tcx>
     for CompileTimeInterpreter<'a, 'mir, 'tcx>
 {
     type MemoryKinds = !;
-    type AllocExtra = ();
     type PointerTag = ();
 
+    type FrameExtra = ();
+    type MemoryExtra = ();
+    type AllocExtra = ();
+
     type MemoryMap = FxHashMap<AllocId, (MemoryKind<!>, Allocation)>;
 
     const STATIC_KIND: Option<!> = None; // no copying of statics allowed
@@ -432,16 +436,18 @@ impl<'a, 'mir, 'tcx> interpret::Machine<'a, 'mir, 'tcx>
     }
 
     fn find_foreign_static(
-        _tcx: TyCtxtAt<'a, 'tcx, 'tcx>,
         _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(
-        alloc: &'_ Allocation
-    ) -> Cow<'_, Allocation<Self::PointerTag>> {
+    fn adjust_static_allocation<'b>(
+        alloc: &'b Allocation,
+        _memory_extra: &(),
+    ) -> Cow<'b, Allocation<Self::PointerTag>> {
         // We do not use a tag so we can just cheaply forward the reference
         Cow::Borrowed(alloc)
     }
@@ -487,6 +493,22 @@ impl<'a, 'mir, 'tcx> interpret::Machine<'a, 'mir, 'tcx>
     ) -> EvalResult<'tcx, Pointer> {
         Ok(ptr)
     }
+
+    #[inline(always)]
+    fn stack_push(
+        _ecx: &mut EvalContext<'a, 'mir, 'tcx, Self>,
+    ) -> EvalResult<'tcx> {
+        Ok(())
+    }
+
+    /// Called immediately before a stack frame gets popped
+    #[inline(always)]
+    fn stack_pop(
+        _ecx: &mut EvalContext<'a, 'mir, 'tcx, Self>,
+        _extra: (),
+    ) -> EvalResult<'tcx> {
+        Ok(())
+    }
 }
 
 /// Project to a field of a (variant of a) const

src/librustc_mir/const_eval.rs

@@ -49,15 +49,15 @@ pub struct EvalContext<'a, 'mir, 'tcx: 'a + 'mir, M: Machine<'a, 'mir, 'tcx>> {
     pub(crate) memory: Memory<'a, 'mir, 'tcx, M>,
 
     /// The virtual call stack.
-    pub(crate) stack: Vec<Frame<'mir, 'tcx, M::PointerTag>>,
+    pub(crate) stack: Vec<Frame<'mir, 'tcx, M::PointerTag, M::FrameExtra>>,
 
     /// A cache for deduplicating vtables
     pub(super) vtables: FxHashMap<(Ty<'tcx>, ty::PolyExistentialTraitRef<'tcx>), AllocId>,
 }
 
 /// A stack frame.
 #[derive(Clone)]
-pub struct Frame<'mir, 'tcx: 'mir, Tag=()> {
+pub struct Frame<'mir, 'tcx: 'mir, Tag=(), Extra=()> {
     ////////////////////////////////////////////////////////////////////////////////
     // Function and callsite information
     ////////////////////////////////////////////////////////////////////////////////
@@ -96,6 +96,9 @@ pub struct Frame<'mir, 'tcx: 'mir, Tag=()> {
 
     /// The index of the currently evaluated statement.
     pub stmt: usize,
+
+    /// Extra data for the machine
+    pub extra: Extra,
 }
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
@@ -196,7 +199,7 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'a, 'mir, 'tcx>> EvalContext<'a, 'mir, 'tc
     }
 
     #[inline(always)]
-    pub fn stack(&self) -> &[Frame<'mir, 'tcx, M::PointerTag>] {
+    pub fn stack(&self) -> &[Frame<'mir, 'tcx, M::PointerTag, M::FrameExtra>] {
         &self.stack
     }
 
@@ -207,12 +210,12 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'a, 'mir, 'tcx>> EvalContext<'a, 'mir, 'tc
     }
 
     #[inline(always)]
-    pub fn frame(&self) -> &Frame<'mir, 'tcx, M::PointerTag> {
+    pub fn frame(&self) -> &Frame<'mir, 'tcx, M::PointerTag, M::FrameExtra> {
         self.stack.last().expect("no call frames exist")
     }
 
     #[inline(always)]
-    pub fn frame_mut(&mut self) -> &mut Frame<'mir, 'tcx, M::PointerTag> {
+    pub fn frame_mut(&mut self) -> &mut Frame<'mir, 'tcx, M::PointerTag, M::FrameExtra> {
         self.stack.last_mut().expect("no call frames exist")
     }
 
@@ -294,7 +297,7 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'a, 'mir, 'tcx>> EvalContext<'a, 'mir, 'tc
 
     pub fn layout_of_local(
         &self,
-        frame: &Frame<'mir, 'tcx, M::PointerTag>,
+        frame: &Frame<'mir, 'tcx, M::PointerTag, M::FrameExtra>,
         local: mir::Local
     ) -> EvalResult<'tcx, TyLayout<'tcx>> {
         let local_ty = frame.mir.local_decls[local].ty;
@@ -424,6 +427,7 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'a, 'mir, 'tcx>> EvalContext<'a, 'mir, 'tc
         ::log_settings::settings().indentation += 1;
 
         // first push a stack frame so we have access to the local substs
+        let extra = M::stack_push(self)?;
         self.stack.push(Frame {
             mir,
             block: mir::START_BLOCK,
@@ -435,6 +439,7 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'a, 'mir, 'tcx>> EvalContext<'a, 'mir, 'tc
             span,
             instance,
             stmt: 0,
+            extra,
         });
 
         // don't allocate at all for trivial constants
@@ -504,6 +509,7 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'a, 'mir, 'tcx>> EvalContext<'a, 'mir, 'tc
         let frame = self.stack.pop().expect(
             "tried to pop a stack frame, but there were none",
         );
+        M::stack_pop(self, frame.extra)?;
         // Abort early if we do not want to clean up: We also avoid validation in that case,
         // because this is CTFE and the final value will be thoroughly validated anyway.
         match frame.return_to_block {

src/librustc_mir/interpret/eval_context.rs

@@ -77,8 +77,16 @@ pub trait Machine<'a, 'mir, 'tcx>: Sized {
     /// The `default()` is used for pointers to consts, statics, vtables and functions.
     type PointerTag: ::std::fmt::Debug + Default + Copy + Eq + Hash + 'static;
 
+    /// Extra data stored in every call frame.
+    type FrameExtra;
+
+    /// Extra data stored in memory.  A reference to this is available when `AllocExtra`
+    /// gets initialized, so you can e.g. have an `Rc` here if there is global state you
+    /// need access to in the `AllocExtra` hooks.
+    type MemoryExtra: Default;
+
     /// Extra data stored in every allocation.
-    type AllocExtra: AllocationExtra<Self::PointerTag>;
+    type AllocExtra: AllocationExtra<Self::PointerTag, Self::MemoryExtra>;
 
     /// Memory's allocation map
     type MemoryMap:
@@ -136,8 +144,9 @@ pub trait Machine<'a, 'mir, 'tcx>: Sized {
     /// 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.)
     fn find_foreign_static(
-        tcx: TyCtxtAt<'a, 'tcx, 'tcx>,
         def_id: DefId,
+        tcx: TyCtxtAt<'a, 'tcx, 'tcx>,
+        memory_extra: &Self::MemoryExtra,
     ) -> EvalResult<'tcx, Cow<'tcx, Allocation<Self::PointerTag, Self::AllocExtra>>>;
 
     /// Called to turn an allocation obtained from the `tcx` into one that has
@@ -147,9 +156,10 @@ pub trait Machine<'a, 'mir, 'tcx>: Sized {
     /// 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(
-        alloc: &'_ Allocation
-    ) -> Cow<'_, Allocation<Self::PointerTag, Self::AllocExtra>>;
+    fn adjust_static_allocation<'b>(
+        alloc: &'b Allocation,
+        memory_extra: &Self::MemoryExtra,
+    ) -> Cow<'b, Allocation<Self::PointerTag, Self::AllocExtra>>;
 
     /// 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.
@@ -207,4 +217,15 @@ pub trait Machine<'a, 'mir, 'tcx>: Sized {
     ) -> EvalResult<'tcx> {
         Ok(())
     }
+
+    /// Called immediately before a new stack frame got pushed
+    fn stack_push(
+        ecx: &mut EvalContext<'a, 'mir, 'tcx, Self>,
+    ) -> EvalResult<'tcx, Self::FrameExtra>;
+
+    /// Called immediately after a stack frame gets popped
+    fn stack_pop(
+        ecx: &mut EvalContext<'a, 'mir, 'tcx, Self>,
+        extra: Self::FrameExtra,
+    ) -> EvalResult<'tcx>;
 }

src/librustc_mir/interpret/machine.rs

@@ -73,6 +73,9 @@ pub struct Memory<'a, 'mir, 'tcx: 'a + 'mir, M: Machine<'a, 'mir, 'tcx>> {
     /// that do not exist any more.
     dead_alloc_map: FxHashMap<AllocId, (Size, Align)>,
 
+    /// Extra data added by the machine.
+    pub extra: M::MemoryExtra,
+
     /// Lets us implement `HasDataLayout`, which is awfully convenient.
     pub(super) tcx: TyCtxtAt<'a, 'tcx, 'tcx>,
 }
@@ -88,13 +91,19 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> HasDataLayout
 
 // FIXME: Really we shouldn't clone memory, ever. Snapshot machinery should instead
 // carefully copy only the reachable parts.
-impl<'a, 'mir, 'tcx: 'a + 'mir, M: Machine<'a, 'mir, 'tcx>>
-    Clone for Memory<'a, 'mir, 'tcx, M>
+impl<'a, 'mir, 'tcx, M>
+    Clone
+for
+    Memory<'a, 'mir, 'tcx, M>
+where
+    M: Machine<'a, 'mir, 'tcx, PointerTag=(), AllocExtra=(), MemoryExtra=()>,
+    M::MemoryMap: AllocMap<AllocId, (MemoryKind<M::MemoryKinds>, Allocation)>,
 {
     fn clone(&self) -> Self {
         Memory {
             alloc_map: self.alloc_map.clone(),
             dead_alloc_map: self.dead_alloc_map.clone(),
+            extra: (),
             tcx: self.tcx,
         }
     }
@@ -103,8 +112,9 @@ impl<'a, 'mir, 'tcx: 'a + 'mir, M: Machine<'a, 'mir, 'tcx>>
 impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
     pub fn new(tcx: TyCtxtAt<'a, 'tcx, 'tcx>) -> Self {
         Memory {
-            alloc_map: Default::default(),
+            alloc_map: M::MemoryMap::default(),
             dead_alloc_map: FxHashMap::default(),
+            extra: M::MemoryExtra::default(),
             tcx,
         }
     }
@@ -133,7 +143,8 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
         align: Align,
         kind: MemoryKind<M::MemoryKinds>,
     ) -> EvalResult<'tcx, Pointer> {
-        Ok(Pointer::from(self.allocate_with(Allocation::undef(size, align), kind)?))
+        let extra = AllocationExtra::memory_allocated(size, &self.extra);
+        Ok(Pointer::from(self.allocate_with(Allocation::undef(size, align, extra), kind)?))
     }
 
     pub fn reallocate(
@@ -309,15 +320,16 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
     /// this machine use the same pointer tag, so it is indirected through
     /// `M::static_with_default_tag`.
     fn get_static_alloc(
-        tcx: TyCtxtAt<'a, 'tcx, 'tcx>,
         id: AllocId,
+        tcx: TyCtxtAt<'a, 'tcx, 'tcx>,
+        memory_extra: &M::MemoryExtra,
     ) -> EvalResult<'tcx, Cow<'tcx, Allocation<M::PointerTag, M::AllocExtra>>> {
         let alloc = tcx.alloc_map.lock().get(id);
         let def_id = match alloc {
             Some(AllocType::Memory(mem)) => {
                 // We got tcx memory. Let the machine figure out whether and how to
                 // turn that into memory with the right pointer tag.
-                return Ok(M::adjust_static_allocation(mem))
+                return Ok(M::adjust_static_allocation(mem, memory_extra))
             }
             Some(AllocType::Function(..)) => {
                 return err!(DerefFunctionPointer)
@@ -331,7 +343,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
         // We got a "lazy" static that has not been computed yet, do some work
         trace!("static_alloc: Need to compute {:?}", def_id);
         if tcx.is_foreign_item(def_id) {
-            return M::find_foreign_static(tcx, def_id);
+            return M::find_foreign_static(def_id, tcx, memory_extra);
         }
         let instance = Instance::mono(tcx.tcx, def_id);
         let gid = GlobalId {
@@ -351,7 +363,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
             let allocation = tcx.alloc_map.lock().unwrap_memory(raw_const.alloc_id);
             // We got tcx memory. Let the machine figure out whether and how to
             // turn that into memory with the right pointer tag.
-            M::adjust_static_allocation(allocation)
+            M::adjust_static_allocation(allocation, memory_extra)
         })
     }
 
@@ -361,7 +373,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
         // `get_static_alloc` that we can actually use directly without inserting anything anywhere.
         // So the error type is `EvalResult<'tcx, &Allocation<M::PointerTag>>`.
         let a = self.alloc_map.get_or(id, || {
-            let alloc = Self::get_static_alloc(self.tcx, id).map_err(Err)?;
+            let alloc = Self::get_static_alloc(id, self.tcx, &self.extra).map_err(Err)?;
             match alloc {
                 Cow::Borrowed(alloc) => {
                     // We got a ref, cheaply return that as an "error" so that the
@@ -390,10 +402,11 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
         id: AllocId,
     ) -> EvalResult<'tcx, &mut Allocation<M::PointerTag, M::AllocExtra>> {
         let tcx = self.tcx;
+        let memory_extra = &self.extra;
         let a = self.alloc_map.get_mut_or(id, || {
             // Need to make a copy, even if `get_static_alloc` is able
             // to give us a cheap reference.
-            let alloc = Self::get_static_alloc(tcx, id)?;
+            let alloc = Self::get_static_alloc(id, tcx, memory_extra)?;
             if alloc.mutability == Mutability::Immutable {
                 return err!(ModifiedConstantMemory);
             }
@@ -601,7 +614,8 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
 /// Interning (for CTFE)
 impl<'a, 'mir, 'tcx, M> Memory<'a, 'mir, 'tcx, M>
 where
-    M: Machine<'a, 'mir, 'tcx, PointerTag=(), AllocExtra=()>,
+    M: Machine<'a, 'mir, 'tcx, PointerTag=(), AllocExtra=(), MemoryExtra=()>,
+    // FIXME: Working around https://github.com/rust-lang/rust/issues/24159
     M::MemoryMap: AllocMap<AllocId, (MemoryKind<M::MemoryKinds>, Allocation)>,
 {
     /// mark an allocation as static and initialized, either mutable or not

src/librustc_mir/interpret/memory.rs

@@ -471,7 +471,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M>
     /// When you know the layout of the local in advance, you can pass it as last argument
     pub fn access_local(
         &self,
-        frame: &super::Frame<'mir, 'tcx, M::PointerTag>,
+        frame: &super::Frame<'mir, 'tcx, M::PointerTag, M::FrameExtra>,
         local: mir::Local,
         layout: Option<TyLayout<'tcx>>,
     ) -> EvalResult<'tcx, OpTy<'tcx, M::PointerTag>> {

src/librustc_mir/interpret/operand.rs

@@ -295,10 +295,12 @@ impl<'tcx, Tag: ::std::fmt::Debug> PlaceTy<'tcx, Tag> {
 // separating the pointer tag for `impl Trait`, see https://github.com/rust-lang/rust/issues/54385
 impl<'a, 'mir, 'tcx, Tag, M> EvalContext<'a, 'mir, 'tcx, M>
 where
+    // FIXME: Working around https://github.com/rust-lang/rust/issues/54385
     Tag: ::std::fmt::Debug+Default+Copy+Eq+Hash+'static,
     M: Machine<'a, 'mir, 'tcx, PointerTag=Tag>,
+    // FIXME: Working around https://github.com/rust-lang/rust/issues/24159
     M::MemoryMap: AllocMap<AllocId, (MemoryKind<M::MemoryKinds>, Allocation<Tag, M::AllocExtra>)>,
-    M::AllocExtra: AllocationExtra<Tag>,
+    M::AllocExtra: AllocationExtra<Tag, M::MemoryExtra>,
 {
     /// Take a value, which represents a (thin or fat) reference, and make it a place.
     /// Alignment is just based on the type.  This is the inverse of `MemPlace::to_ref()`.

src/librustc_mir/interpret/place.rs

@@ -323,6 +323,7 @@ impl_stable_hash_for!(impl<'tcx, 'mir: 'tcx> for struct Frame<'mir, 'tcx> {
     locals,
     block,
     stmt,
+    extra,
 });
 
 impl<'a, 'mir, 'tcx, Ctx> Snapshot<'a, Ctx> for &'a Frame<'mir, 'tcx>
@@ -340,6 +341,7 @@ impl<'a, 'mir, 'tcx, Ctx> Snapshot<'a, Ctx> for &'a Frame<'mir, 'tcx>
             locals,
             block,
             stmt,
+            extra: _,
         } = self;
 
         FrameSnapshot {