#![feature(alloc)]
#![feature(heap_api)]
#![feature(thread_local)]
#![feature(const_fn)]
#![feature(placement_new_protocol)]
#![no_std]

/**
# Jenga: A stack based allocator.

*/

extern crate alloc;

use core::cell::{RefCell, RefMut};
use core::{mem, slice, ptr};
use core::ops::{Deref, DerefMut, Placer, Place, InPlace};
use core::marker::{PhantomData};

pub struct Frame<'a> {
    top:    usize,
    cap:    usize,
    _m:     PhantomData<&'a ()>
}

struct Root(RefCell<Frame<'static>>);
impl Root {
    #[inline(always)]
    fn get_mut(&self) -> RefMut<Frame<'static>> {
        let mut frame = self.0.borrow_mut();
        if frame.top == 0 {
            #[cold]
            frame.init(DEFAULT_SIZE);
        }
        assert!(frame.top != 0);
        frame
    }
}
unsafe impl Sync for Root {} // it isn't sync! don't ever expose this type

#[thread_local]
static ROOT: Root = Root(RefCell::new(Frame { top: 0, cap: 0, _m: PhantomData }));
const  DEFAULT_SIZE:   usize = 1 << 20;

#[inline(always)]
fn padding_for<T>(ptr: usize) -> usize {
    let align = mem::align_of::<T>();
    align - (ptr - 1) % align - 1
}
impl<'a> Frame<'a> {
    fn init(&mut self, size: usize) {
        use alloc;
        
        assert_eq!(self.top, 0);
        let base = unsafe { alloc::heap::allocate(size, 64) } as usize;
        self.top = base;
        self.cap = base + size;
    }
    
    /// allocate `count` elements of T
    #[inline]
    pub fn place<T>(&mut self, count: usize)
     -> Result<FArray<'a, T>, PlaceError> where T: Default
    {
        let ptr = unsafe { self.place_raw(count) }?;
        Ok(FArray::init(ptr, count))
    }
    
    #[inline(always)]
    pub unsafe fn place_raw<T>(&mut self, count: usize) -> Result<*mut T, PlaceError> {
        let size = count * mem::size_of::<T>();
        let pad = padding_for::<T>(self.top);
        let ptr = self.top + pad;
        let new_top = ptr + size;
        
        if new_top < self.cap {
            // update top pointer
            self.top = new_top;
            Ok(ptr as *mut T)
        } else {
            Err(PlaceError::OutOfMemory)
        }
    }
    
    /// drain the iterator into an array.
    /// the iterator may not use the stack allocator.
    #[inline]
    pub fn place_iter<I, T>(&mut self, mut iter: I) -> Result<FArray<T>, PlaceError>
    where I: Iterator<Item=T>
    {
        let old_top = self.top; // in case this fails
        let pad = padding_for::<T>(old_top);
        let mut p = old_top + pad;
        let mut arr = FArray { ptr: p as *mut T, len: 0, _m: PhantomData };
        
        // drain the iterator
        while let Some(e) = iter.next() {
            if p + mem::size_of::<T>() > self.cap {
                self.top = old_top;
                // the data will be dropped now
                #[cold]
                return Err(PlaceError::OutOfMemory);
            }
            unsafe {
                ptr::write(p as *mut T, e);
            }
            p += mem::size_of::<T>();
            arr.len += 1;
        }
        
        Ok(arr)
    }

    /// store something on the stack
    #[inline(always)]
    pub fn push<T>(&mut self) -> Result<FramePlacer<'a, T>, PlaceError> {
        let ptr = unsafe { self.place_raw(1) }?;
        Ok(FramePlacer { ptr: ptr, _m: PhantomData })
    }
    
    /// create a subframe. all allocations that happen within it, will be reset
    /// when it goes out if scope
    #[inline(always)]
    pub fn subframe(&mut self) -> Frame {
        Frame {
            top:    self.top,
            cap:    self.cap,
            _m:     PhantomData
        }
    }
}

pub struct FArray<'a, T: 'a> {
    ptr:    *mut T,
    len:    usize,
    _m:     PhantomData<&'a T>
}
impl<'a, T: Default + 'a> FArray<'a, T> {
    #[inline(always)]
    fn init(ptr: *mut T, len: usize) -> FArray<'a, T> {
        // initialize memory
        for i in 0 .. len {
            unsafe { ptr::write(ptr.offset(i as isize), T::default()) };
        }
        FArray { ptr, len, _m: PhantomData }
    }
}
impl<'a, T> Drop for FArray<'a, T> {
    #[inline(always)]
    fn drop(&mut self) {
        // drop elements
        for i in 0 .. self.len {
            unsafe { ptr::drop_in_place(self.ptr.offset(i as isize)) };
        }
    }
}
impl<'a, T> Deref for FArray<'a, T> {
    type Target = [T];
    #[inline(always)]
    fn deref(&self) -> &[T] {
        unsafe { slice::from_raw_parts(self.ptr, self.len) }
    }
}
impl<'a, T> DerefMut for FArray<'a, T> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut [T] {
        unsafe { slice::from_raw_parts_mut(self.ptr, self.len) }
    }
}

pub struct FBox<'a, T: 'a> {
    ptr: *mut T,
    _m:     PhantomData<&'a T>
}
impl<'a, T: 'a> Drop for FBox<'a, T> {
    #[inline(always)]
    fn drop(&mut self) {
        unsafe { ptr::drop_in_place(self.ptr) }
    }
}

pub struct FramePlacer<'a, T: 'a> {
    ptr: *mut T,
    _m:     PhantomData<&'a T>
}
pub struct FramePlace<'a, T: 'a> {
    ptr:    *mut T,
    _m:     PhantomData<&'a T>
}
impl<'a, T> Placer<T> for FramePlacer<'a, T> {
    type Place = FramePlace<'a, T>;
    #[inline(always)]
    fn make_place(self) -> Self::Place {
        FramePlace { ptr: self.ptr, _m: PhantomData }
    }
}
impl<'a, T> Place<T> for FramePlace<'a, T> {
    #[inline(always)]
    fn pointer(&mut self) -> *mut T {
        self.ptr
    }
}
impl<'a, T> InPlace<T> for FramePlace<'a, T> {
    type Owner = FBox<'a, T>;
    #[inline(always)]
    unsafe fn finalize(self) -> FBox<'a, T> {
        FBox { ptr: self.ptr, _m: PhantomData }
    }
}

/// this function will run the provided closure with pointer to uninitialized memory.
/// drop is not called afterwards.
#[inline(always)]
pub unsafe fn place_raw<F, O, T>(count: usize, f: F) -> Result<O, PlaceError>
    where F: FnOnce(*mut T) -> O
{
    let mut root = ROOT.get_mut();
    let mut sub = root.subframe();
    Ok(f(sub.place_raw(count)?))
}

/// allocate `count` elements of T
#[inline]
pub fn place<F, O, T>(count: usize, f: F) -> Result<O, PlaceError>
    where F: FnOnce(&mut [T]) -> O, T: Default
{   
    let mut root = ROOT.get_mut();
    let mut sub = root.subframe();
    let mut arr = sub.place(count)?;
    Ok(f(&mut arr))
}

/// allocate `count` elements of T
/// the iterator may not use the stack allocator.
#[inline]
pub fn place_iter<T, I, O, F>(it: I, f: F) -> Result<O, PlaceError>
    where I: Iterator<Item=T>, F: FnOnce(&mut FArray<T>) -> O
{   
    let mut root = ROOT.get_mut();
    let mut sub = root.subframe();
    let mut arr = sub.place_iter(it)?;
    Ok(f(&mut arr))
}

#[derive(Debug)]
pub enum InitError {
    AlreadyInitialized
}

/// reserve `size` bytes of memory for the allocator
/// This will fail if it was already initialized
pub fn init_with_capacity(size: usize) -> Result<(), InitError> {
    let mut root = ROOT.0.borrow_mut();
    if root.top != 0 {
        Ok(root.init(size))
    } else {
        Err(InitError::AlreadyInitialized)
    }
}

#[derive(Debug)]
pub enum PlaceError {
    OutOfMemory
}



/// lock the thread local data stack and run the closure with a reference to it
/// calling any non-framed allocations will panic
pub fn frame<F, O>(f: F) -> O
    where F: FnOnce(&mut Frame) -> O
{
    let mut root = ROOT.get_mut();
    let mut sub = root.subframe();
    f(&mut sub)
}