//! Heapless type erasure.
//!
//! [`Any`] is generic over the type of indirection (`&T`, `&mut T`, `Box<T>`)
//! and allows erasing the `T` of the indirection. This is similar to replacing the `T` with
//! [`core::any::Any`]. The main difference to [`core::any::Any`] is 

use core::{any::TypeId, marker::{PhantomData, PhantomPinned}, mem::{MaybeUninit, ManuallyDrop}};

pub trait TypeNameable<'lt> {
    type Name: TypeName<'lt, Nameable = Self>;
}

pub trait TypeName<'lt>: 'static {
    type Nameable: ?Sized + TypeNameable<'lt, Name = Self>;
}

const INDIRECT_SIZE: usize = core::mem::size_of::<usize>() * 2;

#[derive(Clone, Copy)]
#[repr(transparent)]
pub struct RawIndirectAny(MaybeUninit<[u8; INDIRECT_SIZE]>);

pub unsafe trait Indirect<'a> {
    type ForT<T: ?Sized + 'a>: 'a;

    fn into_any<T: ?Sized + 'a>(value: Self::ForT<T>) -> RawIndirectAny;

    unsafe fn from_any<T: ?Sized + 'a>(any: RawIndirectAny) -> Self::ForT<T>;
}

trait Helper {}

/// A `&T` indirection.
pub enum Ref {}

const _: () = assert!(core::mem::size_of::<&dyn Helper>() <= core::mem::size_of::<RawIndirectAny>());

unsafe impl<'a> Indirect<'a> for Ref {
    type ForT<T: ?Sized + 'a> = &'a T;

    fn into_any<T: ?Sized + 'a>(value: &'a T) -> RawIndirectAny {
        unsafe { transmute::<&'a T, RawIndirectAny>(value) }
    }

    unsafe fn from_any<T: ?Sized + 'a>(any: RawIndirectAny) -> &'a T {
        unsafe { transmute::<RawIndirectAny, &'a T>(any) }
    }
}

/// A `&mut T` indirection.
pub enum Mut {}

const _: () = assert!(core::mem::size_of::<&mut dyn Helper>() <= core::mem::size_of::<RawIndirectAny>());

unsafe impl<'a> Indirect<'a> for Mut {
    type ForT<T: ?Sized + 'a> = &'a mut T;

    fn into_any<T: ?Sized + 'a>(value: &'a mut T) -> RawIndirectAny {
        unsafe { transmute::<&'a mut T, RawIndirectAny>(value) }
    }

    unsafe fn from_any<T: ?Sized + 'a>(any: RawIndirectAny) -> &'a mut T {
        unsafe { transmute::<RawIndirectAny, &'a mut T>(any) }
    }
}

#[cfg(feature = "alloc")]
pub use boxed::*;
#[cfg(feature = "alloc")]
mod boxed {
    use super::*;

    #[cfg(not(feature = "std"))]
    use alloc::boxed::Box;

    /// A `Box<T>` indirection.
    pub enum Boxed {}

    const _: () = assert!(core::mem::size_of::<Box<dyn Helper>>() <= core::mem::size_of::<RawIndirectAny>());

    unsafe impl<'a> Indirect<'a> for Boxed {
        type ForT<T: ?Sized + 'a> = Box<T>;

        fn into_any<T: ?Sized + 'a>(value: Box<T>) -> RawIndirectAny {
            unsafe { transmute::<Box<T>, RawIndirectAny>(value) }
        }

        unsafe fn from_any<T: ?Sized + 'a>(any: RawIndirectAny) -> Box<T> {
            unsafe { transmute::<RawIndirectAny, Box<T>>(any) }
        }
    }
}

/// Container for (almost) any indirection.
///
/// An indirection is something like a `&T` that is a pointer.
/// The `'a` lifetime is the lifetime of the indirection, and `'lt`
/// is a lifetime the stored value's `T` can contain.
///
/// This type is designed to allow using a generic containing type with
/// a trait object method. While the value is stored in the [`Any`] it cannot
/// be accessed. The [`Any`] must be downcasted with [`Any::downcast()`] to access
/// the value.
#[must_use]
pub struct Any<'a, 'lt: 'a, I: Indirect<'a>> {
    /// The meta information about the value is stored as a function pointer
    /// to reduce it's size to one pointer.
    info: fn() -> (TypeId, unsafe fn(RawIndirectAny)),

    /// The indirect pointer.
    indirect: RawIndirectAny,

    /// Invariant over `'lt` and holding a `I::ForT`.
    _marker: PhantomData<(I::ForT<()>, *mut &'lt (), PhantomPinned)>,
}

impl<'a, 'lt, I: Indirect<'a>> Drop for Any<'a, 'lt, I> {
    fn drop(&mut self) {
        // We need to drop the stored value.
        
        // Lookup drop function.
        let (_, drop_fn) = (self.info)();

        // SAFETY: self.indirect is never touched again.
        // Additionally, we know that drop_fn is for this self.indirect because it was
        // made by Self::new.
        unsafe { drop_fn(self.indirect) };
    }
}

impl<'a, 'lt, I: Indirect<'a>> Any<'a, 'lt, I> {
    /// Wrap an indirection.
    ///
    /// The inner type `T` of the indirection is erased.
    pub fn new<T: TypeNameable<'lt>>(indirect: I::ForT<T>) -> Self {
        Self {
            info: || (
                TypeId::of::<T::Name>(), 
                |raw| {
                    // SAFETY: This is only called in the drop impl.
                    unsafe { drop(I::from_any::<T>(raw))
                }}
            ),
            indirect: I::into_any(indirect),
            _marker: PhantomData,
        }
    }

    /// Downcast to an indirection with a given `T` type.
    ///
    /// If the type of the stored value is different, then `self` is
    /// returned as is.
    pub fn downcast<T: TypeNameable<'lt>>(self) -> Result<I::ForT<T>, Self> {
        let (id, _) = (self.info)();

        if id == TypeId::of::<T::Name>() {
            Ok(unsafe { I::from_any(self.indirect) })
        } else {
            Err(self)
        }
    }

    /// Type ID of the stored value's `T`.
    pub fn id(&self) -> TypeId {
        (self.info)().0
    }
}

/// # Safety
/// Same rules as [`core::mem::transmute()`].
unsafe fn transmute<T, U>(value: T) -> U {
    // Create union type that can store a `T` or a `U`.
    // We can then use this to convert between them.
    //
    // The repr(C) layout forces no offset between `t` and `u` as talked about here
    // https://rust-lang.github.io/unsafe-code-guidelines/layout/unions.html#c-compatible-layout-repr-c
    #[repr(C)]
    union Transmute<T, U> {
        t: ManuallyDrop<T>,
        u: ManuallyDrop<U>,
    }

    // Create the union in the `T` state.
    let value = Transmute {
        t: ManuallyDrop::new(value),
    };

    // Read from the union in the `U` state.
    // SAFETY: This is safe because the caller has promised that `T` can be transmuted to `U`.
    // The following reference link talks about repr(C) unions being used this way.
    // https://doc.rust-lang.org/reference/items/unions.html#reading-and-writing-union-fields
    ManuallyDrop::into_inner(unsafe { value.u })
}

#[cfg(test)]
mod test {
    use super::*;

    #[derive(Debug, PartialEq)]
    struct X<'a>(&'a mut i32);

    enum Y {}

    impl<'a> TypeNameable<'a> for X<'a> {
        type Name = Y;
    }

    impl<'a> TypeName<'a> for Y {
        type Nameable = X<'a>;
    }

    #[test]
    fn any() {
        let mut x = 42;
        let x = X(&mut x);

        let any = Any::<Ref>::new(&x);

        let Ok(y) = any.downcast() else { panic!() };

        assert_eq!(x, *y);
    }

    #[test]
    #[cfg(feature = "alloc")]
    fn any_box_drop() {
        #[cfg(not(feature = "std"))]
        use alloc::boxed::Box;

        let mut x = 42;
        let x = X(&mut x);

        let _ = Any::<Boxed>::new(Box::new(x));
    }
}