lean_sys/
alloc.rs

1use crate::lean_inc_heartbeat;
2use crate::LEAN_OBJECT_SIZE_DELTA;
3use core::alloc::GlobalAlloc;
4use core::alloc::Layout;
5
6/// A global allocator that uses Lean's allocator. This is useful when
7/// writing FFI libraries for Lean, where people may want to disable rust's `std`.
8/// ```ignore
9/// #![no_std]
10/// #[global_allocator]
11/// static ALLOC : lean_sys::alloc::LeanAlloc = lean_sys::alloc::LeanAlloc;
12/// ```
13pub struct LeanAlloc;
14
15extern "C" {
16    fn aligned_alloc(align: usize, size: usize) -> *mut u8;
17    fn free(ptr: *mut u8);
18}
19
20unsafe impl GlobalAlloc for LeanAlloc {
21    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
22        let alignment = layout.align().max(LEAN_OBJECT_SIZE_DELTA);
23        let offset = layout.size().wrapping_neg() & (alignment - 1);
24        let size = layout.size() + offset;
25
26        #[cfg(feature = "small_allocator")]
27        if alignment == LEAN_OBJECT_SIZE_DELTA && size <= crate::LEAN_MAX_SMALL_OBJECT_SIZE as usize
28        {
29            let slot_idx = crate::lean_get_slot_idx(size as _);
30            return crate::lean_alloc_small(size as _, slot_idx).cast();
31        }
32
33        lean_inc_heartbeat();
34        aligned_alloc(alignment, size)
35    }
36
37    #[allow(unused_variables)]
38    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
39        #[cfg(feature = "small_allocator")]
40        {
41            let alignment = layout.align().max(LEAN_OBJECT_SIZE_DELTA);
42            let offset = layout.size().wrapping_neg() & (alignment - 1);
43            let size = layout.size() + offset;
44            if alignment == LEAN_OBJECT_SIZE_DELTA
45                && size <= crate::LEAN_MAX_SMALL_OBJECT_SIZE as usize
46            {
47                return crate::lean_free_small(ptr.cast());
48            }
49        }
50
51        free(ptr);
52    }
53}
54
55#[cfg(test)]
56mod test {
57    use crate::lean_initialize_locked;
58    extern crate std;
59
60    #[test]
61    fn various_sized_allocations() {
62        unsafe {
63            lean_initialize_locked();
64        }
65        use core::alloc::GlobalAlloc;
66        let alloc = super::LeanAlloc;
67        for size in [
68            1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14, 80, 123, 255, 1024, 2048, 4095, 4096, 4097,
69            16000, 16384,
70        ] {
71            for align in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] {
72                let layout = std::alloc::Layout::from_size_align(size, align).unwrap();
73                let ptr = unsafe { alloc.alloc_zeroed(layout) };
74                assert_eq!(ptr as usize % align, 0);
75                unsafe { alloc.dealloc(ptr, layout) };
76            }
77        }
78    }
79}
lean_sys/alloc.rs

lean_sys/
alloc.rs
