Auto merge of #3526 - Strophox:miri-memory, r=RalfJung

Adjust Allocation Bytes used by Miri to custom MiriAllocBytes

Previously, the `MiriMachine` used `type Bytes = Box<[u8]>` for its allocations.
This PR swaps this out for a custom `MiriAllocBytes` type implemented in `alloc_bytes.rs`.
This is in anticipation of an extension to Miri's FFI, which will require its allocations to take care of alignment (the methods in `impl AllocBytes for Box<[u8]>` ignore this `_align: Align` argument).

Needs https://github.com/rust-lang/rust/pull/124492

src/tools/miri/src/alloc_bytes.rs

@@ -0,0 +1,109 @@
+use std::alloc;
+use std::alloc::Layout;
+use std::borrow::Cow;
+use std::slice;
+
+use rustc_middle::mir::interpret::AllocBytes;
+use rustc_target::abi::{Align, Size};
+
+/// Allocation bytes that explicitly handle the layout of the data they're storing.
+/// This is necessary to interface with native code that accesses the program store in Miri.
+#[derive(Debug)]
+pub struct MiriAllocBytes {
+    /// Stored layout information about the allocation.
+    layout: alloc::Layout,
+    /// Pointer to the allocation contents.
+    /// Invariant:
+    /// * If `self.layout.size() == 0`, then `self.ptr` is some suitably aligned pointer
+    ///   without provenance (and no actual memory was allocated).
+    /// * Otherwise, `self.ptr` points to memory allocated with `self.layout`.
+    ptr: *mut u8,
+}
+
+impl Clone for MiriAllocBytes {
+    fn clone(&self) -> Self {
+        let bytes: Cow<'_, [u8]> = Cow::Borrowed(self);
+        let align = Align::from_bytes(self.layout.align().try_into().unwrap()).unwrap();
+        MiriAllocBytes::from_bytes(bytes, align)
+    }
+}
+
+impl Drop for MiriAllocBytes {
+    fn drop(&mut self) {
+        if self.layout.size() != 0 {
+            // SAFETY: Invariant, `self.ptr` points to memory allocated with `self.layout`.
+            unsafe { alloc::dealloc(self.ptr, self.layout) }
+        }
+    }
+}
+
+impl std::ops::Deref for MiriAllocBytes {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        // SAFETY: `ptr` is non-null, properly aligned, and valid for reading out `self.layout.size()`-many bytes.
+        // Note that due to the invariant this is true even if `self.layout.size() == 0`.
+        unsafe { slice::from_raw_parts(self.ptr, self.layout.size()) }
+    }
+}
+
+impl std::ops::DerefMut for MiriAllocBytes {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        // SAFETY: `ptr` is non-null, properly aligned, and valid for reading out `self.layout.size()`-many bytes.
+        // Note that due to the invariant this is true even if `self.layout.size() == 0`.
+        unsafe { slice::from_raw_parts_mut(self.ptr, self.layout.size()) }
+    }
+}
+
+impl MiriAllocBytes {
+    /// This method factors out how a `MiriAllocBytes` object is allocated,
+    /// specifically given an allocation function `alloc_fn`.
+    /// `alloc_fn` is only used if `size != 0`.
+    /// Returns `Err(layout)` if the allocation function returns a `ptr` that is `ptr.is_null()`.
+    fn alloc_with(
+        size: usize,
+        align: usize,
+        alloc_fn: impl FnOnce(Layout) -> *mut u8,
+    ) -> Result<MiriAllocBytes, Layout> {
+        let layout = Layout::from_size_align(size, align).unwrap();
+        let ptr = if size == 0 {
+            std::ptr::without_provenance_mut(align)
+        } else {
+            let ptr = alloc_fn(layout);
+            if ptr.is_null() {
+                return Err(layout);
+            }
+            ptr
+        };
+        // SAFETY: All `MiriAllocBytes` invariants are fulfilled.
+        Ok(Self { ptr, layout })
+    }
+}
+
+impl AllocBytes for MiriAllocBytes {
+    fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, align: Align) -> Self {
+        let slice = slice.into();
+        let size = slice.len();
+        let align = align.bytes_usize();
+        // SAFETY: `alloc_fn` will only be used if `size != 0`.
+        let alloc_fn = |layout| unsafe { alloc::alloc(layout) };
+        let alloc_bytes = MiriAllocBytes::alloc_with(size, align, alloc_fn)
+            .unwrap_or_else(|layout| alloc::handle_alloc_error(layout));
+        // SAFETY: `alloc_bytes.ptr` and `slice.as_ptr()` are non-null, properly aligned
+        // and valid for the `size`-many bytes to be copied.
+        unsafe { alloc_bytes.ptr.copy_from(slice.as_ptr(), size) };
+        alloc_bytes
+    }
+
+    fn zeroed(size: Size, align: Align) -> Option<Self> {
+        let size = size.bytes_usize();
+        let align = align.bytes_usize();
+        // SAFETY: `alloc_fn` will only be used if `size != 0`.
+        let alloc_fn = |layout| unsafe { alloc::alloc_zeroed(layout) };
+        MiriAllocBytes::alloc_with(size, align, alloc_fn).ok()
+    }
+
+    fn as_mut_ptr(&mut self) -> *mut u8 {
+        self.ptr
+    }
+}

src/tools/miri/src/diagnostics.rs

@@ -459,7 +459,7 @@ pub fn report_error<'tcx, 'mir>(
 
 pub fn report_leaks<'mir, 'tcx>(
     ecx: &InterpCx<'mir, 'tcx, MiriMachine<'mir, 'tcx>>,
-    leaks: Vec<(AllocId, MemoryKind, Allocation<Provenance, AllocExtra<'tcx>>)>,
+    leaks: Vec<(AllocId, MemoryKind, Allocation<Provenance, AllocExtra<'tcx>, MiriAllocBytes>)>,
 ) {
     let mut any_pruned = false;
     for (id, kind, mut alloc) in leaks {

src/tools/miri/src/lib.rs

@@ -13,6 +13,7 @@
 #![feature(lint_reasons)]
 #![feature(trait_upcasting)]
 #![feature(strict_overflow_ops)]
+#![feature(strict_provenance)]
 // Configure clippy and other lints
 #![allow(
     clippy::collapsible_else_if,
@@ -74,6 +75,7 @@
 extern crate rustc_driver;
 
 mod alloc_addresses;
+mod alloc_bytes;
 mod borrow_tracker;
 mod clock;
 mod concurrency;
@@ -107,6 +109,7 @@
 pub use crate::shims::EmulateItemResult;
 
 pub use crate::alloc_addresses::{EvalContextExt as _, ProvenanceMode};
+pub use crate::alloc_bytes::MiriAllocBytes;
 pub use crate::borrow_tracker::stacked_borrows::{
     EvalContextExt as _, Item, Permission, Stack, Stacks,
 };

src/tools/miri/src/machine.rs

@@ -861,7 +861,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
 
     type Provenance = Provenance;
     type ProvenanceExtra = ProvenanceExtra;
-    type Bytes = Box<[u8]>;
+    type Bytes = MiriAllocBytes;
 
     type MemoryMap =
         MonoHashMap<AllocId, (MemoryKind, Allocation<Provenance, Self::AllocExtra, Self::Bytes>)>;
@@ -1087,8 +1087,9 @@ fn adjust_allocation<'b>(
         id: AllocId,
         alloc: Cow<'b, Allocation>,
         kind: Option<MemoryKind>,
-    ) -> InterpResult<'tcx, Cow<'b, Allocation<Self::Provenance, Self::AllocExtra>>> {
-        let kind = kind.expect("we set our GLOBAL_KIND so this cannot be None");
+    ) -> 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");
         if ecx.machine.tracked_alloc_ids.contains(&id) {
             ecx.emit_diagnostic(NonHaltingDiagnostic::CreatedAlloc(
                 id,
@@ -1125,7 +1126,7 @@ fn adjust_allocation<'b>(
             Some(ecx.generate_stacktrace())
         };
 
-        let alloc: Allocation<Provenance, Self::AllocExtra> = alloc.adjust_from_tcx(
+        let alloc: Allocation<Provenance, Self::AllocExtra, Self::Bytes> = alloc.adjust_from_tcx(
             &ecx.tcx,
             AllocExtra {
                 borrow_tracker,

src/tools/miri/src/provenance_gc.rs

@@ -122,7 +122,7 @@ fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
     }
 }
 
-impl VisitProvenance for Allocation<Provenance, AllocExtra<'_>> {
+impl VisitProvenance for Allocation<Provenance, AllocExtra<'_>, MiriAllocBytes> {
     fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         for prov in self.provenance().provenances() {
             prov.visit_provenance(visit);