miri/alloc/
isolated_alloc.rs

1use std::alloc::{self, Layout};
2
3use rustc_index::bit_set::DenseBitSet;
4
5/// How many bytes of memory each bit in the bitset represents.
6const COMPRESSION_FACTOR: usize = 4;
7
8/// A dedicated allocator for interpreter memory contents, ensuring they are stored on dedicated
9/// pages (not mixed with Miri's own memory). This is used in native-lib mode.
10#[derive(Debug)]
11pub struct IsolatedAlloc {
12    /// Pointers to page-aligned memory that has been claimed by the allocator.
13    /// Every pointer here must point to a page-sized allocation claimed via
14    /// the global allocator. These pointers are used for "small" allocations.
15    page_ptrs: Vec<*mut u8>,
16    /// Metadata about which bytes have been allocated on each page. The length
17    /// of this vector must be the same as that of `page_ptrs`, and the domain
18    /// size of the bitset must be exactly `page_size / COMPRESSION_FACTOR`.
19    ///
20    /// Conceptually, each bit of the bitset represents the allocation status of
21    /// one n-byte chunk on the corresponding element of `page_ptrs`. Thus,
22    /// indexing into it should be done with a value one-nth of the corresponding
23    /// offset on the matching `page_ptrs` element (n = `COMPRESSION_FACTOR`).
24    page_infos: Vec<DenseBitSet<usize>>,
25    /// Pointers to multiple-page-sized allocations. These must also be page-aligned,
26    /// with their size stored as the second element of the vector.
27    huge_ptrs: Vec<(*mut u8, usize)>,
28    /// The host (not emulated) page size.
29    page_size: usize,
30}
31
32impl IsolatedAlloc {
33    /// Creates an empty allocator.
34    pub fn new() -> Self {
35        Self {
36            page_ptrs: Vec::new(),
37            huge_ptrs: Vec::new(),
38            page_infos: Vec::new(),
39            // SAFETY: `sysconf(_SC_PAGESIZE)` is always safe to call at runtime
40            // See https://www.man7.org/linux/man-pages/man3/sysconf.3.html
41            page_size: unsafe { libc::sysconf(libc::_SC_PAGESIZE).try_into().unwrap() },
42        }
43    }
44
45    /// For simplicity, we serve small allocations in multiples of COMPRESSION_FACTOR
46    /// bytes with at least that alignment.
47    #[inline]
48    fn normalized_layout(layout: Layout) -> Layout {
49        let align =
50            if layout.align() < COMPRESSION_FACTOR { COMPRESSION_FACTOR } else { layout.align() };
51        let size = layout.size().next_multiple_of(COMPRESSION_FACTOR);
52        Layout::from_size_align(size, align).unwrap()
53    }
54
55    /// Returns the layout used to allocate the pages that hold small allocations.
56    #[inline]
57    fn page_layout(&self) -> Layout {
58        Layout::from_size_align(self.page_size, self.page_size).unwrap()
59    }
60
61    /// If the allocation is greater than a page, then round to the nearest page #.
62    #[inline]
63    fn huge_normalized_layout(layout: Layout, page_size: usize) -> Layout {
64        // Allocate in page-sized chunks
65        let size = layout.size().next_multiple_of(page_size);
66        // And make sure the align is at least one page
67        let align = std::cmp::max(layout.align(), page_size);
68        Layout::from_size_align(size, align).unwrap()
69    }
70
71    /// Determined whether a given normalized (size, align) should be sent to
72    /// `alloc_huge` / `dealloc_huge`.
73    #[inline]
74    fn is_huge_alloc(&self, layout: &Layout) -> bool {
75        layout.align() > self.page_size / 2 || layout.size() >= self.page_size / 2
76    }
77
78    /// Allocates memory as described in `Layout`. This memory should be deallocated
79    /// by calling `dealloc` on this same allocator.
80    ///
81    /// SAFETY: See `alloc::alloc()`
82    pub unsafe fn alloc(&mut self, layout: Layout) -> *mut u8 {
83        // SAFETY: Upheld by caller
84        unsafe { self.allocate(layout, false) }
85    }
86
87    /// Same as `alloc`, but zeroes out the memory.
88    ///
89    /// SAFETY: See `alloc::alloc_zeroed()`
90    pub unsafe fn alloc_zeroed(&mut self, layout: Layout) -> *mut u8 {
91        // SAFETY: Upheld by caller
92        unsafe { self.allocate(layout, true) }
93    }
94
95    /// Abstracts over the logic of `alloc_zeroed` vs `alloc`, as determined by
96    /// the `zeroed` argument.
97    ///
98    /// SAFETY: See `alloc::alloc()`, with the added restriction that `page_size`
99    /// corresponds to the host pagesize.
100    unsafe fn allocate(&mut self, layout: Layout, zeroed: bool) -> *mut u8 {
101        let layout = IsolatedAlloc::normalized_layout(layout);
102        if self.is_huge_alloc(&layout) {
103            // SAFETY: Validity of `layout` upheld by caller; we checked that
104            // the size and alignment are appropriate for being a huge alloc
105            unsafe { self.alloc_huge(layout, zeroed) }
106        } else {
107            for (&mut page, pinfo) in std::iter::zip(&mut self.page_ptrs, &mut self.page_infos) {
108                // SAFETY: The value in `self.page_size` is used to allocate
109                // `page`, with page alignment
110                if let Some(ptr) =
111                    unsafe { Self::alloc_small(self.page_size, layout, page, pinfo, zeroed) }
112                {
113                    return ptr;
114                }
115            }
116
117            // We get here only if there's no space in our existing pages
118            let page_size = self.page_size;
119            // Add another page and allocate from it; this cannot fail since the
120            // new page is empty and we already asserted it fits into a page
121            let (page, pinfo) = self.add_page();
122
123            // SAFETY: See comment on `alloc_from_page` above
124            unsafe { Self::alloc_small(page_size, layout, page, pinfo, zeroed).unwrap() }
125        }
126    }
127
128    /// Used internally by `allocate` to abstract over some logic.
129    ///
130    /// SAFETY: `page` must be a page-aligned pointer to an allocated page,
131    /// where the allocation is (at least) `page_size` bytes.
132    unsafe fn alloc_small(
133        page_size: usize,
134        layout: Layout,
135        page: *mut u8,
136        pinfo: &mut DenseBitSet<usize>,
137        zeroed: bool,
138    ) -> Option<*mut u8> {
139        // Check every alignment-sized block and see if there exists a `size`
140        // chunk of empty space i.e. forall idx . !pinfo.contains(idx / n)
141        for offset in (0..page_size).step_by(layout.align()) {
142            let offset_pinfo = offset / COMPRESSION_FACTOR;
143            let size_pinfo = layout.size() / COMPRESSION_FACTOR;
144            // DenseBitSet::contains() panics if the index is out of bounds
145            if pinfo.domain_size() < offset_pinfo + size_pinfo {
146                break;
147            }
148            if !pinfo.contains_any(offset_pinfo..offset_pinfo + size_pinfo) {
149                pinfo.insert_range(offset_pinfo..offset_pinfo + size_pinfo);
150                // SAFETY: We checked the available bytes after `idx` in the call
151                // to `domain_size` above and asserted there are at least `idx +
152                // layout.size()` bytes available and unallocated after it.
153                // `page` must point to the start of the page, so adding `idx`
154                // is safe per the above.
155                unsafe {
156                    let ptr = page.add(offset);
157                    if zeroed {
158                        // Only write the bytes we were specifically asked to
159                        // zero out, even if we allocated more
160                        ptr.write_bytes(0, layout.size());
161                    }
162                    return Some(ptr);
163                }
164            }
165        }
166        None
167    }
168
169    /// Expands the available memory pool by adding one page.
170    fn add_page(&mut self) -> (*mut u8, &mut DenseBitSet<usize>) {
171        // SAFETY: The system page size, which is the layout size, cannot be 0
172        let page_ptr = unsafe { alloc::alloc(self.page_layout()) };
173        // `page_infos` has to have one bit for each `COMPRESSION_FACTOR`-sized chunk of bytes in the page.
174        assert!(self.page_size % COMPRESSION_FACTOR == 0);
175        self.page_infos.push(DenseBitSet::new_empty(self.page_size / COMPRESSION_FACTOR));
176        self.page_ptrs.push(page_ptr);
177        (page_ptr, self.page_infos.last_mut().unwrap())
178    }
179
180    /// Allocates in multiples of one page on the host system.
181    ///
182    /// SAFETY: Same as `alloc()`.
183    unsafe fn alloc_huge(&mut self, layout: Layout, zeroed: bool) -> *mut u8 {
184        let layout = IsolatedAlloc::huge_normalized_layout(layout, self.page_size);
185        // SAFETY: Upheld by caller
186        let ret =
187            unsafe { if zeroed { alloc::alloc_zeroed(layout) } else { alloc::alloc(layout) } };
188        self.huge_ptrs.push((ret, layout.size()));
189        ret
190    }
191
192    /// Deallocates a pointer from this allocator.
193    ///
194    /// SAFETY: This pointer must have been allocated by calling `alloc()` (or
195    /// `alloc_zeroed()`) with the same layout as the one passed on this same
196    /// `IsolatedAlloc`.
197    pub unsafe fn dealloc(&mut self, ptr: *mut u8, layout: Layout) {
198        let layout = IsolatedAlloc::normalized_layout(layout);
199
200        if self.is_huge_alloc(&layout) {
201            // SAFETY: Partly upheld by caller, and we checked that the size
202            // and align, meaning this must have been allocated via `alloc_huge`
203            unsafe {
204                self.dealloc_huge(ptr, layout);
205            }
206        } else {
207            // SAFETY: It's not a huge allocation, therefore it is a small one.
208            let idx = unsafe { self.dealloc_small(ptr, layout) };
209
210            // This may have been the last allocation on this page. If so, free the entire page.
211            // FIXME: this can lead to threshold effects, we should probably add some form
212            // of hysteresis.
213            if self.page_infos[idx].is_empty() {
214                self.page_infos.remove(idx);
215                let page_ptr = self.page_ptrs.remove(idx);
216                // SAFETY: We checked that there are no outstanding allocations
217                // from us pointing to this page, and we know it was allocated
218                // with this layout
219                unsafe {
220                    alloc::dealloc(page_ptr, self.page_layout());
221                }
222            }
223        }
224    }
225
226    /// Returns the index of the page that this was deallocated from
227    ///
228    /// SAFETY: the pointer must have been allocated with `alloc_small`.
229    unsafe fn dealloc_small(&mut self, ptr: *mut u8, layout: Layout) -> usize {
230        // Offset of the pointer in the current page
231        let offset = ptr.addr() % self.page_size;
232        // And then the page's base address
233        let page_addr = ptr.addr() - offset;
234
235        // Find the page this allocation belongs to.
236        // This could be made faster if the list was sorted -- the allocator isn't fully optimized at the moment.
237        let pinfo = std::iter::zip(&mut self.page_ptrs, &mut self.page_infos)
238            .enumerate()
239            .find(|(_, (page, _))| page.addr() == page_addr);
240        let Some((idx_of_pinfo, (_, pinfo))) = pinfo else {
241            panic!("Freeing in an unallocated page: {ptr:?}\nHolding pages {:?}", self.page_ptrs)
242        };
243        // Mark this range as available in the page.
244        let ptr_idx_pinfo = offset / COMPRESSION_FACTOR;
245        let size_pinfo = layout.size() / COMPRESSION_FACTOR;
246        for idx in ptr_idx_pinfo..ptr_idx_pinfo + size_pinfo {
247            pinfo.remove(idx);
248        }
249        idx_of_pinfo
250    }
251
252    /// SAFETY: Same as `dealloc()` with the added requirement that `layout`
253    /// must ask for a size larger than the host pagesize.
254    unsafe fn dealloc_huge(&mut self, ptr: *mut u8, layout: Layout) {
255        let layout = IsolatedAlloc::huge_normalized_layout(layout, self.page_size);
256        // Find the pointer matching in address with the one we got
257        let idx = self
258            .huge_ptrs
259            .iter()
260            .position(|pg| ptr.addr() == pg.0.addr())
261            .expect("Freeing unallocated pages");
262        // And kick it from the list
263        self.huge_ptrs.remove(idx);
264        // SAFETY: Caller ensures validity of the layout
265        unsafe {
266            alloc::dealloc(ptr, layout);
267        }
268    }
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274
275    /// Helper function to assert that all bytes from `ptr` to `ptr.add(layout.size())`
276    /// are zeroes.
277    ///
278    /// SAFETY: `ptr` must have been allocated with `layout`.
279    unsafe fn assert_zeroes(ptr: *mut u8, layout: Layout) {
280        // SAFETY: Caller ensures this is valid
281        unsafe {
282            for ofs in 0..layout.size() {
283                assert_eq!(0, ptr.add(ofs).read());
284            }
285        }
286    }
287
288    /// Check that small (sub-pagesize) allocations are properly zeroed out.
289    #[test]
290    fn small_zeroes() {
291        let mut alloc = IsolatedAlloc::new();
292        // 256 should be less than the pagesize on *any* system
293        let layout = Layout::from_size_align(256, 32).unwrap();
294        // SAFETY: layout size is the constant above, not 0
295        let ptr = unsafe { alloc.alloc_zeroed(layout) };
296        // SAFETY: `ptr` was just allocated with `layout`
297        unsafe {
298            assert_zeroes(ptr, layout);
299            alloc.dealloc(ptr, layout);
300        }
301    }
302
303    /// Check that huge (> 1 page) allocations are properly zeroed out also.
304    #[test]
305    fn huge_zeroes() {
306        let mut alloc = IsolatedAlloc::new();
307        // 16k is about as big as pages get e.g. on macos aarch64
308        let layout = Layout::from_size_align(16 * 1024, 128).unwrap();
309        // SAFETY: layout size is the constant above, not 0
310        let ptr = unsafe { alloc.alloc_zeroed(layout) };
311        // SAFETY: `ptr` was just allocated with `layout`
312        unsafe {
313            assert_zeroes(ptr, layout);
314            alloc.dealloc(ptr, layout);
315        }
316    }
317
318    /// Check that repeatedly reallocating the same memory will still zero out
319    /// everything properly
320    #[test]
321    fn repeated_allocs() {
322        let mut alloc = IsolatedAlloc::new();
323        // Try both sub-pagesize allocs and those larger than / equal to a page
324        for sz in (1..=(16 * 1024)).step_by(128) {
325            let layout = Layout::from_size_align(sz, 1).unwrap();
326            // SAFETY: all sizes in the range above are nonzero as we start from 1
327            let ptr = unsafe { alloc.alloc_zeroed(layout) };
328            // SAFETY: `ptr` was just allocated with `layout`, which was used
329            // to bound the access size
330            unsafe {
331                assert_zeroes(ptr, layout);
332                ptr.write_bytes(255, sz);
333                alloc.dealloc(ptr, layout);
334            }
335        }
336    }
337
338    /// Checks that allocations of different sizes do not overlap, then for memory
339    /// leaks that might have occurred.
340    #[test]
341    fn check_leaks_and_overlaps() {
342        let mut alloc = IsolatedAlloc::new();
343
344        // Some random sizes and aligns
345        let mut sizes = vec![32; 10];
346        sizes.append(&mut vec![15; 4]);
347        sizes.append(&mut vec![256; 12]);
348        // Give it some multi-page ones too
349        sizes.append(&mut vec![32 * 1024; 4]);
350
351        // Matching aligns for the sizes
352        let mut aligns = vec![16; 12];
353        aligns.append(&mut vec![256; 2]);
354        aligns.append(&mut vec![64; 12]);
355        aligns.append(&mut vec![4096; 4]);
356
357        // Make sure we didn't mess up in the test itself!
358        assert_eq!(sizes.len(), aligns.len());
359
360        // Aggregate the sizes and aligns into a vec of layouts, then allocate them
361        let layouts: Vec<_> = std::iter::zip(sizes, aligns)
362            .map(|(sz, al)| Layout::from_size_align(sz, al).unwrap())
363            .collect();
364        // SAFETY: all sizes specified in `sizes` are nonzero
365        let ptrs: Vec<_> =
366            layouts.iter().map(|layout| unsafe { alloc.alloc_zeroed(*layout) }).collect();
367
368        for (&ptr, &layout) in std::iter::zip(&ptrs, &layouts) {
369            // We requested zeroed allocations, so check that that's true
370            // Then write to the end of the current size, so if the allocs
371            // overlap (or the zeroing is wrong) then `assert_zeroes` will panic.
372            // Also check that the alignment we asked for was respected
373            assert_eq!(ptr.addr().strict_rem(layout.align()), 0);
374            // SAFETY: each `ptr` was allocated with its corresponding `layout`,
375            // which is used to bound the access size
376            unsafe {
377                assert_zeroes(ptr, layout);
378                ptr.write_bytes(255, layout.size());
379                alloc.dealloc(ptr, layout);
380            }
381        }
382
383        // And then verify that no memory was leaked after all that
384        assert!(alloc.page_ptrs.is_empty() && alloc.huge_ptrs.is_empty());
385    }
386}