os_essentials 0.0.3

#![no_std]

use core::alloc::{GlobalAlloc, Layout};
use core::ptr;
use core::sync::atomic::{AtomicPtr, Ordering};
use spin::Mutex;

/// A struct representing a block of memory that is free and can be allocated.
pub struct Node {
    size: usize,                   // The size of the free memory block.
    next: AtomicPtr<Node>,         // A pointer to the next free memory block in the free list.
}

/// A bump allocator with a free list to track freed memory blocks and avoid fragmentation.
pub struct BumpAllocator {
    start: usize,                  // The start address of the allocator's heap.
    end: usize,                    // The end address of the allocator's heap.
    current: usize,                // The current pointer for bump allocation.
    free_list: AtomicPtr<Node>,    // The atomic pointer to the free list of memory blocks.
    total_allocated: usize,        // The total memory allocated by the allocator.
}

impl BumpAllocator {
    /// Creates a new instance of `BumpAllocator` with a given start and end address.
    ///
    /// # Parameters
    /// - `start`: The start address of the heap.
    /// - `end`: The end address of the heap.
    ///
    /// # Returns
    /// A new `BumpAllocator` instance with the specified heap region.
    pub const fn new(start: usize, end: usize) -> Self {
        BumpAllocator {
            start,
            end,
            current: start,
            free_list: AtomicPtr::new(ptr::null_mut()),
            total_allocated: 0,
        }
    }

    /// Allocates memory of the specified layout from the heap.
    ///
    /// # Safety
    /// This function is unsafe because it involves directly manipulating raw pointers.
    ///
    /// # Parameters
    /// - `layout`: The memory layout that specifies the size and alignment of the allocation.
    ///
    /// # Returns
    /// A pointer to the allocated memory, or `null_mut` if allocation fails.
    pub unsafe fn alloc(&mut self, layout: Layout) -> *mut u8 {
        let size = layout.size();
        
        // Check if there are any free blocks in the free list.
        if !self.free_list.load(Ordering::SeqCst).is_null() {
            let mut current_node = self.free_list.load(Ordering::SeqCst);

            // Traverse the free list to find a block that fits the requested size.
            while !current_node.is_null() {
                let node_size = (*current_node).size;

                if node_size >= size {
                    // If a suitable block is found, remove it from the free list.
                    self.free_list.store((*current_node).next.load(Ordering::SeqCst), Ordering::SeqCst);
                    self.total_allocated += size;
                    return current_node as *mut u8; // Return the pointer to the allocated block.
                }

                current_node = (*current_node).next.load(Ordering::SeqCst);
            }
        }

        // No suitable block found; fallback to bump allocation.
        let new_start = self.current;
        let new_end = new_start + size;

        // If there is enough space in the heap, allocate memory by bumping the current pointer.
        if new_end <= self.end {
            self.current = new_end;
            self.total_allocated += size;
            new_start as *mut u8
        } else {
            ptr::null_mut() // Return null if there is not enough space left.
        }
    }

    /// Deallocates a previously allocated block of memory and adds it to the free list.
    ///
    /// # Safety
    /// This function is unsafe because it involves directly manipulating raw pointers.
    ///
    /// # Parameters
    /// - `ptr`: A pointer to the memory block to be deallocated.
    /// - `layout`: The layout of the memory block to be deallocated.
    pub unsafe fn dealloc(&mut self, ptr: *mut u8, layout: Layout) {
        let size = layout.size();
        let new_node = ptr as *mut Node;

        // Add the deallocated block to the free list.
        (*new_node) = Node {
            size,
            next: AtomicPtr::new(self.free_list.load(Ordering::SeqCst)),
        };

        // Update the free list to point to the new node.
        self.free_list.store(new_node, Ordering::SeqCst);
    }

    /// Returns the amount of available memory in the heap.
    ///
    /// # Returns
    /// The number of bytes of available memory.
    pub fn available_memory(&self) -> usize {
        self.end - self.current
    }

    /// Returns the total size of the heap.
    ///
    /// # Returns
    /// The size of the heap in bytes.
    pub fn heap_size(&self) -> usize {
        self.end - self.start
    }

    /// Returns the total memory allocated by the allocator.
    ///
    /// # Returns
    /// The total memory allocated in bytes.
    pub fn used_memory(&self) -> usize {
        self.total_allocated
    }

    /// Returns the amount of remaining memory in the heap.
    ///
    /// # Returns
    /// The number of bytes of remaining memory.
    pub fn remaining_memory(&self) -> usize {
        self.heap_size() - self.total_allocated
    }
}

/// A global allocator instance that uses the bump allocator to allocate and deallocate memory.
pub static GLOBAL_ALLOCATOR: Mutex<BumpAllocator> = Mutex::new(BumpAllocator::new(0x90000, 0xA0000));

/// A wrapper for the global allocator that implements the `GlobalAlloc` trait.
pub struct GlobalAllocatorWrapper;

unsafe impl GlobalAlloc for GlobalAllocatorWrapper {
    /// Allocates memory from the global bump allocator.
    ///
    /// # Safety
    /// This function is unsafe because it involves raw pointer manipulation.
    ///
    /// # Parameters
    /// - `layout`: The layout of the memory to be allocated.
    ///
    /// # Returns
    /// A pointer to the allocated memory.
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let mut allocator = GLOBAL_ALLOCATOR.lock();
        allocator.alloc(layout)
    }

    /// Deallocates a previously allocated block of memory in the global allocator.
    ///
    /// # Safety
    /// This function is unsafe because it involves raw pointer manipulation.
    ///
    /// # Parameters
    /// - `ptr`: A pointer to the memory to be deallocated.
    /// - `layout`: The layout of the memory to be deallocated.
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let mut allocator = GLOBAL_ALLOCATOR.lock();
        allocator.dealloc(ptr, layout);
    }
}