#![allow(dead_code)] // rust for prototyping is no great
#![allow(clippy::disallowed_names)]
#![feature(allocator_api)]

#[macro_use]
pub mod string;

pub mod boxed;
mod linked_list;
mod double_list;
pub mod traits;

use std::{cell::Cell, pin::Pin};

// > Allocator is designed to be implemented on ZSTs
//
// says the rustdoc. But we don't give up so easily
pub struct Arena {
    // I'm actually not sure how pin works in rust, but I'm using my intuition
    // from garbage-collected languages that it needs to be here
    mem: Pin<Box<[u8]>>,
    // Since Allocator methods operate on non-mut reference, we put the mutable
    // stuff in cells
    top: Cell<*mut u8>,
    last: Cell<*mut u8>,
}

impl Arena {
    // Not interesting, create arena with the amount of bytes. I'm surprised
    // it's not unsafe, but then again, creating pointers is a safe operation,
    // it's reading them that can cause issues.
    pub fn new(size: usize) -> Self {
        let mem = vec![0; size].into_boxed_slice();
        let mut mem = Pin::new(mem);
        let top = mem.as_mut_ptr();
        Self {
            mem,
            top: Cell::new(top),
            last: Cell::new(std::ptr::null_mut()),
        }
    }

    pub fn share<'a, T>(&'a self, x: T) -> &'a T {
        let b = Box::new_in(x, self);
        let ptr = Box::leak(b);
        ptr
    }

    pub fn reclaim(mut self) -> Self {
        self.top.set(self.mem.as_mut_ptr());
        self.last.set(std::ptr::null_mut());
        self
    }
}

// Less go babyyyyy
unsafe impl std::alloc::Allocator for Arena {
    fn allocate(&self, layout: std::alloc::Layout) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {
        // We're doing a bump allocator because it's easier and faster
        unsafe {
            // Top points to the free memory
            let top = self.top.get();
            // Align top to the desired allocation - this will be the pointer we
            // return
            let this = top.add(top.align_offset(layout.align()));
            // Next top will point to a location just after the new object
            self.top.set(this.add(layout.size()));
            // Now do a fucking NonNull dance. Made more complicated by the fact
            // we have to return a slice not a pointer
            let p = std::ptr::NonNull::new_unchecked(this);
            let p = std::ptr::NonNull::slice_from_raw_parts(p, layout.size());

            // Optimization for reallocating
            self.last.set(this);

            Ok(p)
        }
    }

    unsafe fn shrink(&self, p: std::ptr::NonNull<u8>, _old: std::alloc::Layout, new_layout: std::alloc::Layout) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {
        // We do a small optimization: just shrink it in place
        let r = std::ptr::NonNull::slice_from_raw_parts(p, new_layout.size());
        Ok(r)
    }

    unsafe fn deallocate(&self, _ptr: std::ptr::NonNull<u8>, _layout: std::alloc::Layout) {
        // Since it's a bump allocator, we don't do anything on deallocation.
    }

    unsafe fn grow(
        &self,
        ptr: std::ptr::NonNull<u8>,
        old_layout: std::alloc::Layout,
        new_layout: std::alloc::Layout
    ) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {
        if ptr.as_ptr() == self.last.get() {
            // is it possible that reallocation changes alignment?
            assert_eq!(old_layout.align(), new_layout.align());
            // update top to point after reallocation
            let top = self.top.get();
            self.top.set( top.add(new_layout.size() - old_layout.size()) );

            let p = std::ptr::NonNull::slice_from_raw_parts(ptr, new_layout.size());
            Ok(p)
        } else {
            self.allocate(new_layout)
        }
    }
}

impl Drop for Arena {
    fn drop(&mut self) {
        // Now here's why borrow checker is great: all pointer to the arena
        // memory are annotated with its lifetime, so we can't drop the arena
        // earlier than any pointers to it! Even `Box::leak` gives us a
        // reference with the lifetime of arena. The only thing we risk is not
        // running the destructor, which is perfectly safe.

        // I'm writing this inside a `Drop` implementation because at first I
        // didn't realize this fact and there used to be code here.
    }
}

fn playtime() {
    // Simple but non trivial usage: return pointer from a function
    fn rets() -> Box<i64, Arena> {
        let arena = Arena::new(8);
        Box::new_in(228, arena)
    }

    let number = rets();
    println!("{}", *number); // Follow along and deduce what it prints!

    // A more complicated usage: pass arena as a parameter
    let arena = Arena::new(24);
    fn uses(a: &Arena) {
        let foo = Box::new_in(322_u64, a);
        let bar = Box::new_in("eight ch", a);
        println!("{} {}", *bar, *foo);
    }
    uses(&arena);
    // Follow along: is the type of pointer created here same or different from
    // example above?

    // This more complicated example gives you a hint to the correct answer
    let arena = Arena::new(16);
    fn uses_rets<'a>(a: &'a Arena) -> Box<&'a str, &'a Arena> {
        Box::new_in("oh wee", a)
    }
    let woot = uses_rets(&arena);
    println!("{}", *woot);
}

fn use_arena() {
    let arena = Arena::new(4096);
    let foo = Box::new_in("kekus", &arena);
    let bar = Box::new_in(12, &arena);
    // here's a fucking problem: I can't to_string with a custom allocator.
    // There isn't even new_in for a string

    println!("ha {} ha {}", *foo, *bar);


    let s = format_in!(&arena, "ha {} ha {}", *foo, *bar).unwrap();
    println!("{}", s);
}

fn main() {
    playtime();
    use_arena()
}