add associated type for allocbytes

Author: nia-e

compiler/rustc_middle/src/mir/interpret/allocation.rs

@@ -27,12 +27,20 @@ use crate::ty;
 
 /// Functionality required for the bytes of an `Allocation`.
 pub trait AllocBytes: Clone + fmt::Debug + Deref<Target = [u8]> + DerefMut<Target = [u8]> {
+    /// Used for giving extra metadata on creation. Miri uses this to allow
+    /// machine memory to be allocated separately from other allocations.
+    type ByteMetadata: Default + Clone + fmt::Debug;
+
     /// Create an `AllocBytes` from a slice of `u8`.
-    fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, _align: Align) -> Self;
+    fn from_bytes<'a>(
+        slice: impl Into<Cow<'a, [u8]>>,
+        _align: Align,
+        _dsc: Self::ByteMetadata,
+    ) -> Self;
 
     /// Create a zeroed `AllocBytes` of the specified size and alignment.
     /// Returns `None` if we ran out of memory on the host.
-    fn zeroed(size: Size, _align: Align) -> Option<Self>;
+    fn zeroed(size: Size, _align: Align, _dsc: Self::ByteMetadata) -> Option<Self>;
 
     /// Gives direct access to the raw underlying storage.
     ///
@@ -51,11 +59,13 @@ pub trait AllocBytes: Clone + fmt::Debug + Deref<Target = [u8]> + DerefMut<Targe
 
 /// Default `bytes` for `Allocation` is a `Box<u8>`.
 impl AllocBytes for Box<[u8]> {
-    fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, _align: Align) -> Self {
+    type ByteMetadata = ();
+
+    fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, _align: Align, _dsc: ()) -> Self {
         Box::<[u8]>::from(slice.into())
     }
 
-    fn zeroed(size: Size, _align: Align) -> Option<Self> {
+    fn zeroed(size: Size, _align: Align, _dsc: ()) -> Option<Self> {
         let bytes = Box::<[u8]>::try_new_zeroed_slice(size.bytes().try_into().ok()?).ok()?;
         // SAFETY: the box was zero-allocated, which is a valid initial value for Box<[u8]>
         let bytes = unsafe { bytes.assume_init() };
@@ -203,7 +213,7 @@ where
 
         let len = decoder.read_usize();
         let bytes = if all_zero { vec![0u8; len] } else { decoder.read_raw_bytes(len).to_vec() };
-        let bytes = Bytes::from_bytes(bytes, align);
+        let bytes = Bytes::from_bytes(bytes, align, Bytes::ByteMetadata::default());
 
         let provenance = Decodable::decode(decoder);
         let init_mask = Decodable::decode(decoder);
@@ -396,7 +406,7 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
         align: Align,
         mutability: Mutability,
     ) -> Self {
-        let bytes = Bytes::from_bytes(slice, align);
+        let bytes = Bytes::from_bytes(slice, align, Bytes::ByteMetadata::default());
         let size = Size::from_bytes(bytes.len());
         Self {
             bytes,
@@ -424,7 +434,7 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
         // deterministic. However, we can be non-deterministic here because all uses of const
         // evaluation (including ConstProp!) will make compilation fail (via hard error
         // or ICE) upon encountering a `MemoryExhausted` error.
-        let bytes = Bytes::zeroed(size, align).ok_or_else(fail)?;
+        let bytes = Bytes::zeroed(size, align, Bytes::ByteMetadata::default()).ok_or_else(fail)?;
 
         Ok(Allocation {
             bytes,

src/tools/miri/src/alloc/alloc_bytes.rs

@@ -0,0 +1,166 @@
+use std::alloc::Layout;
+use std::borrow::Cow;
+use std::{alloc, slice};
+
+use rustc_abi::{Align, Size};
+use rustc_middle::mir::interpret::AllocBytes;
+
+#[cfg(target_os = "linux")]
+use crate::alloc::discrete_alloc::MachineAlloc;
+use crate::helpers::ToU64 as _;
+
+#[derive(Clone, Copy, Debug, Default)]
+pub enum MiriByteMdata {
+    #[default]
+    Global,
+    Isolated,
+}
+
+/// 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` was allocated with the equivalent layout with size 1.
+    /// * Otherwise, `self.ptr` points to memory allocated with `self.layout`.
+    ptr: *mut u8,
+    /// Whether this instance of `MiriAllocBytes` had its allocation created by calling `alloc::alloc()`
+    /// (`Global`) or the discrete allocator (`Isolated`)
+    dsc: MiriByteMdata,
+}
+
+impl Clone for MiriAllocBytes {
+    fn clone(&self) -> Self {
+        let bytes: Cow<'_, [u8]> = Cow::Borrowed(self);
+        let align = Align::from_bytes(self.layout.align().to_u64()).unwrap();
+        MiriAllocBytes::from_bytes(bytes, align, self.dsc)
+    }
+}
+
+impl Drop for MiriAllocBytes {
+    fn drop(&mut self) {
+        // We have to reconstruct the actual layout used for allocation.
+        // (`Deref` relies on `size` so we can't just always set it to at least 1.)
+        let alloc_layout = if self.layout.size() == 0 {
+            Layout::from_size_align(1, self.layout.align()).unwrap()
+        } else {
+            self.layout
+        };
+
+        // SAFETY: Invariant, `self.ptr` points to memory allocated with `self.layout`.
+        unsafe {
+            match self.dsc {
+                MiriByteMdata::Global => alloc::dealloc(self.ptr, alloc_layout),
+                MiriByteMdata::Isolated => {
+                    #[cfg(target_os = "linux")]
+                    {MachineAlloc::dealloc(self.ptr, alloc_layout)}
+                    #[cfg(not(target_os = "linux"))]
+                    {unreachable!()}
+                },
+            }
+        }
+    }
+}
+
+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, given a specific allocation function.
+    /// If `size == 0` we allocate using a different `alloc_layout` with `size = 1`, to ensure each allocation has a unique address.
+    /// Returns `Err(alloc_layout)` if the allocation function returns a `ptr` where `ptr.is_null()`.
+    fn alloc_with(
+        size: u64,
+        align: u64,
+        dsc: MiriByteMdata,
+        alloc_fn: impl FnOnce(Layout) -> *mut u8,
+    ) -> Result<MiriAllocBytes, ()> {
+        let size = usize::try_from(size).map_err(|_| ())?;
+        let align = usize::try_from(align).map_err(|_| ())?;
+        let layout = Layout::from_size_align(size, align).map_err(|_| ())?;
+        // When size is 0 we allocate 1 byte anyway, to ensure each allocation has a unique address.
+        let alloc_layout =
+            if size == 0 { Layout::from_size_align(1, align).unwrap() } else { layout };
+        let ptr = alloc_fn(alloc_layout);
+        if ptr.is_null() {
+            Err(())
+        } else {
+            // SAFETY: All `MiriAllocBytes` invariants are fulfilled.
+            Ok(Self { ptr, layout, dsc })
+        }
+    }
+}
+
+impl AllocBytes for MiriAllocBytes {
+    type ByteMetadata = MiriByteMdata;
+
+    fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, align: Align, dsc: MiriByteMdata) -> Self {
+        let slice = slice.into();
+        let size = slice.len();
+        let align = align.bytes();
+        // SAFETY: `alloc_fn` will only be used with `size != 0`.
+        let alloc_fn = |layout| unsafe {
+            match dsc {
+                MiriByteMdata::Global => alloc::alloc(layout),
+                MiriByteMdata::Isolated => {
+                    #[cfg(target_os = "linux")]
+                    {MachineAlloc::alloc(layout)}
+                    #[cfg(not(target_os = "linux"))]
+                    {unreachable!()}
+                },
+            }
+        };
+        let alloc_bytes = MiriAllocBytes::alloc_with(size.to_u64(), align, dsc, alloc_fn)
+            .unwrap_or_else(|()| {
+                panic!("Miri ran out of memory: cannot create allocation of {size} bytes")
+            });
+        // 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, dsc: MiriByteMdata) -> Option<Self> {
+        let size = size.bytes();
+        let align = align.bytes();
+        // SAFETY: `alloc_fn` will only be used with `size != 0`.
+        let alloc_fn = |layout| unsafe {
+            match dsc {
+                MiriByteMdata::Global => alloc::alloc_zeroed(layout),
+                MiriByteMdata::Isolated => {
+                    #[cfg(target_os = "linux")]
+                    {MachineAlloc::alloc_zeroed(layout)}
+                    #[cfg(not(target_os = "linux"))]
+                    {unreachable!()}
+                },
+            }
+        };
+        MiriAllocBytes::alloc_with(size, align, dsc, alloc_fn).ok()
+    }
+
+    fn as_mut_ptr(&mut self) -> *mut u8 {
+        self.ptr
+    }
+
+    fn as_ptr(&self) -> *const u8 {
+        self.ptr
+    }
+}