Inject a scalable slab allocator

ostd/libs/ostd-macros/src/lib.rs

@@ -68,11 +68,12 @@ pub fn test_main(_attr: TokenStream, item: TokenStream) -> TokenStream {
 /// A macro attribute for the global frame allocator.
 ///
 /// The attributed static variable will be used to provide frame allocation
-/// for the kernel. The variable should have type `ostd::mm::GlobalFrameAllocator`.
+/// for the kernel.
 ///
 /// # Example
 ///
 /// ```ignore
+/// use core::alloc::Layout;
 /// use ostd::{mm::{frame::GlobalFrameAllocator, Paddr}, global_frame_allocator};
 ///
 /// // Of course it won't work because all allocations will fail.
@@ -102,6 +103,78 @@ pub fn global_frame_allocator(_attr: TokenStream, item: TokenStream) -> TokenStr
     .into()
 }
 
+/// A macro attribute to register the global heap allocator.
+///
+/// The attributed static variable will be used to provide heap allocation
+/// for the kernel.
+///
+/// This attribute is not to be confused with Rust's built-in
+/// [`global_allocator`] attribute, which applies to a static variable
+/// implementing the unsafe `GlobalAlloc` trait. In contrast, the
+/// [`global_heap_allocator`] attribute does not require the heap allocator to
+/// implement an unsafe trait. [`global_heap_allocator`] eventually relies on
+/// [`global_allocator`] to customize Rust's heap allocator.
+///
+/// # Example
+///
+/// ```ignore
+/// use core::alloc::{AllocError, Layout};
+/// use ostd::{mm::heap::{GlobalHeapAllocator, HeapSlot}, global_heap_allocator};
+///
+/// // Of course it won't work and all allocations will fail.
+/// // It's just an example.
+/// #[global_heap_allocator]
+/// static ALLOCATOR: MyHeapAllocator = MyHeapAllocator;
+///
+/// struct MyHeapAllocator;
+///
+/// impl GlobalHeapAllocator for MyHeapAllocator {
+///     fn alloc(&self, _layout: Layout) -> Result<HeapSlot, AllocError> { None }
+///     fn dealloc(&self, _slot: HeapSlot) -> Result<(), AllocError> {}
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn global_heap_allocator(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    // Make a `static __GLOBAL_HEAP_ALLOCATOR_REF: &'static dyn GlobalHeapAllocator`
+    // That points to the annotated static variable.
+    let item = parse_macro_input!(item as syn::ItemStatic);
+    let static_name = &item.ident;
+
+    quote!(
+        #[no_mangle]
+        static __GLOBAL_HEAP_ALLOCATOR_REF: &'static dyn ostd::mm::heap::GlobalHeapAllocator = &#static_name;
+        #item
+    )
+    .into()
+}
+
+/// A macro attribute to provide the heap slot type given the layout.
+///
+/// The users must decide the size and the type of the heap slot to serve an
+/// allocation with the layout. The function should return `None` if the layout
+/// is not supported.
+///
+/// The annotated function should be idempotent, i.e., the result should be the
+/// same for the same layout. OSDK enforces this by only allowing the function
+/// to be `const`.
+#[proc_macro_attribute]
+pub fn global_heap_allocator_slot_type_map(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    // Rewrite the input `const fn __any_name__(layout: Layout) -> Option<SlotInfo> { ... }` to
+    // `const extern "Rust" fn __GLOBAL_HEAP_SLOT_SIZE_FROM_LAYOUT(layout: Layout) -> Option<SlotInfo> { ... }`.
+    // Reject if the input is not a `const fn`.
+    let item = parse_macro_input!(item as syn::ItemFn);
+    assert!(
+        item.sig.constness.is_some(),
+        "the annotated function must be `const`"
+    );
+
+    quote!(
+        #[export_name = "__GLOBAL_HEAP_SLOT_SIZE_FROM_LAYOUT"]
+        #item
+    )
+    .into()
+}
+
 /// A macro attribute for the panic handler.
 ///
 /// The attributed function will be used to override OSTD's default