UniqueBox test

use core::fmt;
use std::{
    cell::Cell,
    mem,
    ops::Deref,
    ptr::{NonNull, self},
};

use roc_std::{RocBox, Storage};

use super::{roc_alloc, roc_dealloc};

#[repr(C)]
pub struct UniqueBox<T> {
    contents: NonNull<T>,
}

impl<T> UniqueBox<T> {
    pub fn new(contents: T) -> Self {

        let contents_ptr = Self::alloc_self();
        let contents = unsafe {
            // Write the data to the memory.
            *contents_ptr = contents;
            // No need to check if the ptr is non null it is checked in alloc_self.
            NonNull::new_unchecked(contents_ptr)
        };

        Self { contents }
    }

    // Make the allocation necessary for the data
    fn alloc_self() -> *mut T {
        let alignment = Self::alloc_alignment();
        let bytes = mem::size_of::<T>() + alignment;
        // Initialize the reference count.
        let refcount_one = Storage::new_reference_counted();

        unsafe {
           let ptr = roc_alloc(bytes, alignment as u32);
            if ptr.is_null() {
                todo!("Call roc_panic with the info that an allocation failed");
            }
            // Write the ref count.
            ptr.cast::<Storage>().write(refcount_one);
            // Return the pointer to the data.
            ptr.cast::<u8>().add(alignment).cast::<T>()
        }
    }

    #[inline(always)]
    fn alloc_alignment() -> usize {
        mem::align_of::<T>().max(mem::align_of::<Storage>())
    }

    pub fn into_inner(self) -> T {
        unsafe {ptr::read(self.contents.as_ptr() as *const T)}
    }

    fn storage(&self) -> &Cell<Storage> {
        let alignment = Self::alloc_alignment();
        unsafe {
            &*self
                .contents
                .as_ptr()
                .cast::<u8>()
                .sub(alignment)
                .cast::<Cell<Storage>>()
        }
    }
}

unsafe impl<T> Send for UniqueBox<T> {}

impl<T> Deref for UniqueBox<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        unsafe { self.contents.as_ref() }
    }
}

impl<T> Eq for UniqueBox<T> where T: Eq {}

impl<T, U> PartialEq<UniqueBox<U>> for UniqueBox<T> where T: PartialEq<U> {
    fn eq(&self, other: &UniqueBox<U>) -> bool {
        self.deref() == other.deref()
    }
}

impl<T, U> PartialOrd<UniqueBox<U>> for UniqueBox<T> where T: PartialOrd<U> {
    fn partial_cmp(&self, other: &UniqueBox<U>) -> Option<std::cmp::Ordering> {
        let self_contents = unsafe { self.contents.as_ref() };
        let other_contents = unsafe { other.contents.as_ref() };
        self_contents.partial_cmp(other_contents)
    }
}

impl<T> Ord for UniqueBox<T> where T: Ord {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        let self_contents = unsafe { self.contents.as_ref() };
        let other_contents = unsafe { other.contents.as_ref() };
        self_contents.cmp(other_contents)
    }
}

impl<T> fmt::Debug for UniqueBox<T>
where
    T: fmt::Debug,
{
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        self.deref().fmt(f)
    }
}

impl<T> Clone for UniqueBox<T> {
    fn clone(&self) -> Self {
        let contents_ptr = Self::alloc_self();
        let contents = unsafe {
            *contents_ptr = ptr::read(self.contents.as_ptr());
            NonNull::new_unchecked(contents_ptr)
        };
        Self { contents }
    }
}

impl<T> From<RocBox<T>> for UniqueBox<T> {
    fn from(value: RocBox<T>) -> Self {
        let alignment = Self::alloc_alignment();
        // Get the inner value
        let inner = value.into_inner();
        // Get the storage of the RocBox. (Not sure it works)
        let value_storage: &Cell<Storage> = {
            unsafe {
                &*(&inner as *const T)
                    .cast::<u8>()
                    .sub(alignment)
                    .cast::<Cell<Storage>>()
            }
        };
        let new_storage = value_storage.get();
        // Check the reference count. If it is one we take ownership
        // else we duplicate the value and create a new reference
        // count at one.
        let contents = if new_storage.is_unique() {
            value_storage.set(new_storage);
            NonNull::new(&mut inner as *mut T).unwrap()
        } else {
            let contents_ptr = Self::alloc_self();
            unsafe {
                *contents_ptr = ptr::read(value.deref() as *const T);
                NonNull::new_unchecked(contents_ptr)
            }
        };

        Self { contents }
    }
}

impl<T> Into<RocBox<T>> for UniqueBox<T> {
    fn into(self) -> RocBox<T> {
        // Intrenal representation is equivalent between the two types
        // so it should be fine.
        unsafe {
            mem::transmute(self)
        }
    }
}

impl<T> Drop for UniqueBox<T> {
    fn drop(&mut self) {
        // We don't need to check the reference count because
        // it is assumed that it is one of UniqueBox.
        let contents = self.contents;
        let alignment = Self::alloc_alignment();
        unsafe {
            let contents_ptr = contents.as_ptr();
            mem::drop::<T>(ptr::read(contents_ptr));
            // Release the memory.
            roc_dealloc(contents.as_ptr().cast::<u8>().sub(alignment).cast(), alignment as u32);
        }
    }
}