//! 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 region_iter;
pub use self::region_iter::AllocatorRegionIterator;
pub mod align;
pub use self::align::align_up;
pub mod locked;
pub use self::locked::LockedAllocator;
#[cfg(test)]
mod region_iter_test;
#[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 mut 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 == 0 {
                    break;
                }

                last = self.get_region_for_ptr(last.next as *mut u8).unwrap();
            }

            // Store `next` as our new region
            last.next = addr;
            region.prev = last.data - size_of::<Region>();
            region.next = 0;
        } 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 an iterator over the memory regions known to this Allocator.
    pub fn iter(&self) -> AllocatorRegionIterator {
        AllocatorRegionIterator::new(&self)
    }

    /// 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;
            }

            match region {
                Some(r) => {
                    region = self.get_region_for_ptr(r.next as *mut u8);
                },

                None => { break; },
            }
        }

        region
    }

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

                if region.next == 0 {
                    break;
                }

                region = unsafe { &mut *(region.next as *mut Region) };
            }
        }

        None
    }

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

                if region.next == 0 {
                    break;
                }

                region = unsafe { &mut *(region.next as *mut Region) };
            }
        }

        None
    }

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

        // Recursively backwards combine regions.
        //
        // If the region before the current one is adjacent in memory, and is
        // marked as free, combine it with this region.
        loop {
            // get previous region
            if region.prev == 0 { break; }
            let prev = self.get_region_for_ptr(region.prev as *mut u8);
            if prev.is_none() { break; }
            let mut prev = prev.unwrap();
            if prev.get_used() { break; }

            // if it's immediately before this region ...
            let expect_region = prev.data + prev.size();
            match unsafe { region.header_ptr() } {
                Some(r) if r.as_ptr() as usize == expect_region => {}
                Some(_) | None => {
                    break;
                }
            };

            // ... combine
            prev.size = prev.size() + region.size() + size_of::<Region>();
            prev.padding_start = 0;
            prev.padding_end = 0;
            prev.next = region.next;

            // set the `next` value correctly so the new region is in the list
            if let Some(ref mut pp) = self.get_region_for_ptr(prev.prev as *mut u8) {
                pp.next = unsafe { prev.header_ptr().unwrap().as_ptr() } as usize;
            }

            // get this region out of the list for the next iteration
            let hdrptr = unsafe { prev.header_ptr().unwrap().as_ptr() };
            region = self.get_region_for_ptr(hdrptr).expect("combine: failed to get next 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 = self.get_region_for_ptr(region.next as *mut u8)
                    .expect("combine: failed to get next region");
                if nextregion.get_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;
            }
        }
    }

    /// 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 = region.data + layout.size();
            let new_size = region.size() - layout.size();
            let mut new_region = Region::new(data_addr, new_size);
            new_region.next = region.next;
            new_region.prev = region.data - size_of::<Region>();

            // 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 { region.data_ptr().unwrap() })
    }

    /// 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_data_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);

        unsafe {
            self.combine_regions(region.header_ptr().unwrap().as_ptr());
        }
    }
}