Rollup merge of #125633 - RalfJung:miri-no-copy, r=saethlin

miri: avoid making a full copy of all new allocations

Hopefully fixes https://github.com/rust-lang/miri/issues/3637

r? ``@saethlin``

5 files changed, 52 insertions, 26 deletions

diff --git a/compiler/rustc_const_eval/src/const_eval/dummy_machine.rs b/compiler/rustc_const_eval/src/const_eval/dummy_machine.rs
index 62979af8a60..9a98677a844 100644
--- a/compiler/rustc_const_eval/src/const_eval/dummy_machine.rs
+++ b/compiler/rustc_const_eval/src/const_eval/dummy_machine.rs
@@ -174,7 +174,7 @@ impl<'tcx> interpret::Machine<'tcx> for DummyMachine {
         unimplemented!()
     }
 
-    fn init_frame_extra(
+    fn init_frame(
         _ecx: &mut InterpCx<'tcx, Self>,
         _frame: interpret::Frame<'tcx, Self::Provenance>,
     ) -> interpret::InterpResult<'tcx, interpret::Frame<'tcx, Self::Provenance, Self::FrameExtra>>
diff --git a/compiler/rustc_const_eval/src/const_eval/machine.rs b/compiler/rustc_const_eval/src/const_eval/machine.rs
index d13ba4a2b95..79a161d3f03 100644
--- a/compiler/rustc_const_eval/src/const_eval/machine.rs
+++ b/compiler/rustc_const_eval/src/const_eval/machine.rs
@@ -643,7 +643,7 @@ impl<'tcx> interpret::Machine<'tcx> for CompileTimeInterpreter<'tcx> {
     }
 
     #[inline(always)]
-    fn init_frame_extra(
+    fn init_frame(
         ecx: &mut InterpCx<'tcx, Self>,
         frame: Frame<'tcx>,
     ) -> InterpResult<'tcx, Frame<'tcx>> {
diff --git a/compiler/rustc_const_eval/src/interpret/eval_context.rs b/compiler/rustc_const_eval/src/interpret/eval_context.rs
index 5cd50a92870..7c2100fcbe3 100644
--- a/compiler/rustc_const_eval/src/interpret/eval_context.rs
+++ b/compiler/rustc_const_eval/src/interpret/eval_context.rs
@@ -819,7 +819,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             tracing_span: SpanGuard::new(),
             extra: (),
         };
-        let frame = M::init_frame_extra(self, pre_frame)?;
+        let frame = M::init_frame(self, pre_frame)?;
         self.stack_mut().push(frame);
 
         // Make sure all the constants required by this frame evaluate successfully (post-monomorphization check).
diff --git a/compiler/rustc_const_eval/src/interpret/machine.rs b/compiler/rustc_const_eval/src/interpret/machine.rs
index 5713e7bd82b..4ae0aca5a0c 100644
--- a/compiler/rustc_const_eval/src/interpret/machine.rs
+++ b/compiler/rustc_const_eval/src/interpret/machine.rs
@@ -329,27 +329,41 @@ pub trait Machine<'tcx>: Sized {
         ptr: Pointer<Self::Provenance>,
     ) -> Option<(AllocId, Size, Self::ProvenanceExtra)>;
 
-    /// Called to adjust allocations to the Provenance and AllocExtra of this machine.
+    /// Called to adjust global allocations to the Provenance and AllocExtra of this machine.
     ///
     /// If `alloc` contains pointers, then they are all pointing to globals.
     ///
-    /// 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
-    /// globals, and allocations created by Miri during evaluation.
-    ///
-    /// `kind` is the kind of the allocation being adjusted; it can be `None` when
-    /// it's a global and `GLOBAL_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 adjust things), 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_allocation<'b>(
+    fn adjust_global_allocation<'b>(
         ecx: &InterpCx<'tcx, Self>,
         id: AllocId,
-        alloc: Cow<'b, Allocation>,
-        kind: Option<MemoryKind<Self::MemoryKind>>,
-    ) -> InterpResult<'tcx, Cow<'b, Allocation<Self::Provenance, Self::AllocExtra, Self::Bytes>>>;
+        alloc: &'b Allocation,
+    ) -> InterpResult<'tcx, Cow<'b, Allocation<Self::Provenance, Self::AllocExtra, Self::Bytes>>>
+    {
+        // The default implementation does a copy; CTFE machines have a more efficient implementation
+        // based on their particular choice for `Provenance`, `AllocExtra`, and `Bytes`.
+        let kind = Self::GLOBAL_KIND
+            .expect("if GLOBAL_KIND is None, adjust_global_allocation must be overwritten");
+        let alloc = alloc.adjust_from_tcx(&ecx.tcx, |ptr| ecx.global_root_pointer(ptr))?;
+        let extra =
+            Self::init_alloc_extra(ecx, id, MemoryKind::Machine(kind), alloc.size(), alloc.align)?;
+        Ok(Cow::Owned(alloc.with_extra(extra)))
+    }
+
+    /// Initialize the extra state of an allocation.
+    ///
+    /// This is guaranteed to be called exactly once on all allocations that are accessed by the
+    /// program.
+    fn init_alloc_extra(
+        ecx: &InterpCx<'tcx, Self>,
+        id: AllocId,
+        kind: MemoryKind<Self::MemoryKind>,
+        size: Size,
+        align: Align,
+    ) -> InterpResult<'tcx, Self::AllocExtra>;
 
     /// Return a "root" pointer for the given allocation: the one that is used for direct
     /// accesses to this static/const/fn allocation, or the one returned from the heap allocator.
@@ -473,7 +487,7 @@ pub trait Machine<'tcx>: Sized {
     }
 
     /// Called immediately before a new stack frame gets pushed.
-    fn init_frame_extra(
+    fn init_frame(
         ecx: &mut InterpCx<'tcx, Self>,
         frame: Frame<'tcx, Self::Provenance>,
     ) -> InterpResult<'tcx, Frame<'tcx, Self::Provenance, Self::FrameExtra>>;
@@ -590,13 +604,23 @@ pub macro compile_time_machine(<$tcx: lifetime>) {
     }
 
     #[inline(always)]
-    fn adjust_allocation<'b>(
+    fn adjust_global_allocation<'b>(
         _ecx: &InterpCx<$tcx, Self>,
         _id: AllocId,
-        alloc: Cow<'b, Allocation>,
-        _kind: Option<MemoryKind<Self::MemoryKind>>,
+        alloc: &'b Allocation,
     ) -> InterpResult<$tcx, Cow<'b, Allocation<Self::Provenance>>> {
-        Ok(alloc)
+        // Overwrite default implementation: no need to adjust anything.
+        Ok(Cow::Borrowed(alloc))
+    }
+
+    fn init_alloc_extra(
+        _ecx: &InterpCx<$tcx, Self>,
+        _id: AllocId,
+        _kind: MemoryKind<Self::MemoryKind>,
+        _size: Size,
+        _align: Align,
+    ) -> InterpResult<$tcx, Self::AllocExtra> {
+        Ok(())
     }
 
     fn extern_static_pointer(
diff --git a/compiler/rustc_const_eval/src/interpret/memory.rs b/compiler/rustc_const_eval/src/interpret/memory.rs
index 40bbfaa92c6..521f28b7123 100644
--- a/compiler/rustc_const_eval/src/interpret/memory.rs
+++ b/compiler/rustc_const_eval/src/interpret/memory.rs
@@ -239,7 +239,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
 
     pub fn allocate_raw_ptr(
         &mut self,
-        alloc: Allocation,
+        alloc: Allocation<M::Provenance, (), M::Bytes>,
         kind: MemoryKind<M::MemoryKind>,
     ) -> InterpResult<'tcx, Pointer<M::Provenance>> {
         let id = self.tcx.reserve_alloc_id();
@@ -248,8 +248,11 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             M::GLOBAL_KIND.map(MemoryKind::Machine),
             "dynamically allocating global memory"
         );
-        let alloc = M::adjust_allocation(self, id, Cow::Owned(alloc), Some(kind))?;
-        self.memory.alloc_map.insert(id, (kind, alloc.into_owned()));
+        // We have set things up so we don't need to call `adjust_from_tcx` here,
+        // so we avoid copying the entire allocation contents.
+        let extra = M::init_alloc_extra(self, id, kind, alloc.size(), alloc.align)?;
+        let alloc = alloc.with_extra(extra);
+        self.memory.alloc_map.insert(id, (kind, alloc));
         M::adjust_alloc_root_pointer(self, Pointer::from(id), Some(kind))
     }
 
@@ -583,11 +586,10 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         };
         M::before_access_global(self.tcx, &self.machine, id, alloc, def_id, is_write)?;
         // We got tcx memory. Let the machine initialize its "extra" stuff.
-        M::adjust_allocation(
+        M::adjust_global_allocation(
             self,
             id, // always use the ID we got as input, not the "hidden" one.
-            Cow::Borrowed(alloc.inner()),
-            M::GLOBAL_KIND.map(MemoryKind::Machine),
+            alloc.inner(),
         )
     }