Use very simple bump allocator as global allocator

Signed-off-by: Graham MacDonald <[email protected]>

aarch64/src/main.rs

@@ -22,13 +22,16 @@ mod vm;
 
 use crate::kmem::from_virt_to_physaddr;
 use crate::vm::kernel_root;
+use alloc::boxed::Box;
 use core::ffi::c_void;
 use core::ptr;
 use port::fdt::DeviceTree;
 use port::mem::PhysRange;
-use port::println;
+use port::{println, vmalloc};
 use vm::PageTable;
 
+extern crate alloc;
+
 #[cfg(not(test))]
 core::arch::global_asm!(include_str!("l.S"));
 
@@ -145,6 +148,13 @@ pub extern "C" fn main9(dtb_va: usize) {
 
     println!("looping now");
 
+    vmalloc::print_status();
+    {
+        let _b = Box::new("ddododo");
+        vmalloc::print_status();
+    }
+    vmalloc::print_status();
+
     #[allow(clippy::empty_loop)]
     loop {}
 }

aarch64/src/runtime.rs

@@ -5,11 +5,15 @@ extern crate alloc;
 use crate::kmem::physaddr_as_virt;
 use crate::registers::rpi_mmio;
 use crate::uartmini::MiniUart;
-use alloc::alloc::{GlobalAlloc, Layout};
+use alloc::alloc::Layout;
 use core::fmt::Write;
 use core::panic::PanicInfo;
 use port::devcons::PanicConsole;
 use port::mem::VirtRange;
+use port::vmalloc;
+
+#[global_allocator]
+static ALLOCATOR: vmalloc::VmAllocator = vmalloc::VmAllocator {};
 
 // TODO
 //  - Add qemu integration test
@@ -39,17 +43,3 @@ pub fn panic(info: &PanicInfo) -> ! {
 fn oom(_layout: Layout) -> ! {
     panic!("oom");
 }
-
-struct FakeAlloc;
-
-unsafe impl GlobalAlloc for FakeAlloc {
-    unsafe fn alloc(&self, _layout: Layout) -> *mut u8 {
-        panic!("fake alloc");
-    }
-    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
-        panic!("fake dealloc");
-    }
-}
-
-#[global_allocator]
-static FAKE_ALLOCATOR: FakeAlloc = FakeAlloc {};

port/src/bumpalloc.rs

@@ -0,0 +1,160 @@
+use core::alloc::{AllocError, Allocator, Layout};
+use core::cell::UnsafeCell;
+use core::ptr::NonNull;
+use core::sync::atomic::{AtomicUsize, Ordering::Relaxed};
+
+#[cfg(not(test))]
+use crate::println;
+
+/// Bump allocator to be used for earliest allocations in r9.  These allocations
+/// can never be freed - attempting to do so will panic.
+#[repr(C, align(4096))]
+pub struct Bump<const SIZE_BYTES: usize, const MAX_SUPPORTED_ALIGN: usize> {
+    bytes: UnsafeCell<[u8; SIZE_BYTES]>,
+    next_offset: AtomicUsize,
+    wasted: AtomicUsize,
+}
+
+unsafe impl<const SIZE_BYTES: usize, const MAX_SUPPORTED_ALIGN: usize> Send
+    for Bump<SIZE_BYTES, MAX_SUPPORTED_ALIGN>
+{
+}
+
+unsafe impl<const SIZE_BYTES: usize, const MAX_SUPPORTED_ALIGN: usize> Sync
+    for Bump<SIZE_BYTES, MAX_SUPPORTED_ALIGN>
+{
+}
+
+impl<const SIZE_BYTES: usize, const MAX_SUPPORTED_ALIGN: usize>
+    Bump<SIZE_BYTES, MAX_SUPPORTED_ALIGN>
+{
+    pub const fn new(init_value: u8) -> Self {
+        Self {
+            bytes: UnsafeCell::new([init_value; SIZE_BYTES]),
+            next_offset: AtomicUsize::new(0),
+            wasted: AtomicUsize::new(0),
+        }
+    }
+
+    pub fn print_status(&self) {
+        let allocated = self.next_offset.load(Relaxed);
+        let remaining = SIZE_BYTES - allocated;
+        let wasted = self.wasted.load(Relaxed);
+        println!(
+            "Bump: allocated: {allocated} free: {remaining} total: {SIZE_BYTES} wasted: {wasted}"
+        );
+    }
+
+    /// Test helper to get the offset of the result in the buffer
+    #[cfg(test)]
+    fn result_offset(&self, result: Result<NonNull<[u8]>, AllocError>) -> Option<isize> {
+        unsafe {
+            result
+                .ok()
+                .map(|bytes| bytes.byte_offset_from(NonNull::new_unchecked(self.bytes.get())))
+        }
+    }
+}
+
+unsafe impl<const SIZE_BYTES: usize, const MAX_SUPPORTED_ALIGN: usize> Allocator
+    for Bump<SIZE_BYTES, MAX_SUPPORTED_ALIGN>
+{
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        let size = layout.size();
+        let align = layout.align();
+
+        if align > MAX_SUPPORTED_ALIGN {
+            return Err(AllocError {});
+        }
+
+        let mut wasted = 0;
+        let mut alloc_offset = 0;
+        let alloc_ok = self
+            .next_offset
+            .fetch_update(Relaxed, Relaxed, |last_offset| {
+                let align_mask = !(align - 1);
+                alloc_offset = if last_offset & !align_mask != 0 {
+                    (last_offset + align) & align_mask
+                } else {
+                    last_offset
+                };
+                wasted = alloc_offset - last_offset;
+
+                let new_offset = alloc_offset + size;
+                if new_offset > SIZE_BYTES {
+                    None
+                } else {
+                    Some(new_offset)
+                }
+            })
+            .is_err();
+
+        if alloc_ok {
+            Err(AllocError {})
+        } else {
+            self.wasted.fetch_add(wasted, Relaxed);
+            Ok(unsafe { NonNull::new_unchecked(self.bytes.get().byte_add(alloc_offset)) })
+        }
+    }
+
+    unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: Layout) {
+        // panic!("Can't deallocate from Bump allocator (ptr: {:p}, layout: {:?})", ptr, layout)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::mem::PAGE_SIZE_4K;
+
+    use super::*;
+
+    #[test]
+    fn bump_new() {
+        let bump = Bump::<PAGE_SIZE_4K, PAGE_SIZE_4K>::new(0);
+        let result = unsafe { bump.allocate(Layout::from_size_align_unchecked(4096, 4096)) };
+        assert!(result.is_ok());
+        assert_eq!(bump.result_offset(result), Some(0));
+        assert_eq!(bump.wasted.load(Relaxed), 0);
+        assert_eq!(bump.next_offset.load(Relaxed), 4096);
+
+        // Next should fail - out of space
+        let result = unsafe { bump.allocate(Layout::from_size_align_unchecked(1, 1)) };
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn bump_alignment() {
+        let bump = Bump::<{ 3 * PAGE_SIZE_4K }, PAGE_SIZE_4K>::new(0);
+
+        // Small allocation
+        let mut expected_waste = 0;
+        let result = unsafe { bump.allocate(Layout::from_size_align_unchecked(16, 1)) };
+        assert!(result.is_ok());
+        assert_eq!(bump.result_offset(result), Some(0));
+        assert_eq!(bump.wasted.load(Relaxed), expected_waste);
+        assert_eq!(bump.next_offset.load(Relaxed), 16);
+
+        // Align next allocation to 4096, wasting space
+        expected_waste += 4096 - 16;
+        let result = unsafe { bump.allocate(Layout::from_size_align_unchecked(16, 4096)) };
+        assert!(result.is_ok());
+        assert_eq!(bump.result_offset(result), Some(4096));
+        assert_eq!(bump.wasted.load(Relaxed), expected_waste);
+        assert_eq!(bump.next_offset.load(Relaxed), 4096 + 16);
+
+        // Align next allocation to 4096, wasting space
+        expected_waste += 4096 - 16;
+        let result = unsafe { bump.allocate(Layout::from_size_align_unchecked(4096, 4096)) };
+        assert!(result.is_ok());
+        assert_eq!(bump.result_offset(result), Some(2 * 4096));
+        assert_eq!(bump.wasted.load(Relaxed), expected_waste);
+        assert_eq!(bump.next_offset.load(Relaxed), 3 * 4096);
+    }
+
+    #[test]
+    fn align_too_high() {
+        let bump = Bump::<PAGE_SIZE_4K, PAGE_SIZE_4K>::new(0);
+        let result = unsafe { bump.allocate(Layout::from_size_align_unchecked(4096, 8192)) };
+        assert!(result.is_err());
+    }
+}

port/src/lib.rs

@@ -1,12 +1,17 @@
 #![allow(clippy::upper_case_acronyms)]
 #![cfg_attr(not(any(test)), no_std)]
+#![feature(allocator_api)]
 #![feature(maybe_uninit_slice)]
 #![feature(step_trait)]
 #![forbid(unsafe_op_in_unsafe_fn)]
 
 pub mod bitmapalloc;
+pub mod bumpalloc;
 pub mod dat;
 pub mod devcons;
 pub mod fdt;
 pub mod mcslock;
 pub mod mem;
+pub mod vmalloc;
+
+extern crate alloc;

port/src/vmalloc.rs

@@ -0,0 +1,27 @@
+use crate::{bumpalloc::Bump, mem::PAGE_SIZE_4K};
+use alloc::alloc::{GlobalAlloc, Layout};
+use core::{alloc::Allocator, ptr::null_mut};
+
+#[cfg(not(test))]
+use crate::println;
+
+static BUMP_ALLOC: Bump<{ 32 * 256 * PAGE_SIZE_4K }, PAGE_SIZE_4K> = Bump::new(0);
+
+pub struct VmAllocator {}
+
+unsafe impl GlobalAlloc for VmAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        println!("vmalloc::alloc");
+
+        let result = BUMP_ALLOC.allocate(layout);
+        result.map_or(null_mut(), |b| b.as_ptr() as *mut u8)
+    }
+
+    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
+        println!("vmalloc::dealloc");
+    }
+}
+
+pub fn print_status() {
+    BUMP_ALLOC.print_status();
+}