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}