From cbee8b7dde708164081fdb979e2b96740ba516a6 Mon Sep 17 00:00:00 2001
From: Nathan Ringo <nathan@remexre.com>
Date: Wed, 4 Sep 2024 02:19:32 -0500
Subject: Enables paging.

---
 crates/alloc_buddy/src/lib.rs           | 43 ++++++++++++++++++++++++++++++++-
 crates/alloc_buddy/src/tree.rs          |  5 ++++
 crates/alloc_buddy/tests/hosted_test.rs |  2 +-
 3 files changed, 48 insertions(+), 2 deletions(-)

(limited to 'crates/alloc_buddy')

diff --git a/crates/alloc_buddy/src/lib.rs b/crates/alloc_buddy/src/lib.rs
index fc5d245..083d8d5 100644
--- a/crates/alloc_buddy/src/lib.rs
+++ b/crates/alloc_buddy/src/lib.rs
@@ -232,9 +232,29 @@ impl<
         Ok(buddy_allocator)
     }
 
+    /// Tries to allocate a subregion of the smallest size class that will fit the given value.
+    #[inline(always)]
+    pub fn alloc<T>(&mut self) -> Result<NonNull<T>, AllocError> {
+        let ptr = self.alloc_of_size_class(
+            (size_of::<T>().next_power_of_two().trailing_zeros() as usize)
+                .saturating_sub(PAGE_SIZE_BITS),
+        )?;
+        Ok(ptr.cast())
+    }
+
+    /// Tries to allocate a subregion of a particular size from the allocator.
+    #[requires(size.is_power_of_two())]
+    #[requires((PAGE_SIZE_BITS..PAGE_SIZE_BITS + SIZE_CLASS_COUNT).contains(
+        &(size.trailing_zeros() as usize)
+    ))]
+    #[inline(always)]
+    pub fn alloc_of_size(&mut self, size: usize) -> Result<NonNull<u8>, AllocError> {
+        self.alloc_of_size_class(size.trailing_zeros() as usize - PAGE_SIZE_BITS)
+    }
+
     /// Tries to allocate a subregion of a particular size class from the allocator.
     #[requires(size_class < SIZE_CLASS_COUNT)]
-    pub fn alloc(&mut self, size_class: usize) -> Result<NonNull<u8>, AllocError> {
+    pub fn alloc_of_size_class(&mut self, size_class: usize) -> Result<NonNull<u8>, AllocError> {
         if let Some(ptr) = self.free_list(size_class).pop() {
             // Fast-path: try allocating from the right size class's free list.
 
@@ -309,6 +329,18 @@ impl<
         }
     }
 
+    /// Tries to allocate a subregion of the smallest size class that will fit the given value, and
+    /// zeroes it out.
+    #[inline(always)]
+    pub fn alloc_zeroed<T>(&mut self) -> Result<NonNull<T>, AllocError> {
+        let size_class = (size_of::<T>().next_power_of_two().trailing_zeros() as usize)
+            .saturating_sub(PAGE_SIZE_BITS);
+        let ptr = self.alloc_of_size_class(size_class)?;
+        // SAFETY: The allocator gives back this much memory.
+        unsafe { ptr.write_bytes(0, PAGE_SIZE << size_class) };
+        Ok(ptr.cast())
+    }
+
     /// Returns a subregion of a particular size class to the allocator.
     ///
     /// # Safety
@@ -483,6 +515,15 @@ impl<
     }
 }
 
+unsafe impl<
+        'allocator,
+        const PAGE_SIZE: usize,
+        const PAGE_SIZE_BITS: usize,
+        const SIZE_CLASS_COUNT: usize,
+    > Send for BuddyAllocator<'allocator, PAGE_SIZE, PAGE_SIZE_BITS, SIZE_CLASS_COUNT>
+{
+}
+
 #[derive(Debug)]
 struct MetadataLayout<
     const PAGE_SIZE: usize,
diff --git a/crates/alloc_buddy/src/tree.rs b/crates/alloc_buddy/src/tree.rs
index e04605b..19415f8 100644
--- a/crates/alloc_buddy/src/tree.rs
+++ b/crates/alloc_buddy/src/tree.rs
@@ -98,3 +98,8 @@ impl<const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> Tree<PAGE_SIZE, PAGE_S
         (tree_addr_lo..tree_addr_hi).contains(&addr)
     }
 }
+
+unsafe impl<const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> Send
+    for Tree<PAGE_SIZE, PAGE_SIZE_BITS>
+{
+}
diff --git a/crates/alloc_buddy/tests/hosted_test.rs b/crates/alloc_buddy/tests/hosted_test.rs
index d515f6e..90245a7 100644
--- a/crates/alloc_buddy/tests/hosted_test.rs
+++ b/crates/alloc_buddy/tests/hosted_test.rs
@@ -103,7 +103,7 @@ impl Action {
             Action::Alloc {
                 sentinel_value,
                 size_class,
-            } => match buddy.alloc(size_class) {
+            } => match buddy.alloc_of_size_class(size_class) {
                 Ok(ptr) => unsafe {
                     let slice = slice::from_raw_parts_mut(ptr.as_ptr(), PAGE_SIZE << size_class);
                     slice.fill(sentinel_value);
-- 
cgit v1.2.3