get_global_alloc_bytes refactor

Author: Strophox

src/tools/miri/src/alloc_addresses/mod.rs

@@ -356,16 +356,27 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
 
     // This returns some prepared `MiriAllocBytes`, either because `addr_from_alloc_id` reserved
     // memory space in the past, or by doing the pre-allocation right upon being called.
-    fn take_prepared_alloc_bytes(&self, id: AllocId, kind: MemoryKind) -> InterpResult<'tcx, MiriAllocBytes> {
+    fn get_global_alloc_bytes(&self, id: AllocId, kind: MemoryKind, bytes: &[u8], align: Align) -> InterpResult<'tcx, MiriAllocBytes> {
         let ecx = self.eval_context_ref();
-        // This additional call ensures that some `MiriAllocBytes` are prepared.
-        ecx.addr_from_alloc_id(id, kind)?;
-        let mut global_state = ecx.machine.alloc_addresses.borrow_mut();
-        let prepared_alloc_bytes = global_state
-            .prepared_alloc_bytes
-            .remove(&id)
-            .unwrap_or_else(|| panic!("alloc bytes for {id:?} has not been prepared"));
-        Ok(prepared_alloc_bytes)
+        Ok(if ecx.machine.native_lib.is_some() {
+            // This additional call ensures that some `MiriAllocBytes` are always prepared.
+            ecx.addr_from_alloc_id(id, kind)?;
+            let mut global_state = ecx.machine.alloc_addresses.borrow_mut();
+            // The memory we need here will have already been allocated during an earlier call to
+            // `global_root_pointer` for this allocation. So don't create a new `MiriAllocBytes` here, instead
+            // fetch the previously prepared bytes from `prepared_alloc_bytes`.
+            let mut prepared_alloc_bytes = global_state
+                .prepared_alloc_bytes
+                .remove(&id)
+                .unwrap_or_else(|| panic!("alloc bytes for {id:?} have not been prepared"));
+            assert!(prepared_alloc_bytes.as_ptr().is_aligned_to(align.bytes_usize()));
+            assert_eq!(prepared_alloc_bytes.len(), bytes.len());
+            // SAFETY: `alloc_bytes` and `bytes` span the same number of bytes.
+            unsafe { prepared_alloc_bytes.as_mut_ptr().copy_from(bytes.as_ptr(), bytes.len()) };
+            prepared_alloc_bytes
+        } else {
+            MiriAllocBytes::from_bytes(std::borrow::Cow::Borrowed(&*bytes), align)
+        })
     }
 
     /// When a pointer is used for a memory access, this computes where in which allocation the

src/tools/miri/src/lib.rs

@@ -14,6 +14,7 @@
 #![feature(strict_overflow_ops)]
 #![feature(strict_provenance)]
 #![feature(exposed_provenance)]
+#![feature(pointer_is_aligned_to)]
 // Configure clippy and other lints
 #![allow(
     clippy::collapsible_else_if,

src/tools/miri/src/machine.rs

@@ -1242,27 +1242,13 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> {
     {
         // 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");
-        // This closure is special in that it relies on `take_prepared_alloc_bytes` to give it some
-        // reserved memory space in case Miri needs to handle addresses especially for native calls.
-        let alloc = alloc.adjust_from_tcx(&ecx.tcx, |bytes, align| {
-                Ok(if ecx.machine.native_lib.is_some() {
-                    assert_eq!(alloc.align, align);
-                    let kind = Self::GLOBAL_KIND.unwrap().into();
-                    let mut alloc_bytes = ecx.take_prepared_alloc_bytes(id, kind)?;
-                    assert_eq!(alloc_bytes.len(), bytes.len());
-                    // SAFETY: `alloc_bytes` and `bytes` span the same number of bytes.
-                    unsafe { alloc_bytes.as_mut_ptr().copy_from(bytes.as_ptr(), bytes.len()) };
-                    alloc_bytes
-                } else {
-                    MiriAllocBytes::from_bytes(Cow::Borrowed(&*bytes), align)
-                })
-            }, 
+        let kind = Self::GLOBAL_KIND.unwrap().into();
+        let alloc = alloc.adjust_from_tcx(&ecx.tcx,
+            |bytes, align| ecx.get_global_alloc_bytes(id, kind, bytes, align),
             |ptr| ecx.global_root_pointer(ptr),
         )?;
         let extra =
-            Self::init_alloc_extra(ecx, id, MemoryKind::Machine(kind), alloc.size(), alloc.align)?;
+            Self::init_alloc_extra(ecx, id, kind, alloc.size(), alloc.align)?;
         Ok(Cow::Owned(alloc.with_extra(extra)))
     }