miri/alloc/
alloc_bytes.rs

1use std::alloc::Layout;
2use std::borrow::Cow;
3use std::cell::RefCell;
4use std::rc::Rc;
5use std::{alloc, slice};
6
7use rustc_abi::{Align, Size};
8use rustc_middle::mir::interpret::AllocBytes;
9
10use crate::alloc::isolated_alloc::IsolatedAlloc;
11use crate::helpers::ToU64 as _;
12
13#[derive(Clone, Debug)]
14pub enum MiriAllocParams {
15    Global,
16    Isolated(Rc<RefCell<IsolatedAlloc>>),
17}
18
19/// Allocation bytes that explicitly handle the layout of the data they're storing.
20/// This is necessary to interface with native code that accesses the program store in Miri.
21#[derive(Debug)]
22pub struct MiriAllocBytes {
23    /// Stored layout information about the allocation.
24    layout: alloc::Layout,
25    /// Pointer to the allocation contents.
26    /// Invariant:
27    /// * If `self.layout.size() == 0`, then `self.ptr` was allocated with the equivalent layout with size 1.
28    /// * Otherwise, `self.ptr` points to memory allocated with `self.layout`.
29    ptr: *mut u8,
30    /// Whether this instance of `MiriAllocBytes` had its allocation created by calling `alloc::alloc()`
31    /// (`Global`) or the discrete allocator (`Isolated`)
32    params: MiriAllocParams,
33}
34
35impl Clone for MiriAllocBytes {
36    fn clone(&self) -> Self {
37        let bytes: Cow<'_, [u8]> = Cow::Borrowed(self);
38        let align = Align::from_bytes(self.layout.align().to_u64()).unwrap();
39        MiriAllocBytes::from_bytes(bytes, align, self.params.clone())
40    }
41}
42
43impl Drop for MiriAllocBytes {
44    fn drop(&mut self) {
45        // We have to reconstruct the actual layout used for allocation.
46        // (`Deref` relies on `size` so we can't just always set it to at least 1.)
47        let alloc_layout = if self.layout.size() == 0 {
48            Layout::from_size_align(1, self.layout.align()).unwrap()
49        } else {
50            self.layout
51        };
52
53        // SAFETY: Invariant, `self.ptr` points to memory allocated with `self.layout`.
54        unsafe {
55            match self.params.clone() {
56                MiriAllocParams::Global => alloc::dealloc(self.ptr, alloc_layout),
57                MiriAllocParams::Isolated(alloc) =>
58                    alloc.borrow_mut().dealloc(self.ptr, alloc_layout),
59            }
60        }
61    }
62}
63
64impl std::ops::Deref for MiriAllocBytes {
65    type Target = [u8];
66
67    fn deref(&self) -> &Self::Target {
68        // SAFETY: `ptr` is non-null, properly aligned, and valid for reading out `self.layout.size()`-many bytes.
69        // Note that due to the invariant this is true even if `self.layout.size() == 0`.
70        unsafe { slice::from_raw_parts(self.ptr, self.layout.size()) }
71    }
72}
73
74impl std::ops::DerefMut for MiriAllocBytes {
75    fn deref_mut(&mut self) -> &mut Self::Target {
76        // SAFETY: `ptr` is non-null, properly aligned, and valid for reading out `self.layout.size()`-many bytes.
77        // Note that due to the invariant this is true even if `self.layout.size() == 0`.
78        unsafe { slice::from_raw_parts_mut(self.ptr, self.layout.size()) }
79    }
80}
81
82impl MiriAllocBytes {
83    /// This method factors out how a `MiriAllocBytes` object is allocated, given a specific allocation function.
84    /// If `size == 0` we allocate using a different `alloc_layout` with `size = 1`, to ensure each allocation has a unique address.
85    /// Returns `Err(alloc_layout)` if the allocation function returns a `ptr` where `ptr.is_null()`.
86    fn alloc_with(
87        size: u64,
88        align: u64,
89        params: MiriAllocParams,
90        alloc_fn: impl FnOnce(Layout, &MiriAllocParams) -> *mut u8,
91    ) -> Result<MiriAllocBytes, ()> {
92        let size = usize::try_from(size).map_err(|_| ())?;
93        let align = usize::try_from(align).map_err(|_| ())?;
94        let layout = Layout::from_size_align(size, align).map_err(|_| ())?;
95        // When size is 0 we allocate 1 byte anyway, to ensure each allocation has a unique address.
96        let alloc_layout =
97            if size == 0 { Layout::from_size_align(1, align).unwrap() } else { layout };
98        let ptr = alloc_fn(alloc_layout, &params);
99        if ptr.is_null() {
100            Err(())
101        } else {
102            // SAFETY: All `MiriAllocBytes` invariants are fulfilled.
103            Ok(Self { ptr, layout, params })
104        }
105    }
106}
107
108impl AllocBytes for MiriAllocBytes {
109    type AllocParams = MiriAllocParams;
110
111    fn from_bytes<'a>(
112        slice: impl Into<Cow<'a, [u8]>>,
113        align: Align,
114        params: MiriAllocParams,
115    ) -> Self {
116        let slice = slice.into();
117        let size = slice.len();
118        let align = align.bytes();
119        // SAFETY: `alloc_fn` will only be used with `size != 0`.
120        let alloc_fn = |layout, params: &MiriAllocParams| unsafe {
121            match params {
122                MiriAllocParams::Global => alloc::alloc(layout),
123                MiriAllocParams::Isolated(alloc) => alloc.borrow_mut().alloc(layout),
124            }
125        };
126        let alloc_bytes = MiriAllocBytes::alloc_with(size.to_u64(), align, params, alloc_fn)
127            .unwrap_or_else(|()| {
128                panic!("Miri ran out of memory: cannot create allocation of {size} bytes")
129            });
130        // SAFETY: `alloc_bytes.ptr` and `slice.as_ptr()` are non-null, properly aligned
131        // and valid for the `size`-many bytes to be copied.
132        unsafe { alloc_bytes.ptr.copy_from(slice.as_ptr(), size) };
133        alloc_bytes
134    }
135
136    fn zeroed(size: Size, align: Align, params: MiriAllocParams) -> Option<Self> {
137        let size = size.bytes();
138        let align = align.bytes();
139        // SAFETY: `alloc_fn` will only be used with `size != 0`.
140        let alloc_fn = |layout, params: &MiriAllocParams| unsafe {
141            match params {
142                MiriAllocParams::Global => alloc::alloc_zeroed(layout),
143                MiriAllocParams::Isolated(alloc) => alloc.borrow_mut().alloc_zeroed(layout),
144            }
145        };
146        MiriAllocBytes::alloc_with(size, align, params, alloc_fn).ok()
147    }
148
149    fn as_mut_ptr(&mut self) -> *mut u8 {
150        self.ptr
151    }
152
153    fn as_ptr(&self) -> *const u8 {
154        self.ptr
155    }
156}