diff --git a/Cargo.toml b/Cargo.toml index 386a9ed..af3c5db 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,6 +10,7 @@ linked_list_allocator = { version = "0.6.6", default-features = false } [features] panic_console = [] +debug_allocations = [] [dev-dependencies] corepack = { version = "0.4.0", default-features = false, features = ["alloc"] } diff --git a/src/entry_point.rs b/src/entry_point.rs index 2fe5c40..545d163 100644 --- a/src/entry_point.rs +++ b/src/entry_point.rs @@ -368,22 +368,82 @@ pub unsafe extern "C" fn rust_start(app_start: usize, stacktop: usize, app_heap_ } } +#[cfg(feature = "debug_allocations")] +use core::fmt::Write; +#[cfg(feature = "debug_allocations")] +use core::sync::atomic; +#[cfg(feature = "debug_allocations")] +use core::sync::atomic::AtomicUsize; + #[cfg(any(target_arch = "arm", target_arch = "riscv32"))] #[global_allocator] -static ALLOCATOR: TockAllocator = TockAllocator; +static ALLOCATOR: TockAllocator = TockAllocator::new(); static mut HEAP: Heap = Heap::empty(); -struct TockAllocator; +// With the "debug_allocations" feature, we use `AtomicUsize` to store the +// statistics because: +// - it is `Sync`, so we can use it in a static object (the allocator), +// - it implements interior mutability, so we can use it in the allocator +// methods (that take an immutable `&self` reference). +struct TockAllocator { + #[cfg(feature = "debug_allocations")] + count: AtomicUsize, + #[cfg(feature = "debug_allocations")] + size: AtomicUsize, +} + +impl TockAllocator { + const fn new() -> TockAllocator { + TockAllocator { + #[cfg(feature = "debug_allocations")] + count: AtomicUsize::new(0), + #[cfg(feature = "debug_allocations")] + size: AtomicUsize::new(0), + } + } +} unsafe impl GlobalAlloc for TockAllocator { unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - HEAP.allocate_first_fit(layout) + let ptr = HEAP + .allocate_first_fit(layout) .ok() - .map_or(0 as *mut u8, |allocation| allocation.as_ptr()) + .map_or(ptr::null_mut(), NonNull::as_ptr); + #[cfg(feature = "debug_allocations")] + { + self.count.fetch_add(1, atomic::Ordering::SeqCst); + self.size.fetch_add(layout.size(), atomic::Ordering::SeqCst); + writeln!( + crate::console::Console::new(), + "alloc[{}, {}] = {:?} ({} ptrs, {} bytes)", + layout.size(), + layout.align(), + ptr, + self.count.load(atomic::Ordering::SeqCst), + self.size.load(atomic::Ordering::SeqCst) + ) + .unwrap(); + } + ptr } unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + #[cfg(feature = "debug_allocations")] + { + self.count.fetch_sub(1, atomic::Ordering::SeqCst); + self.size.fetch_sub(layout.size(), atomic::Ordering::SeqCst); + writeln!( + crate::console::Console::new(), + "dealloc[{}, {}] = {:?} ({} ptrs, {} bytes)", + layout.size(), + layout.align(), + ptr, + self.count.load(atomic::Ordering::SeqCst), + self.size.load(atomic::Ordering::SeqCst) + ) + .unwrap(); + } HEAP.deallocate(NonNull::new_unchecked(ptr), layout) } }