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``

8 files changed, 92 insertions, 85 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(),
         )
     }
 
diff --git a/compiler/rustc_middle/src/mir/interpret/allocation.rs b/compiler/rustc_middle/src/mir/interpret/allocation.rs
index 791e87735f4..b0f8a047b82 100644
--- a/compiler/rustc_middle/src/mir/interpret/allocation.rs
+++ b/compiler/rustc_middle/src/mir/interpret/allocation.rs
@@ -266,19 +266,6 @@ impl AllocRange {
 
 // The constructors are all without extra; the extra gets added by a machine hook later.
 impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
-    /// Creates an allocation from an existing `Bytes` value - this is needed for miri FFI support
-    pub fn from_raw_bytes(bytes: Bytes, align: Align, mutability: Mutability) -> Self {
-        let size = Size::from_bytes(bytes.len());
-        Self {
-            bytes,
-            provenance: ProvenanceMap::new(),
-            init_mask: InitMask::new(size, true),
-            align,
-            mutability,
-            extra: (),
-        }
-    }
-
     /// Creates an allocation initialized by the given bytes
     pub fn from_bytes<'a>(
         slice: impl Into<Cow<'a, [u8]>>,
@@ -342,18 +329,30 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
             Err(x) => x,
         }
     }
+
+    /// Add the extra.
+    pub fn with_extra<Extra>(self, extra: Extra) -> Allocation<Prov, Extra, Bytes> {
+        Allocation {
+            bytes: self.bytes,
+            provenance: self.provenance,
+            init_mask: self.init_mask,
+            align: self.align,
+            mutability: self.mutability,
+            extra,
+        }
+    }
 }
 
 impl Allocation {
     /// Adjust allocation from the ones in `tcx` to a custom Machine instance
-    /// with a different `Provenance`, `Extra` and `Byte` type.
-    pub fn adjust_from_tcx<Prov: Provenance, Extra, Bytes: AllocBytes, Err>(
-        self,
+    /// with a different `Provenance` and `Byte` type.
+    pub fn adjust_from_tcx<Prov: Provenance, Bytes: AllocBytes, Err>(
+        &self,
         cx: &impl HasDataLayout,
-        extra: Extra,
         mut adjust_ptr: impl FnMut(Pointer<CtfeProvenance>) -> Result<Pointer<Prov>, Err>,
-    ) -> Result<Allocation<Prov, Extra, Bytes>, Err> {
-        let mut bytes = self.bytes;
+    ) -> Result<Allocation<Prov, (), Bytes>, Err> {
+        // Copy the data.
+        let mut bytes = Bytes::from_bytes(Cow::Borrowed(&*self.bytes), self.align);
         // Adjust provenance of pointers stored in this allocation.
         let mut new_provenance = Vec::with_capacity(self.provenance.ptrs().len());
         let ptr_size = cx.data_layout().pointer_size.bytes_usize();
@@ -369,12 +368,12 @@ impl Allocation {
         }
         // Create allocation.
         Ok(Allocation {
-            bytes: AllocBytes::from_bytes(Cow::Owned(Vec::from(bytes)), self.align),
+            bytes,
             provenance: ProvenanceMap::from_presorted_ptrs(new_provenance),
-            init_mask: self.init_mask,
+            init_mask: self.init_mask.clone(),
             align: self.align,
             mutability: self.mutability,
-            extra,
+            extra: self.extra,
         })
     }
 }
diff --git a/src/tools/miri/src/concurrency/thread.rs b/src/tools/miri/src/concurrency/thread.rs
index 8e8e1182bff..06af72dadbd 100644
--- a/src/tools/miri/src/concurrency/thread.rs
+++ b/src/tools/miri/src/concurrency/thread.rs
@@ -862,14 +862,15 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             if tcx.is_foreign_item(def_id) {
                 throw_unsup_format!("foreign thread-local statics are not supported");
             }
-            let allocation = this.ctfe_query(|tcx| tcx.eval_static_initializer(def_id))?;
-            let mut allocation = allocation.inner().clone();
+            let alloc = this.ctfe_query(|tcx| tcx.eval_static_initializer(def_id))?;
+            // We make a full copy of this allocation.
+            let mut alloc = alloc.inner().adjust_from_tcx(&this.tcx, |ptr| this.global_root_pointer(ptr))?;
             // This allocation will be deallocated when the thread dies, so it is not in read-only memory.
-            allocation.mutability = Mutability::Mut;
+            alloc.mutability = Mutability::Mut;
             // Create a fresh allocation with this content.
-            let new_alloc = this.allocate_raw_ptr(allocation, MiriMemoryKind::Tls.into())?;
-            this.machine.threads.set_thread_local_alloc(def_id, new_alloc);
-            Ok(new_alloc)
+            let ptr = this.allocate_raw_ptr(alloc, MiriMemoryKind::Tls.into())?;
+            this.machine.threads.set_thread_local_alloc(def_id, ptr);
+            Ok(ptr)
         }
     }
 
diff --git a/src/tools/miri/src/machine.rs b/src/tools/miri/src/machine.rs
index fc1a655c123..c9b390eb67e 100644
--- a/src/tools/miri/src/machine.rs
+++ b/src/tools/miri/src/machine.rs
@@ -1,7 +1,6 @@
 //! Global machine state as well as implementation of the interpreter engine
 //! `Machine` trait.
 
-use std::borrow::Cow;
 use std::cell::RefCell;
 use std::collections::hash_map::Entry;
 use std::fmt;
@@ -1086,40 +1085,33 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> {
         }
     }
 
-    fn adjust_allocation<'b>(
+    fn init_alloc_extra(
         ecx: &MiriInterpCx<'tcx>,
         id: AllocId,
-        alloc: Cow<'b, Allocation>,
-        kind: Option<MemoryKind>,
-    ) -> InterpResult<'tcx, Cow<'b, Allocation<Self::Provenance, Self::AllocExtra, Self::Bytes>>>
-    {
-        let kind = kind.expect("we set our STATIC_KIND so this cannot be None");
+        kind: MemoryKind,
+        size: Size,
+        align: Align,
+    ) -> InterpResult<'tcx, Self::AllocExtra> {
         if ecx.machine.tracked_alloc_ids.contains(&id) {
-            ecx.emit_diagnostic(NonHaltingDiagnostic::CreatedAlloc(
-                id,
-                alloc.size(),
-                alloc.align,
-                kind,
-            ));
+            ecx.emit_diagnostic(NonHaltingDiagnostic::CreatedAlloc(id, size, align, kind));
         }
 
-        let alloc = alloc.into_owned();
         let borrow_tracker = ecx
             .machine
             .borrow_tracker
             .as_ref()
-            .map(|bt| bt.borrow_mut().new_allocation(id, alloc.size(), kind, &ecx.machine));
+            .map(|bt| bt.borrow_mut().new_allocation(id, size, kind, &ecx.machine));
 
-        let race_alloc = ecx.machine.data_race.as_ref().map(|data_race| {
+        let data_race = ecx.machine.data_race.as_ref().map(|data_race| {
             data_race::AllocState::new_allocation(
                 data_race,
                 &ecx.machine.threads,
-                alloc.size(),
+                size,
                 kind,
                 ecx.machine.current_span(),
             )
         });
-        let buffer_alloc = ecx.machine.weak_memory.then(weak_memory::AllocState::new_allocation);
+        let weak_memory = ecx.machine.weak_memory.then(weak_memory::AllocState::new_allocation);
 
         // If an allocation is leaked, we want to report a backtrace to indicate where it was
         // allocated. We don't need to record a backtrace for allocations which are allowed to
@@ -1130,17 +1122,6 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> {
             Some(ecx.generate_stacktrace())
         };
 
-        let alloc: Allocation<Provenance, Self::AllocExtra, Self::Bytes> = alloc.adjust_from_tcx(
-            &ecx.tcx,
-            AllocExtra {
-                borrow_tracker,
-                data_race: race_alloc,
-                weak_memory: buffer_alloc,
-                backtrace,
-            },
-            |ptr| ecx.global_root_pointer(ptr),
-        )?;
-
         if matches!(kind, MemoryKind::Machine(kind) if kind.should_save_allocation_span()) {
             ecx.machine
                 .allocation_spans
@@ -1148,7 +1129,7 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> {
                 .insert(id, (ecx.machine.current_span(), None));
         }
 
-        Ok(Cow::Owned(alloc))
+        Ok(AllocExtra { borrow_tracker, data_race, weak_memory, backtrace })
     }
 
     fn adjust_alloc_root_pointer(
@@ -1357,7 +1338,7 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> {
     }
 
     #[inline(always)]
-    fn init_frame_extra(
+    fn init_frame(
         ecx: &mut InterpCx<'tcx, Self>,
         frame: Frame<'tcx, Provenance>,
     ) -> InterpResult<'tcx, Frame<'tcx, Provenance, FrameExtra<'tcx>>> {