#![feature(thread_local)]
#![feature(const_fn)]
#![feature(placement_new_protocol)]
#![feature(placement_in_syntax)]
#![feature(nonzero)]
#![cfg_attr(not(feature = "std"), no_std)]
#![cfg_attr(feature = "alloc", feature(alloc, heap_api))]


/*!
Jenga: A stack based allocator.

## General principles: 
There are two ways of making allocations:

 a) by calling one of the [place](fn.place.html) or
    [place_iter](fn.place_iter.html) methods,
    which then allow access to the data in the provided closure.
 b) by calling [frame](fn.frame.html) which allows access to
    the thread_local [Frame](struct.Frame.html),
    from which futher allocations can be made.
    This method might be slightly faster as no additional checks are nessesary.

## Memory Pool

As a default size, 1MB of memory will be allocated whenever the allocator is
used first. This can be changed by calling
[init_with_capacity](fn.init_with_capacity.html) first.
*/

#[cfg(feature = "std")]
extern crate core;
#[cfg(feature = "alloc")]
extern crate alloc;

#[cfg(feature = "std")]
use std::io;

use core::{cell, mem, slice, ptr, ops, marker, fmt, nonzero, str};

use cell::{RefCell, RefMut};
use ops::{Deref, DerefMut, Placer, Place, InPlace};
use marker::{PhantomData};
use nonzero::NonZero;

#[derive(Debug)]
pub struct InitError {}

#[derive(Debug)]
pub struct OutOfMemory {}

/**
An allocation Frame.

Data can be allocated using:

 - `place` which allocates b bytes of type T and returns a mutable view to it.
 - `place_raw` if you really want unsafe code.
 - `place_iter` will drain the given iterator and return a mutable view to the data.
 - `cursor` gives you a gowing buffer that supports `io::Write`.
 - `writer` does the same, but supports `fmt::Write` and is can be turned into `&str`
 
All returned buffers are valid for the lifetime of the `Frame`.
*/

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]
            #[cfg(feature = "alloc")]
            frame.init(DEFAULT_SIZE);
            
            #[cfg(not(feature = "alloc"))]
            panic!("Root frame not initialized");
        }
        assert!(frame.top != 0);
        frame
    }
}
unsafe impl Sync for Root {} // it isn't sync! don't ever expose this type

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

fn get_root() -> RefMut<'static, Frame<'static>> {
    unsafe { ROOT.get_mut() }
}

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


#[inline(always)]
fn padding_for<T>(ptr: usize) -> usize {
    let align = mem::align_of::<T>();
    align - (ptr - 1) % align - 1
}
impl<'a> Frame<'a> {
    #[cfg(feature = "alloc")]
    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;
    }
    
    /// Turn a slice of memory into an allocator Frame.
    pub fn from_slice(s: &'a mut [u8]) -> Frame<'a> {
        let ptr = s.as_mut_ptr() as usize;
        let len = s.len();
        
        Frame {
            top:    ptr,
            cap:    ptr + len,
            _m:     PhantomData
        }
    }
    
    /// Allocate `count` elements of `T`.
    /// Returns `OutOfMemory` if there is not enough space left.
    #[inline]
    pub fn try_place<T>(&mut self, count: usize)
     -> Result<FArray<'a, T>, OutOfMemory> where T: Default
    {
        let ptr = try!(unsafe { self.try_place_raw(count) });
        Ok(FArray::init(ptr.get(), count))
    }
    
    /// allocate `count` elements of `T`.
    /// panics if there is not enough space left.
    #[inline]
    pub fn place<T>(&mut self, count: usize) -> FArray<'a, T> where T: Default
    {
        let ptr = unsafe { self.place_raw(count) };
        FArray::init(ptr, count)
    }
    
    /// Allocate `count` elements of `T`.
    /// Returns `OutOfMemory` if there is not enough space left.
    /// The caller has to ensure the data does not outlife the `Frame`
    #[inline(always)]
    pub unsafe fn try_place_raw<T>(&mut self, count: usize)
     -> Result<NonZero<*mut T>, OutOfMemory>
    {
        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(NonZero::new(ptr as *mut T))
        } else {
            #[cold]
            Err(OutOfMemory{})
        }
    }
    
    /// Allocate `count` elements of `T`.
    /// panics if there is not enough space left.
    /// The caller has to ensure the data does not outlife the `Frame`
    #[inline(always)]
    pub unsafe fn place_raw<T>(&mut self, count: usize) -> *mut T {
        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 {
            panic!("Out of Memory");
        }
        // update top pointer
        self.top = new_top;
        ptr as *mut T
    }
    
    /// Drain the provided Iterator and return the resulting data.
    /// The iterator may not use the stack allocator.
    /// panics if the iterator attempts to use the stack allocator.
    #[inline]
    pub fn place_iter<I, T>(&mut self, mut iter: I) -> Result<FArray<'a, T>, OutOfMemory>
    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(OutOfMemory{});
            }
            unsafe {
                ptr::write(p as *mut T, e);
            }
            p += mem::size_of::<T>();
            arr.len += 1;
        }
        self.top = p;
        
        Ok(arr)
    }

    /// Store something on the stack. To be used with the place syntax.
    #[inline(always)]
    pub fn push<T>(&mut self) -> FramePlacer<'a, T> {
        let ptr = unsafe { self.place_raw(1) };
        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
        }
    }
    
    /// Create a growing Buffer that supports io::Write
    /// call `finish()` on the buffer to complete it.
    #[cfg(feature = "std")]
    #[inline(always)]
    pub fn cursor<'b>(&'b mut self) -> Cursor<'a, 'b> {
        Cursor {
            ptr:    self.top as *mut u8,
            len:    0,
            cap:    self.cap - self.top,
            frame:  self
        }
    }
    
    /// Create a growing Buffer that supports fmt::Write
    /// call `finish()` on the buffer to complete it.
    #[inline(always)]
    pub fn writer<'b>(&'b mut self) -> Writer<'a, 'b> {
        Writer {
            ptr:    self.top as *mut u8,
            len:    0,
            cap:    self.cap - self.top,
            frame:  self
        }
    }
}

/// 'a: lifetime of the frame
/// 'b: the borrow of &mut Frame
#[cfg(feature = "std")]
pub struct Cursor<'a, 'b> where 'a: 'b  {
    ptr:    *mut u8,
    len:    usize,
    cap:    usize, // in bytes,
    frame:  &'b mut Frame<'a>
}
#[cfg(feature = "std")]
impl<'a, 'b> Cursor<'a, 'b> {
    pub fn finish(self) -> &'a mut [u8] {
        self.frame.top += self.len;
        
        unsafe {
            slice::from_raw_parts_mut(self.ptr, self.len)
        }
    }
}
#[cfg(feature = "std")]
impl<'a, 'b> io::Write for Cursor<'a, 'b> {
    #[inline(always)]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if self.len + buf.len() <= self.cap {
            unsafe {
                ptr::copy(
                    buf.as_ptr(),
                    self.ptr.offset(self.len as isize),
                    buf.len()
                );
            }
            self.len += buf.len();
            Ok(buf.len())
        } else {
            Err(io::ErrorKind::Other.into())
        }
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

/// 'a: lifetime of the frame
/// 'b: the borrow of &mut Frame
pub struct Writer<'a, 'b> where 'a: 'b {
    ptr:    *mut u8,
    len:    usize,
    cap:    usize, // in bytes,
    frame:  &'b mut Frame<'a>
}
impl<'a, 'b> Writer<'a, 'b> {
    pub fn finish(self) -> &'a str {
        self.frame.top += self.len;
        
        unsafe {
            let slice = slice::from_raw_parts(self.ptr, self.len);
            str::from_utf8_unchecked(slice)
        }
    }
}
impl<'a, 'b> fmt::Write for Writer<'a, 'b> {
    #[inline(always)]
    fn write_str(&mut self, s: &str) -> Result<(), fmt::Error> {
        if self.len + s.len() <= self.cap {
            unsafe {
                ptr::copy(s.as_ptr(), self.ptr.offset(self.len as isize), s.len());
            }
            self.len += s.len();
            Ok(())
        } else {
            Err(fmt::Error::default())
        }
    }
    #[inline(always)]
    fn write_char(&mut self, c: char) -> Result<(), fmt::Error> {
        if self.len + 4 <= self.cap {
            let mut slice = unsafe {
                slice::from_raw_parts_mut(
                    self.ptr.offset(self.len as isize),
                    self.cap
                )
            };
            self.len += c.encode_utf8(slice).len();
            Ok(())
        } else {
            Err(fmt::Error::default())
        }
    }
}

/// An array that owns the data.
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) }
    }
}
impl<'a, T: fmt::Debug> fmt::Debug for FArray<'a, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        (self as &[T]).fmt(f)
    }
}

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) }
    }
}
impl<'a, T> Deref for FBox<'a, T> {
    type Target = T;
    #[inline(always)]
    fn deref(&self) -> &T {
        unsafe { &*self.ptr }
    }
}
impl<'a, T> DerefMut for FBox<'a, T> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut T {
        unsafe { &mut *self.ptr }
    }
}
impl<'a, T, U> PartialEq<U> for FBox<'a, T> where T: PartialEq<U> {
    fn eq(&self, rhs: &U) -> bool {
        (self as &T).eq(rhs)
    }
}
impl<'a, T: fmt::Debug> fmt::Debug for FBox<'a, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        (self as &T).fmt(f)
    }
}

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, OutOfMemory>
    where F: FnOnce(*mut T) -> O
{
    let mut root = ROOT.get_mut();
    let mut sub = root.subframe();
    Ok(f(sub.try_place_raw(count)?.get()))
}

/// allocate `count` elements of T
#[inline]
pub fn place<F, O, T>(count: usize, f: F) -> Result<O, OutOfMemory>
    where F: FnOnce(&mut [T]) -> O, T: Default
{   
    let mut root = get_root();
    let mut sub = root.subframe();
    let mut arr = sub.try_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, OutOfMemory>
    where I: Iterator<Item=T>, F: FnOnce(&mut FArray<T>) -> O
{   
    let mut root = get_root();
    let mut sub = root.subframe();
    let mut arr = sub.place_iter(it)?;
    Ok(f(&mut arr))
}


/// 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 = get_root();
    let mut sub = root.subframe();
    f(&mut sub)
}

#[test]
fn test_place() {
    frame(|f| {
        let p1 = f.push();
        let p2 = f.push();
        let a = p1 <- "hi";
        let b = p2 <- 42;
        
        assert_eq!(a, "hi");
        assert_eq!(b, 42);
    })
}