1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
use core::alloc::{GlobalAlloc, Layout};
pub use slab_allocator_rs::LockedHeap;
use spin::Lazy;

// A wrapper around `slab_allocator_rs::LockedHeap` that initializes the heap lazily.
pub struct LazyHeap(Lazy<LockedHeap>);

impl LazyHeap {
    /// Create a new `LazyHeap` with the given initialization function.
    pub const fn new(init: fn() -> LockedHeap) -> Self {
        Self(Lazy::new(init))
    }

    /// Create a new `LazyHeap` with the default initialization function.
    pub const fn empty() -> Self {
        Self(Lazy::new(|| LockedHeap::empty()))
    }

    /// Initialize the heap with the given range.
    pub fn init(&self, begin: usize, len: usize) {
        unsafe {
            self.0.init(begin, len);
        }
    }
}

unsafe impl GlobalAlloc for LazyHeap {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        unsafe { self.0.alloc(layout) }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { self.0.dealloc(ptr, layout) }
    }
}