address review

Author: nia-e

src/alloc/isolated_alloc.rs

@@ -2,6 +2,9 @@ use std::alloc::{self, Layout};
 
 use rustc_index::bit_set::DenseBitSet;
 
+/// How many bytes of memory each bit in the bitset represents.
+const COMPRESSION_FACTOR: usize = 4;
+
 /// A dedicated allocator for interpreter memory contents, ensuring they are stored on dedicated
 /// pages (not mixed with Miri's own memory). This is very useful for native-lib mode.
 #[derive(Debug)]
@@ -10,19 +13,19 @@ pub struct IsolatedAlloc {
     /// Every pointer here must point to a page-sized allocation claimed via
     /// the global allocator.
     page_ptrs: Vec<*mut u8>,
-    /// Pointers to multi-page-sized allocations. These must also be page-aligned,
+    /// Pointers to multiple-page-sized allocations. These must also be page-aligned,
     /// with their size stored as the second element of the vector.
     huge_ptrs: Vec<(*mut u8, usize)>,
     /// Metadata about which bytes have been allocated on each page. The length
     /// of this vector must be the same as that of `page_ptrs`, and the domain
-    /// size of the bitset must be exactly `page_size / 8`.
+    /// size of the bitset must be exactly `page_size / COMPRESSION_FACTOR`.
     ///
     /// Conceptually, each bit of the bitset represents the allocation status of
-    /// one 8-byte chunk on the corresponding element of `page_ptrs`. Thus,
-    /// indexing into it should be done with a value one-eighth of the corresponding
-    /// offset on the matching `page_ptrs` element.
+    /// one n-byte chunk on the corresponding element of `page_ptrs`. Thus,
+    /// indexing into it should be done with a value one-nth of the corresponding
+    /// offset on the matching `page_ptrs` element (n = `COMPRESSION_FACTOR`).
     page_infos: Vec<DenseBitSet<usize>>,
-    /// The host (not emulated) page size, or 0 if it has not yet been set.
+    /// The host (not emulated) page size.
     page_size: usize,
 }
 
@@ -39,38 +42,57 @@ impl IsolatedAlloc {
 
     /// Expands the available memory pool by adding one page.
     fn add_page(&mut self, page_size: usize) -> (*mut u8, &mut DenseBitSet<usize>) {
-        assert_ne!(page_size, 0);
-
         let page_layout = unsafe { Layout::from_size_align_unchecked(page_size, page_size) };
         let page_ptr = unsafe { alloc::alloc(page_layout) };
         // `page_infos` has to have one-eighth as many bits as a page has bytes
         // (or one-64th as many bytes)
-        self.page_infos.push(DenseBitSet::new_empty(page_size / 8));
+        self.page_infos.push(DenseBitSet::new_empty(page_size / COMPRESSION_FACTOR));
         self.page_ptrs.push(page_ptr);
         (page_ptr, self.page_infos.last_mut().unwrap())
     }
 
-    /// For simplicity, we allocate in multiples of 8 bytes with at least that
-    /// alignment.
+    /// For simplicity, we serve allocations in multiples of COMPRESSION_FACTOR
+    /// bytes with at least that alignment.
     #[inline]
     fn normalized_layout(layout: Layout) -> (usize, usize) {
-        let align = if layout.align() < 8 { 8 } else { layout.align() };
-        let size = layout.size().next_multiple_of(8);
+        let align =
+            if layout.align() < COMPRESSION_FACTOR { COMPRESSION_FACTOR } else { layout.align() };
+        let size = layout.size().next_multiple_of(COMPRESSION_FACTOR);
         (size, align)
     }
 
+    /// If the allocation is greater than a page, then round to the nearest page #.
+    /// Since we pass this into the global allocator, it's more useful to return
+    /// a `Layout` instead of a pair of usizes.
+    #[inline]
+    fn huge_normalized_layout(layout: Layout, page_size: usize) -> Layout {
+        // Allocate in page-sized chunks
+        let size = layout.size().next_multiple_of(page_size);
+        // Align probably *shouldn't* ever be greater than the page size,
+        // but just in case
+        let align = std::cmp::max(layout.align(), page_size);
+        Layout::from_size_align(size, align).unwrap()
+    }
+
+    /// Determined whether a given (size, align) should be sent to `alloc_huge` /
+    /// `dealloc_huge`.
+    #[inline]
+    fn is_huge_alloc(size: usize, align: usize, page_size: usize) -> bool {
+        align >= page_size || size >= page_size
+    }
+
     /// Allocates memory as described in `Layout`. This memory should be deallocated
     /// by calling `dealloc` on this same allocator.
     ///
     /// SAFETY: See `alloc::alloc()`
-    pub fn alloc(&mut self, layout: Layout) -> *mut u8 {
+    pub unsafe fn alloc(&mut self, layout: Layout) -> *mut u8 {
         unsafe { self.allocate(layout, false) }
     }
 
     /// Same as `alloc`, but zeroes out the memory.
     ///
     /// SAFETY: See `alloc::alloc_zeroed()`
-    pub fn alloc_zeroed(&mut self, layout: Layout) -> *mut u8 {
+    pub unsafe fn alloc_zeroed(&mut self, layout: Layout) -> *mut u8 {
         unsafe { self.allocate(layout, true) }
     }
 
@@ -80,8 +102,9 @@ impl IsolatedAlloc {
     /// SAFETY: See `alloc::alloc()`, with the added restriction that `page_size`
     /// corresponds to the host pagesize.
     unsafe fn allocate(&mut self, layout: Layout, zeroed: bool) -> *mut u8 {
-        if layout.align() > self.page_size || layout.size() > self.page_size {
-            unsafe { self.alloc_multi_page(layout, zeroed) }
+        let (size, align) = IsolatedAlloc::normalized_layout(layout);
+        if IsolatedAlloc::is_huge_alloc(size, align, self.page_size) {
+            unsafe { self.alloc_huge(layout, zeroed) }
         } else {
             for (&mut page, pinfo) in std::iter::zip(&mut self.page_ptrs, &mut self.page_infos) {
                 if let Some(ptr) =
@@ -93,6 +116,8 @@ impl IsolatedAlloc {
 
             // We get here only if there's no space in our existing pages
             let page_size = self.page_size;
+            // Add another page and allocate from it; this cannot fail since the
+            // new page is empty and we already asserted it fits into a page
             let (page, pinfo) = self.add_page(page_size);
             unsafe { Self::alloc_from_page(page_size, layout, page, pinfo, zeroed).unwrap() }
         }
@@ -112,18 +137,18 @@ impl IsolatedAlloc {
         let (size, align) = IsolatedAlloc::normalized_layout(layout);
 
         // Check every alignment-sized block and see if there exists a `size`
-        // chunk of empty space i.e. forall idx . !pinfo.contains(idx / 8)
+        // chunk of empty space i.e. forall idx . !pinfo.contains(idx / n)
         for idx in (0..page_size).step_by(align) {
-            let idx_pinfo = idx / 8;
-            let size_pinfo = size / 8;
+            let idx_pinfo = idx / COMPRESSION_FACTOR;
+            let size_pinfo = size / COMPRESSION_FACTOR;
             // DenseBitSet::contains() panics if the index is out of bounds
             if pinfo.domain_size() < idx_pinfo + size_pinfo {
                 break;
             }
             // FIXME: is there a more efficient way to check whether the entire range is unset
             // in the bitset?
             let range_avail = !(idx_pinfo..idx_pinfo + size_pinfo).any(|idx| pinfo.contains(idx));
-            if pred {
+            if range_avail {
                 pinfo.insert_range(idx_pinfo..idx_pinfo + size_pinfo);
                 unsafe {
                     let ptr = page.add(idx);
@@ -142,7 +167,8 @@ impl IsolatedAlloc {
     /// Allocates in multiples of one page on the host system.
     ///
     /// SAFETY: Same as `alloc()`.
-    unsafe fn alloc_multi_page(&mut self, layout: Layout, zeroed: bool) -> *mut u8 {
+    unsafe fn alloc_huge(&mut self, layout: Layout, zeroed: bool) -> *mut u8 {
+        let layout = IsolatedAlloc::huge_normalized_layout(layout, self.page_size);
         let ret =
             unsafe { if zeroed { alloc::alloc_zeroed(layout) } else { alloc::alloc(layout) } };
         self.huge_ptrs.push((ret, layout.size()));
@@ -157,14 +183,16 @@ impl IsolatedAlloc {
     pub unsafe fn dealloc(&mut self, ptr: *mut u8, layout: Layout) {
         let (size, align) = IsolatedAlloc::normalized_layout(layout);
 
-        let ptr_idx = ptr.addr() % self.page_size;
-        let page_addr = ptr.addr() - ptr_idx;
-
-        if align > self.page_size || size > self.page_size {
+        if IsolatedAlloc::is_huge_alloc(size, align, self.page_size) {
             unsafe {
-                self.dealloc_multi_page(ptr, layout);
+                self.dealloc_huge(ptr, layout);
             }
         } else {
+            // Offset of the pointer in the current page
+            let ptr_idx = ptr.addr() % self.page_size;
+            // And then the page's base address
+            let page_addr = ptr.addr() - ptr_idx;
+
             // Find the page this allocation belongs to.
             // This could be made faster if the list was sorted -- the allocator isn't fully optimized at the moment.
             let pinfo = std::iter::zip(&mut self.page_ptrs, &mut self.page_infos)
@@ -175,49 +203,52 @@ impl IsolatedAlloc {
                     self.page_ptrs
                 )
             };
-            let ptr_idx_pinfo = ptr_idx / 8;
-            let size_pinfo = size / 8;
+            // Since each bit of the bitset represents COMPRESSION_FACTOR bytes,
+            // we need to divide by that when getting the indices
+            let ptr_idx_pinfo = ptr_idx / COMPRESSION_FACTOR;
+            let size_pinfo = size / COMPRESSION_FACTOR;
             for idx in ptr_idx_pinfo..ptr_idx_pinfo + size_pinfo {
                 pinfo.remove(idx);
             }
-        }
 
-        let mut free = vec![];
-        let page_layout =
-            unsafe { Layout::from_size_align_unchecked(self.page_size, self.page_size) };
-        for (idx, pinfo) in self.page_infos.iter().enumerate() {
-            if pinfo.is_empty() {
-                free.push(idx);
-            }
-        }
-        free.reverse();
-        for idx in free {
-            let _ = self.page_infos.remove(idx);
-            unsafe {
-                alloc::dealloc(self.page_ptrs.remove(idx), page_layout);
+            // We allocated all the pages with this layout
+            let page_layout =
+                unsafe { Layout::from_size_align_unchecked(self.page_size, self.page_size) };
+            // Only 1 page might have been freed after a dealloc, so if it exists,
+            // find it and free it (and adjust the vectors)
+            if let Some(free_idx) = self.page_infos.iter().position(|pinfo| pinfo.is_empty()) {
+                self.page_infos.remove(free_idx);
+                unsafe {
+                    alloc::dealloc(self.page_ptrs.remove(free_idx), page_layout);
+                }
             }
         }
     }
 
     /// SAFETY: Same as `dealloc()` with the added requirement that `layout`
     /// must ask for a size larger than the host pagesize.
-    unsafe fn dealloc_multi_page(&mut self, ptr: *mut u8, layout: Layout) {
+    unsafe fn dealloc_huge(&mut self, ptr: *mut u8, layout: Layout) {
+        let layout = IsolatedAlloc::huge_normalized_layout(layout, self.page_size);
+        // Find the pointer matching in address with the one we got
         let idx = self
             .huge_ptrs
             .iter()
             .position(|pg| ptr.addr() == pg.0.addr())
             .expect("Freeing unallocated pages");
-        let ptr = self.huge_ptrs.remove(idx).0;
+        // And kick it from the list
+        self.huge_ptrs.remove(idx);
         unsafe {
             alloc::dealloc(ptr, layout);
         }
     }
 }
-/*
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    /// Helper function to assert that all bytes from `ptr` to `ptr.add(layout.size())`
+    /// are zeroes.
     fn assert_zeroes(ptr: *mut u8, layout: Layout) {
         unsafe {
             for ofs in 0..layout.size() {
@@ -226,85 +257,101 @@ mod tests {
         }
     }
 
+    /// Check that small (sub-pagesize) allocations are properly zeroed out.
     #[test]
     fn small_zeroes() {
+        let mut alloc = IsolatedAlloc::new();
+        // 256 should be less than the pagesize on *any* system
         let layout = Layout::from_size_align(256, 32).unwrap();
-        // allocate_zeroed
-        let ptr = unsafe { IsolatedAlloc::alloc_zeroed(layout, 0) };
+        let ptr = unsafe { alloc.alloc_zeroed(layout) };
         assert_zeroes(ptr, layout);
         unsafe {
-            IsolatedAlloc::dealloc(ptr, layout, 0);
+            alloc.dealloc(ptr, layout);
         }
     }
 
+    /// Check that huge (> 1 page) allocations are properly zeroed out also.
     #[test]
-    fn big_zeroes() {
+    fn huge_zeroes() {
+        let mut alloc = IsolatedAlloc::new();
+        // 16k is about as big as pages get e.g. on macos aarch64
         let layout = Layout::from_size_align(16 * 1024, 128).unwrap();
-        let ptr = unsafe { IsolatedAlloc::alloc_zeroed(layout, 1) };
+        let ptr = unsafe { alloc.alloc_zeroed(layout) };
         assert_zeroes(ptr, layout);
         unsafe {
-            IsolatedAlloc::dealloc(ptr, layout, 1);
+            alloc.dealloc(ptr, layout);
         }
     }
 
+    /// Check that repeatedly reallocating the same memory will still zero out
+    /// everything properly
     #[test]
     fn repeated_allocs() {
+        let mut alloc = IsolatedAlloc::new();
+        // Try both sub-pagesize allocs and those larger than / equal to a page
         for sz in (1..=(16 * 1024)).step_by(128) {
             let layout = Layout::from_size_align(sz, 1).unwrap();
-            let ptr = unsafe { IsolatedAlloc::alloc_zeroed(layout, 2) };
+            let ptr = unsafe { alloc.alloc_zeroed(layout) };
             assert_zeroes(ptr, layout);
             unsafe {
                 ptr.write_bytes(255, sz);
-                IsolatedAlloc::dealloc(ptr, layout, 2);
+                alloc.dealloc(ptr, layout);
             }
         }
     }
 
+    /// Checks that allocations of different sizes do not overlap.
     #[test]
     fn no_overlaps() {
-        no_overlaps_inner(3);
+        let mut alloc = IsolatedAlloc::new();
+        no_overlaps_inner(&mut alloc);
     }
 
-    fn no_overlaps_inner(id: u64) {
+    /// Allows us to reuse this bit for `no_overlaps` and `check_leaks`.
+    fn no_overlaps_inner(alloc: &mut IsolatedAlloc) {
         // Some random sizes and aligns
         let mut sizes = vec![32; 10];
         sizes.append(&mut vec![15; 4]);
         sizes.append(&mut vec![256; 12]);
         // Give it some multi-page ones too
         sizes.append(&mut vec![32 * 1024; 4]);
 
+        // Matching aligns for the sizes
         let mut aligns = vec![16; 12];
         aligns.append(&mut vec![256; 2]);
         aligns.append(&mut vec![64; 12]);
         aligns.append(&mut vec![4096; 4]);
 
+        // Make sure we didn't mess up in the test itself!
         assert_eq!(sizes.len(), aligns.len());
+
+        // Aggregate the sizes and aligns into a vec of layouts, then allocate them
         let layouts: Vec<_> = std::iter::zip(sizes, aligns)
             .map(|(sz, al)| Layout::from_size_align(sz, al).unwrap())
             .collect();
-        let ptrs: Vec<_> = layouts
-            .iter()
-            .map(|layout| unsafe { IsolatedAlloc::alloc_zeroed(*layout, id) })
-            .collect();
+        let ptrs: Vec<_> =
+            layouts.iter().map(|layout| unsafe { alloc.alloc_zeroed(*layout) }).collect();
 
         for (&ptr, &layout) in std::iter::zip(&ptrs, &layouts) {
-            // Make sure we don't allocate overlapping ranges
+            // We requested zeroed allocations, so check that that's true
+            // Then write to the end of the current size, so if the allocs
+            // overlap (or the zeroing is wrong) then `assert_zeroes` will panic
             unsafe {
                 assert_zeroes(ptr, layout);
                 ptr.write_bytes(255, layout.size());
-                IsolatedAlloc::dealloc(ptr, layout, id);
+                alloc.dealloc(ptr, layout);
             }
         }
     }
 
+    /// Check for memory leaks after repeated allocations and deallocations.
     #[test]
     fn check_leaks() {
-        // Generate some noise first
-        no_overlaps_inner(4);
-        let alloc = ALLOCATOR.lock().unwrap();
+        let mut alloc = IsolatedAlloc::new();
 
-        // Should get auto-deleted if the allocations are empty
-        assert!(!alloc.allocators.contains_key(&4));
+        // Generate some noise first so leaks can manifest
+        no_overlaps_inner(&mut alloc);
+        // And then verify that no memory was leaked
+        assert!(alloc.page_ptrs.is_empty() && alloc.huge_ptrs.is_empty());
     }
 }
-*/