allocator/
allocator.rs

1#![feature(allocator_api, const_mut_refs, const_ptr_offset)]
2
3use buddyalloc::Heap;
4use std::{
5    alloc::{AllocError, Allocator, Layout},
6    ptr::NonNull,
7    sync::Mutex,
8};
9
10/// Declare a simple heap locked behind a Mutex.
11#[derive(Debug)]
12struct LockedHeap<const N: usize>(Mutex<Heap<N>>);
13
14/// Implement Rust's [Allocator] trait for the locked heap.
15unsafe impl<const N: usize> Allocator for LockedHeap<N> {
16    fn allocate(&self, layout: std::alloc::Layout) -> Result<NonNull<[u8]>, AllocError> {
17        let mut heap = self.0.lock().map_err(|_| AllocError)?;
18
19        let ptr = heap.allocate(layout).map_err(|_| AllocError)?;
20
21        // SAFETY: The pointer is guaranteed to not be NULL if the heap didn't return an error.
22        Ok(unsafe { NonNull::new_unchecked(std::slice::from_raw_parts_mut(ptr, layout.size())) })
23    }
24
25    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: std::alloc::Layout) {
26        let mut heap = match self.0.lock() {
27            Ok(h) => h,
28            Err(_) => return,
29        };
30
31        heap.deallocate(ptr.as_ptr(), layout);
32    }
33}
34
35fn main() {
36    // Allocate the backing memory for our heap. This memory _MUST_
37    // be aligned by at least 4096.
38    let layout = Layout::from_size_align(16384, 4096).unwrap();
39    let mem = unsafe { std::alloc::alloc(layout) };
40
41    // Construct our locked heap, with a minimum block size of 16 (16384 >> 10).
42    let heap: LockedHeap<10> = LockedHeap(Mutex::new(
43        unsafe { Heap::new(NonNull::new(mem).unwrap(), 16384) }.unwrap(),
44    ));
45    let mut vec = Vec::with_capacity_in(16, &heap);
46
47    vec.push(0usize);
48    vec.push(1usize);
49    vec.push(2usize);
50    vec.push(3usize);
51
52    println!("{:?}", vec);
53
54    // Drop the heap and vector before freeing the backing memory.
55    drop(vec);
56    drop(heap);
57
58    unsafe {
59        std::alloc::dealloc(mem, layout);
60    }
61}