//! A simple memory allocator, designed for polymorphOS.

#![feature(const_fn)]
#![feature(allocator_api)]
#![feature(alloc_layout_extra)]
#![cfg_attr(not(test), no_std)]

extern crate alloc;

use alloc::alloc::{AllocErr, Layout};
use core::mem::{self, size_of};
use core::ptr::NonNull;

pub mod region;
pub use self::region::{Region, RegionFlags};
pub mod align;
pub use self::align::align_up;
pub mod locked;
pub use self::locked::LockedAllocator;
#[cfg(test)]
mod test;

/// Simple memory allocator
pub struct Allocator {
    /// Address of the header for the first known memory region.
    pub first_region: usize,
}

impl Allocator {
    /// Creates an empty Allocator with no regions.
    pub const fn empty() -> Allocator {
        Allocator { first_region: 0 }
    }

    /// The minimum size of an allocation.
    pub fn min_allocation_size() -> usize {
        size_of::<usize>() * 2
    }

    /// Adds a new memory region to this Allocator with the given `addr` and
    /// `size`.
    ///
    /// # Unsafety
    ///
    /// This function must only be called once for a given `addr` and `size`
    /// pair, and should never be called with a subregion of memory from
    /// `addr` to `addr + size - 1` given in a previous call.
    pub unsafe fn add_region(&mut self, addr: usize, size: usize) {
        if size <= Region::min_size() {
            return;
        }

        // Construct region header
        let region = Region::new(addr, size);

        if self.first_region != 0 {
            // If this is not our first region, iterate through all regions
            // until we find one with no `next` address, and add our new region
            // to that
            let mut last = self.get_region_first().unwrap();
            loop {
                if last.next().is_none() {
                    break;
                }

                last = last.next().unwrap();
            }

            // Store `next` as our new region
            last.next = addr;
        } else {
            // Else, if this is our first region, just store the new region
            // address
            self.first_region = addr;
        }

        // Store the region at the right place in memory
        let ptr = addr as *mut Region;
        ptr.write(region);
    }

    /// Returns the region header for the first known region.
    pub fn get_region_first(&self) -> Option<&mut Region> {
        if self.first_region == 0 {
            return None;
        }

        unsafe { Some(&mut *(self.first_region as *mut Region)) }
    }

    /// Returns the region header for the first known region that is not set
    /// as used, and is large enough to contain `size`.
    pub fn get_region_first_free(&self, size: usize) -> Option<&mut Region> {
        let mut region = self.get_region_first();
        loop {
            if let Some(ref r) = region {
                if r.get_used() == false && r.size() >= size {
                    break;
                }
            } else {
                break;
            }

            if region.is_some() {
                region = region.unwrap().next();
            } else {
                break;
            }
        }

        region
    }

    /// Returns the region header associated with a given data pointer (`ptr`).
    pub fn get_region_for_ptr(&self, ptr: *mut u8) -> Option<&mut Region> {
        let mut region = self.get_region_first();
        loop {
            if let Some(ref r) = region {
                if unsafe { r.data_ptr() } == ptr {
                    return region;
                }
            }

            if region.is_some() {
                region = region.unwrap().next();
            } else {
                break;
            }
        }

        None
    }

    /// Combines free regions that are adjacent to the region with data pointer
    /// `ptr`.
    pub fn combine_regions(&mut self, ptr: *mut u8) {
        let region = self.get_region_for_ptr(ptr).expect("failed to get region");

        // Recursively forward combine regions.
        //
        // If the region after the current one is adjacent in memory, and is
        // marked as free, combine it with this region.
        loop {
            if region.next == region.data + region.size() {
                let nextregion = region.next().expect("failed to get next region");
                let used = nextregion.get_used();
                if used {
                    break;
                }

                let size = region.size() + nextregion.size() + size_of::<Region>();
                let next = nextregion.next;
                core::mem::drop(nextregion);

                region.size = size;
                region.next = next;
            } else {
                break;
            }
        }

        // TODO: Backwards combine regions
    }

    /// Allocates a region of memory using the size and alignment from the
    /// given `layout`, and returns a pointer to it as a `NonNull<u8>`.
    ///
    /// # Errors
    ///
    /// Returns an `AllocErr` if there was not enough free memory to allocate
    /// a region for the requested layout.
    pub fn allocate(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
        let mut size = layout.size();
        if size < Self::min_allocation_size() {
            size = Self::min_allocation_size();
        }

        size = align_up(size, mem::align_of::<Region>());
        let layout = Layout::from_size_align(size, layout.align()).unwrap();
        let size_pad = layout.padding_needed_for(layout.align()) + layout.size();
        let region = self.get_region_first_free(size_pad);

        // No free regions big enough for this allocation
        if region.is_none() {
            return Err(AllocErr);
        }

        // Mark region as used and add data
        let mut region = region.unwrap();
        region.set_used(true);
        region.padding_start = size_pad - layout.size();

        // If there's enough space after this region to create a new region,
        // do that
        if (layout.size() + Region::min_size()) <= region.size() {
            let data_addr = unsafe { region.data_ptr() as usize } + layout.size();
            let new_size = region.size() - layout.size();
            let mut new_region = Region::new(data_addr, new_size);
            new_region.next = region.next;

            // write pointer and update next region pointer
            unsafe {
                let ptr = data_addr as *mut Region;
                ptr.write(new_region);
                region.next = ptr as usize;
            }

            region.size = layout.size();
        }

        // Successful allocation!
        Ok(unsafe { NonNull::new_unchecked(region.data_ptr()) })
    }

    /// Deallocate the region with the data pointer `ptr`, and recombines
    /// regions around it.
    pub fn deallocate(&mut self, ptr: NonNull<u8>, _layout: Layout) {
        let region = self
            .get_region_for_ptr(ptr.as_ptr())
            .expect("failed to get region for given pointer");

        let size = region.size();
        region.padding_start = 0;
        region.padding_end = 0;
        region.size = size;
        region.set_used(false);

        self.combine_regions(ptr.as_ptr());
    }
}