use bitsets

Author: nia-e

src/alloc/isolated_alloc.rs

@@ -3,9 +3,15 @@ use std::collections::HashMap;
 use std::hash::{BuildHasherDefault, DefaultHasher};
 use std::sync;
 
+use rustc_index::bit_set::DenseBitSet;
+
 static ALLOCATOR: sync::Mutex<IsolatedAlloc> = sync::Mutex::new(IsolatedAlloc::empty());
 
 pub struct IsolatedAlloc {
+    /// A map of machine ID to allocator. If running in multi-seeded mode,
+    /// each machine should have its own pool of memory that can be accessed
+    /// separately. We use the normal `HashMap` type so that it's available
+    /// in a `const` context.
     #[allow(rustc::default_hash_types)]
     allocators: HashMap<u64, IsolatedAllocInner, BuildHasherDefault<DefaultHasher>>,
     /// The host (not emulated) page size, or 0 if it has not yet been set.
@@ -24,14 +30,14 @@ pub struct IsolatedAllocInner {
     /// 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 length of
-    /// the boxed slice must be exactly `page_size / 8`.
+    /// 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`.
     ///
-    /// Conceptually, each bit of the `u8` represents the allocation status of one
-    /// byte on the corresponding element of `page_ptrs`; in practice, we only allocate
-    /// in 8-byte chunks currently, so the `u8`s are only ever 0 (fully free) or
-    /// 255 (fully allocated).
-    page_infos: Vec<Box<[u8]>>,
+    /// 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.
+    page_infos: Vec<DenseBitSet<usize>>,
 }
 
 // SAFETY: We only point to heap-allocated data
@@ -42,6 +48,7 @@ impl IsolatedAlloc {
     /// allow this function to be `const`; it is updated to its real value on
     /// the first call to `alloc()` or `alloc_zeroed()`.
     const fn empty() -> Self {
+        // We need this to be `const`
         #[allow(rustc::default_hash_types)]
         Self { allocators: HashMap::with_hasher(BuildHasherDefault::new()), page_size: 0 }
     }
@@ -108,14 +115,15 @@ impl IsolatedAllocInner {
     }
 
     /// Expands the available memory pool by adding one page.
-    fn add_page(&mut self, page_size: usize) -> (*mut u8, &mut Box<[u8]>) {
+    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) };
         // We don't overwrite the bytes we hand out so make sure they're zeroed by default!
         let page_ptr = unsafe { alloc::alloc(page_layout) };
-        // `page_infos` has to be one-eighth of the pagesize per the field docs
-        self.page_infos.push(vec![0u8; page_size / 8].into_boxed_slice());
+        // `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_ptrs.push(page_ptr);
         (page_ptr, self.page_infos.last_mut().unwrap())
     }
@@ -160,23 +168,28 @@ impl IsolatedAllocInner {
         page_size: usize,
         layout: Layout,
         page: *mut u8,
-        pinfo: &mut Box<[u8]>,
+        pinfo: &mut DenseBitSet<usize>,
         zeroed: bool,
     ) -> Option<*mut u8> {
         let (size, align) = IsolatedAllocInner::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)
         for idx in (0..page_size).step_by(align) {
             let idx_pinfo = idx / 8;
             let size_pinfo = size / 8;
-            if pinfo.len() < idx_pinfo + size_pinfo {
+            // DenseBitSet::contains() panics if the index is out of bounds
+            if pinfo.domain_size() < idx_pinfo + size_pinfo {
                 break;
             }
-            if pinfo[idx_pinfo..idx_pinfo + size_pinfo].iter().all(|v| *v == 0) {
-                pinfo[idx_pinfo..idx_pinfo + size_pinfo].fill(255);
+            let pred = !(idx_pinfo..idx_pinfo + size_pinfo).any(|idx| pinfo.contains(idx));
+            if pred {
+                pinfo.insert_range(idx_pinfo..idx_pinfo + size_pinfo);
                 unsafe {
                     let ptr = page.add(idx);
                     if zeroed {
-                        // Only write the bytes we were specifically asked to zero out
+                        // Only write the bytes we were specifically asked to
+                        // zero out, even if we allocated more
                         ptr.write_bytes(0, layout.size());
                     }
                     return Some(ptr);
@@ -223,14 +236,15 @@ impl IsolatedAllocInner {
             };
             let ptr_idx_pinfo = ptr_idx / 8;
             let size_pinfo = size / 8;
-            // Everything is always aligned to at least 8 bytes so this is ok
-            pinfo[ptr_idx_pinfo..ptr_idx_pinfo + size_pinfo].fill(0);
+            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(page_size, page_size) };
         for (idx, pinfo) in self.page_infos.iter().enumerate() {
-            if pinfo.iter().all(|p| *p == 0) {
+            if pinfo.is_empty() {
                 free.push(idx);
             }
         }