feat!: Add unstable Allocator API

Changes toolchain requirement to `nightly`.

Author: berkus

benches/buddy_alloc.rs

@@ -1,7 +1,12 @@
+#![feature(allocator_api)]
+
 #[macro_use]
 extern crate criterion;
 
-use criterion::{Criterion, Throughput};
+use {
+    criterion::{Criterion, Throughput},
+    std::alloc::{Allocator, Layout},
+};
 
 use buddy_alloc::buddy_alloc::{BuddyAlloc, BuddyAllocParam};
 
@@ -20,18 +25,21 @@ fn with_allocator<F: FnOnce(BuddyAlloc)>(f: F) {
 
 fn bench_alloc(allocator: &mut BuddyAlloc, alloc_size: usize) {
     for _i in 0..(ALLOC_SIZE / alloc_size) {
-        allocator.malloc(alloc_size);
+        allocator
+            .allocate(Layout::from_size_align(alloc_size, 1).unwrap())
+            .unwrap();
     }
 }
 
 fn bench_alloc_then_free(allocator: &mut BuddyAlloc, alloc_size: usize) {
     let count = ALLOC_SIZE / alloc_size;
+    let layout = Layout::from_size_align(alloc_size, 1).unwrap();
     let mut ptrs = Vec::with_capacity(count);
     for _i in 0..count {
-        ptrs.push(allocator.malloc(alloc_size));
+        ptrs.push(allocator.allocate(layout).unwrap());
     }
     for _i in 0..count {
-        allocator.free(ptrs.pop().unwrap());
+        unsafe { allocator.deallocate(ptrs.pop().unwrap().cast(), layout) };
     }
 }
 

rust-toolchain

@@ -1 +1 @@
-stable
+nightly

src/buddy_alloc.rs

@@ -6,6 +6,11 @@
 
 #![allow(clippy::needless_range_loop)]
 
+use core::{
+    alloc::{AllocError, Allocator, Layout},
+    ptr::NonNull,
+};
+
 const OOM_MSG: &str = "requires more memory space to initialize BuddyAlloc";
 const LEAF_ALIGN_ERROR_MSG: &str = "leaf size must be aligned to 16 bytes";
 /// required to align to 16 bytes, since Node takes 16 bytes on 64-bits machine.
@@ -299,11 +304,56 @@ impl BuddyAlloc {
         self.unavailable = end_addr - base_addr;
     }
 
-    pub fn malloc(&mut self, nbytes: usize) -> *mut u8 {
+    /// available bytes
+    pub fn available_bytes(&self) -> usize {
+        self.end_addr - self.unavailable - self.base_addr
+    }
+
+    fn entry(&self, i: usize) -> &Entry {
+        debug_assert!(i < self.entries_size, "index out of range");
+        unsafe { self.entries.add(i).as_ref().expect("entry") }
+    }
+
+    /// find k for p
+    fn find_k_for_p(&self, p: *const u8) -> usize {
+        for k in 0..(self.entries_size - 1) {
+            if bit_isset(self.entry(k + 1).split, self.block_index(k + 1, p)) {
+                debug_assert!(bit_isset(self.entry(k).alloc, self.block_index(k, p)));
+                return k;
+            }
+        }
+        0
+    }
+
+    /// block index of p under k
+    fn block_index(&self, k: usize, p: *const u8) -> usize {
+        if (p as usize) < self.base_addr {
+            // TODO handle this outside
+            panic!("out of memory");
+        }
+        let n = p as usize - self.base_addr;
+        // equal to: n / block_size_2base(k, self.leaf2base);
+        let index = (n >> k) >> self.leaf2base;
+        debug_assert!(index < nblock(k, self.entries_size));
+        index
+    }
+
+    /// block addr of index under k
+    fn block_addr(&self, k: usize, i: usize) -> usize {
+        // equal to: i * block_size_2base(k, self.leaf2base);
+        let n = (i << k) << self.leaf2base;
+        self.base_addr + n
+    }
+}
+
+unsafe impl Allocator for BuddyAlloc {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        let nbytes = layout.size();
+        // TODO: alignment!
         let fk = first_up_k(nbytes, 1 << self.leaf2base);
         let mut k = match (fk..self.entries_size).find(|&k| !Node::is_empty(self.entry(k).free)) {
             Some(k) => k,
-            None => return core::ptr::null_mut(),
+            None => return Err(AllocError),
         };
         let p: *mut u8 = Node::pop(self.entry(k).free) as *mut u8;
         bit_set(self.entry(k).alloc, self.block_index(k, p));
@@ -321,10 +371,15 @@ impl BuddyAlloc {
             p as usize,
             "misalignment"
         );
-        p
+
+        Ok(NonNull::slice_from_raw_parts(
+            unsafe { NonNull::new_unchecked(p) },
+            layout.size(),
+        ))
     }
 
-    pub fn free(&mut self, mut p: *mut u8) {
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, _layout: Layout) {
+        let mut p = ptr.as_ptr();
         let mut k = self.find_k_for_p(p);
         while k < (self.entries_size - 1) {
             let block_index = self.block_index(k, p);
@@ -355,45 +410,4 @@ impl BuddyAlloc {
         debug_assert!(!bit_isset(self.entry(k).alloc, self.block_index(k, p)));
         Node::push(self.entry(k).free, p);
     }
-
-    /// available bytes
-    pub fn available_bytes(&self) -> usize {
-        self.end_addr - self.unavailable - self.base_addr
-    }
-
-    fn entry(&self, i: usize) -> &Entry {
-        debug_assert!(i < self.entries_size, "index out of range");
-        unsafe { self.entries.add(i).as_ref().expect("entry") }
-    }
-
-    /// find k for p
-    fn find_k_for_p(&self, p: *const u8) -> usize {
-        for k in 0..(self.entries_size - 1) {
-            if bit_isset(self.entry(k + 1).split, self.block_index(k + 1, p)) {
-                debug_assert!(bit_isset(self.entry(k).alloc, self.block_index(k, p)));
-                return k;
-            }
-        }
-        0
-    }
-
-    /// block index of p under k
-    fn block_index(&self, k: usize, p: *const u8) -> usize {
-        if (p as usize) < self.base_addr {
-            // TODO handle this outside
-            panic!("out of memory");
-        }
-        let n = p as usize - self.base_addr;
-        // equal to: n / block_size_2base(k, self.leaf2base);
-        let index = (n >> k) >> self.leaf2base;
-        debug_assert!(index < nblock(k, self.entries_size));
-        index
-    }
-
-    /// block addr of index under k
-    fn block_addr(&self, k: usize, i: usize) -> usize {
-        // equal to: i * block_size_2base(k, self.leaf2base);
-        let n = (i << k) << self.leaf2base;
-        self.base_addr + n
-    }
 }

src/freelist_alloc.rs

@@ -1,6 +1,12 @@
 //! Freelist allocator
 //! Optimized for fixed small memory block.
 
+use core::{
+    alloc::{AllocError, Allocator, Layout},
+    cell::RefCell,
+    ptr::NonNull,
+};
+
 /// Fixed size 64 Bytes, can't allocate more in one allocation.
 pub const BLOCK_SIZE: usize = 64;
 
@@ -65,7 +71,7 @@ pub struct FreelistAlloc {
     base_addr: usize,
     /// memory end addr
     end_addr: usize,
-    free: *mut Node,
+    free: RefCell<*mut Node>,
 }
 
 impl FreelistAlloc {
@@ -94,36 +100,47 @@ impl FreelistAlloc {
         FreelistAlloc {
             base_addr,
             end_addr,
-            free,
+            free: RefCell::new(free),
         }
     }
 
     pub fn contains_ptr(&self, p: *mut u8) -> bool {
         let addr = p as usize;
         addr >= self.base_addr && addr < self.end_addr
     }
+}
 
-    pub fn malloc(&mut self, nbytes: usize) -> *mut u8 {
-        if nbytes > BLOCK_SIZE || self.free.is_null() {
-            return core::ptr::null_mut();
+unsafe impl Allocator for FreelistAlloc {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        let nbytes = layout.size();
+        // TODO: alignment!
+        if nbytes > BLOCK_SIZE || self.free.borrow().is_null() {
+            return Err(AllocError);
         }
 
-        let is_last = Node::is_empty(self.free);
-        let p = Node::pop(self.free) as *mut u8;
+        let is_last = Node::is_empty(self.free.borrow().as_const());
+        let p = Node::pop(*self.free.borrow_mut()) as *mut u8;
         if is_last {
-            self.free = core::ptr::null_mut();
+            self.free.replace(core::ptr::null_mut());
         }
-        p
+        Ok(NonNull::slice_from_raw_parts(
+            unsafe { NonNull::new_unchecked(p) },
+            layout.size(),
+        ))
     }
 
-    pub fn free(&mut self, p: *mut u8) {
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, _layout: Layout) {
+        let p = ptr.as_ptr();
         debug_assert!(self.contains_ptr(p));
-        if self.free.is_null() {
+        let f = self.free.borrow();
+        if f.is_null() {
             let n = p.cast();
             Node::init(n);
-            self.free = n;
+            drop(f);
+            self.free.replace(n);
         } else {
-            Node::push(self.free, p);
+            drop(f);
+            Node::push(*self.free.borrow_mut(), p);
         }
     }
 }

src/lib.rs

@@ -1,4 +1,8 @@
 #![cfg_attr(not(test), no_std)]
+#![feature(allocator_api)]
+#![feature(nonnull_slice_from_raw_parts)]
+#![feature(ptr_const_cast)]
+#![feature(slice_ptr_get)]
 
 pub mod buddy_alloc;
 pub mod freelist_alloc;

src/non_threadsafe_alloc.rs

@@ -1,10 +1,17 @@
 //! NonThreadSafeAlloc
 //! An allocator that does not support thread-safe
 
-use crate::buddy_alloc::{BuddyAlloc, BuddyAllocParam};
-use crate::freelist_alloc::{FreelistAlloc, FreelistAllocParam, BLOCK_SIZE};
-use core::alloc::{GlobalAlloc, Layout};
-use core::cell::RefCell;
+use {
+    crate::{
+        buddy_alloc::{BuddyAlloc, BuddyAllocParam},
+        freelist_alloc::{FreelistAlloc, FreelistAllocParam, BLOCK_SIZE},
+    },
+    core::{
+        alloc::{AllocError, Allocator, GlobalAlloc, Layout},
+        cell::RefCell,
+        ptr::NonNull,
+    },
+};
 
 /// Use buddy allocator if request bytes is large than this,
 /// otherwise use freelist allocator
@@ -50,34 +57,49 @@ impl NonThreadsafeAlloc {
     }
 }
 
-unsafe impl GlobalAlloc for NonThreadsafeAlloc {
-    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        let bytes = layout.size();
+// ==== Allocator api ====
+unsafe impl Allocator for NonThreadsafeAlloc {
+    /// Allocate a memory block from the pool.
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
         // use BuddyAlloc if size is larger than MAX_FREELIST_ALLOC_SIZE
-        if bytes > MAX_FREELIST_ALLOC_SIZE {
-            self.fetch_buddy_alloc(|alloc| alloc.malloc(bytes))
+        if layout.size() > MAX_FREELIST_ALLOC_SIZE {
+            unsafe { self.fetch_buddy_alloc(|alloc| alloc.allocate(layout)) }
         } else {
             // try freelist alloc, fallback to BuddyAlloc if failed
-            let mut p = self.fetch_freelist_alloc(|alloc| alloc.malloc(bytes));
-            if p.is_null() {
-                p = self.fetch_buddy_alloc(|alloc| alloc.malloc(bytes));
+            unsafe {
+                self.fetch_freelist_alloc(|alloc| alloc.allocate(layout))
+                    .or_else(|_| self.fetch_buddy_alloc(|alloc| alloc.allocate(layout)))
             }
-            p
         }
     }
-    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
         let freed = self.fetch_freelist_alloc(|alloc| {
-            if alloc.contains_ptr(ptr) {
-                alloc.free(ptr);
+            if alloc.contains_ptr(ptr.as_ptr()) {
+                alloc.deallocate(ptr, layout);
                 true
             } else {
                 false
             }
         });
         if !freed {
-            self.fetch_buddy_alloc(|alloc| alloc.free(ptr));
+            self.fetch_buddy_alloc(|alloc| alloc.deallocate(ptr, layout));
         }
     }
 }
 
+// ==== GlobalAlloc api ====
 unsafe impl Sync for NonThreadsafeAlloc {}
+
+unsafe impl GlobalAlloc for NonThreadsafeAlloc {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        self.allocate(layout)
+            .map_or(core::ptr::null_mut(), |p| p.as_mut_ptr())
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        if !ptr.is_null() {
+            self.deallocate(NonNull::new_unchecked(ptr), layout)
+        }
+    }
+}