//! Extended type erasure support.
//!
//! The [`core::any`] module, and [`Any`][core::any::Any] in particular,
//! provide great utilities to perform type erasure. However, there are some
//! limitations. Namely `Any` requires the type be `'static`. This doesn't
//! work for [`treaty`][crate] as it needs the `'ctx` lifetime.
//! Also `Any` doesn't allow type erasing `!Sized` types as this would
//! result in a "double fat pointer"?!, and those don't exist.
//!
//! This module solves both problems. First, [`LtAny`] is a lifetime containing
//! counterpart to [`Any`][core::any::Any]. [`LtAny`] allows one lifetime which 
//! [`treaty`][crate] uses for the `'ctx` lifetime shared by walkers and visitors.
//! Second, [`IndirectLtAny`] is able to type erase a pointer like type even
//! if it's a `!Sized` type. This allows for type erasing borrows of trait objects
//! like `&dyn Trait`.
//!
//! For a type to be compatible with this module it needs to implement [`TypeNameable`]
//! to give it a unique [`TypeId`][core::any::TypeId]. This can be done manually 
//! without unsafe code. However, its recommended to use the provided [`nameable`]
//! macro when possible.

pub mod static_wrapper;

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

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

/// A type with another type acting as its name.
///
/// The `'a` lifetime is the lifetime `Self` must outlive.
/// The `'lt` lifetime is some arbitrary lifetime `Self` can use in it's definition.
/// 
/// The [`nameable`] allows implementing this trait and [`TypeName`] with minimal effort.
/// For types that are `'static` they can be wrapped by the included wrappers in 
/// the [`static_wrapper`] module.
///
/// This trait is circular with [`TypeName`] on [`Self::Name`].
/// As a result, both must be implemented with matching implementations,
/// and each type can only be used as the name for one type.
pub trait TypeNameable<'a, 'lt>: 'a {
    /// The type acting to name `Self`.
    ///
    /// This name type is unique to this `Self` type.
    type Name: ?Sized + TypeName<'a, 'lt, Nameable = Self>;
}

/// The pair trait to [`TypeNameable`].
///
/// This trait is implemented by types acting as names.
/// Each name type can only be used to name one type.
///
/// This type must be `'static` so it works with [`TypeId`][core::any::TypeId].
pub trait TypeName<'a, 'lt>: 'static {
    /// The type this type names.
    ///
    /// In some sense, this is the lifetime poisoned form of `Self`.
    type Nameable: ?Sized + TypeNameable<'a, 'lt, Name = Self>;
}

/// Implement [`TypeNameable`] and generate a unique name type.
///
/// ```
/// use treaty::any::{nameable, TypeNameable, TypeName};
///
/// pub struct MyType<T>(pub T);
///
/// nameable! {
///     pub struct Name['a, 'lt, T];
///     impl [T::Name] for MyType<T> where { T: TypeNameable<'a, 'lt>, T::Name: Sized }
///     impl [T] where MyType<T::Nameable> { T: TypeName<'a, 'lt>, T::Nameable: Sized }
/// }
/// ```
///
/// The generated `Name` struct will not be nameable outside the macro.
/// Its only purpose is to be a unique type to act as a name.
/// The first `impl` is for [`TypeNameable`] on the type given between the `for` and `where`.
/// The first list of generics are what will be passed to the generics of `Name`.
/// The second `impl` is for [`TypeName`] in `Name`. The type given after the `where`
/// needs to match the type given in the first `impl`. However, in the second impl
/// the `T` is a generic for a name not the generic for the type. That's why 
/// `T::Nameable` is used instead of `T`.
#[doc(hidden)]
#[macro_export]
macro_rules! nameable {
    {
        $vis:vis struct $name:ident[$a:lifetime, $lt:lifetime $(, $($generic:ident),* $(,)?)?];

        impl $([$($name_generics:tt)*])? for $type:ty where {$($nameable_bound:tt)*}
    } => {
        $crate::any::nameable! {
            $vis struct $name[$a, $lt $(, $($generic),*)?];

            impl $([$($name_generics)*])? for $type where {$($nameable_bound)*}
            impl $([$($name_generics)*])? where $type {$($nameable_bound)*}
        }
    };
    {
        $vis:vis struct $name:ident[$a:lifetime, $lt:lifetime $(, $($generic:ident),* $(,)?)?];

        impl $([$($name_generics:tt)*])? for $type:ty where {$($nameable_bound:tt)*}

        impl $([$($nameable_generics:tt)*])? where $nameable_type:ty {$($name_bound:tt)*}
    } => {
        const _: () = {
            $vis struct $name $(< $($generic: ?Sized),* >)?(
                ::core::marker::PhantomData<fn() -> ($( $(*const $generic,)* )?)>
            );

            impl<$a, $lt $(, $($generic)*)?>
                $crate::any::TypeNameable<$a, $lt> for $type
            where
                $($nameable_bound)*
            {
                type Name = $name $(<$($name_generics)*>)?;
            }

            impl<$a, $lt $(, $($generic),*)?>
                $crate::any::TypeName<$a, $lt> for Name$(<$($nameable_generics)*>)?
            where
                $($name_bound)*
            {
                type Nameable = $nameable_type;
            }
        };
    };
}
#[doc(inline)]
pub use nameable;

nameable! {
    pub struct Name['a, 'lt, T];
    impl [T::Name] for &'lt T where { 
        T: TypeNameable<'a, 'lt> + ?Sized, T::Name: Sized, 'lt: 'a 
    }
    impl [T] where &'lt T::Nameable { 
        T: TypeName<'a, 'lt>, T::Nameable: 'lt, 'lt: 'a 
    }
}

nameable! {
    pub struct Name['a, 'lt, T];
    impl [T::Name] for &'lt mut T where { 
        T: TypeNameable<'a, 'lt> + ?Sized, T::Name: Sized, 'lt: 'a 
    }
    impl [T] where &'lt mut T::Nameable { 
        T: TypeName<'a, 'lt>, T::Nameable: 'lt, 'lt: 'a 
    }
}

nameable! {
    pub struct Name['a, 'lt, T];
    impl [T::Name] for *const T where { 
        T: TypeNameable<'a, 'lt> + ?Sized, T::Name: Sized
    }
    impl [T] where *const T::Nameable { 
        T: TypeName<'a, 'lt>
    }
}

nameable! {
    pub struct Name['a, 'lt, T];
    impl [T::Name] for *mut T where { 
        T: TypeNameable<'a, 'lt> + ?Sized, T::Name: Sized
    }
    impl [T] where *mut T::Nameable { 
        T: TypeName<'a, 'lt>
    }
}

#[cfg(feature = "alloc")]
nameable! {
    pub struct Name['a, 'lt, T];
    impl [T::Name] for Box<T> where { 
        T: TypeNameable<'a, 'lt> + ?Sized, T::Name: Sized
    }
    impl [T] where Box<T::Nameable> { 
        T: TypeName<'a, 'lt>
    }
}

/// [`TypeId`][core::any::TypeId] with a lifetime generic `'lt`.
///
/// This allows comparing types that contain zero or one lifetimes.
/// When `LtTypeId::of::<A>() == LtTypeId::of::<B>()` then `A` is `B`.
#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Copy, Clone, Debug)]
pub struct LtTypeId<'lt> {
    /// Invariant over `'lt` for the eq check to be correct.
    /// The borrow checker is checking that the lifetimes of the type
    /// IDs are the same instead of doing it at runtime, which we can't do.
    _marker: PhantomData<fn(&'lt ()) -> &'lt ()>,

    /// The type ID of the name type of the type.
    name_id: core::any::TypeId,
}

impl<'lt> LtTypeId<'lt> {
    /// Get the ID of a type.
    ///
    /// The type must implement [`TypeNameable`]. Note, the `'a` lifetime is **not**
    /// tracked by the [`LtTypeId`], only the `'lt` lifetime is.
    pub fn of<'a, T: ?Sized + TypeNameable<'a, 'lt>>() -> Self {
        LtTypeId {
            _marker: PhantomData,
            name_id: core::any::TypeId::of::<T::Name>(),
        }
    }
}

/// [`Any`][core::any::Any] with a lifetime generic `'lt`.
///
/// This trait is implemented on all types that implement
/// [`TypeNameable`]. It is not possible to implement this trait manually.
///
/// Like [`Any`][core::any::Any] this trait can be used as a trait object with
/// downcasting.
pub trait LtAny<'lt>: sealed::Sealed<'lt> {
    /// Get the [`LtTypeId`] of the type of `self`.
    fn type_id(&self) -> LtTypeId<'lt>;
}

// Prevent any impls except the following.
mod sealed {
    use super::*;

    pub trait Sealed<'lt> {}

    impl<'a, 'lt, T: ?Sized + TypeNameable<'a, 'lt>> Sealed<'lt> for T {}
}

impl<'a, 'lt, T: ?Sized + TypeNameable<'a, 'lt>> LtAny<'lt> for T {
    fn type_id(&self) -> LtTypeId<'lt> {
        LtTypeId::of::<T>()
    }
}

impl<'a, 'lt> dyn LtAny<'lt> + 'a {
    /// Check if `self` is of type `T`.
    pub fn is<T: ?Sized + TypeNameable<'a, 'lt>>(&self) -> bool {
        LtTypeId::of::<T>() == self.type_id()
    }

    /// Downcast a `&dyn LtAny<'lt>` into a `&T`.
    pub fn downcast_ref<T: TypeNameable<'a, 'lt>>(&self) -> Option<&T> {
        if self.is::<T>() {
            Some(unsafe { &*(self as *const dyn LtAny<'lt> as *const T) })
        } else {
            None
        }
    }

    /// Downcast a `&mut dyn LtAny<'lt>` into a `&mut T`.
    pub fn downcast_mut<T: TypeNameable<'a, 'lt>>(&mut self) -> Option<&mut T> {
        if self.is::<T>() {
            Some(unsafe { &mut *(self as *mut dyn LtAny<'lt> as *mut T) })
        } else {
            None
        }
    }

    /// Downcast a `Box<dyn LtAny<'lt>>` into a `Box<T>`.
    #[cfg(feature = "alloc")]
    pub fn downcast_box<T: TypeNameable<'a, 'lt>>(self: Box<Self>) -> Result<Box<T>, Box<Self>> {
        if self.is::<T>() {
            Ok(unsafe {
                let raw: *mut dyn LtAny<'lt> = Box::into_raw(self);
                Box::from_raw(raw as *mut T)
            })
        } else {
            Err(self)
        }
    }
}

/// Dynamic trait lookup.
///
/// This trait allows looking up the trait object form of `self` for a 
/// given trait object `id`. This is similar to upcasting to the trait given
/// by `id` if [`AnyTrait`] had every trait as a super bound.
///
/// ```
/// use treaty::any::{AnyTrait, any_trait, nameable};
///
/// // Create a test value.
/// let my_num = MyNum(42);
///
/// // Cast to be a AnyTrait trait object.
/// // Now we don't know the type.
/// let anything: &dyn AnyTrait<'_> = &my_num;
///
/// // We can still upcast to an impl of ToNum.
/// let to_num_object: &dyn ToNum = anything.upcast().unwrap();
///
/// assert_eq!(to_num_object.num(), 42);
///
/// // === Type Setup ===
///
/// // An example trait.
/// trait ToNum {
///     fn num(&self) -> i32;
/// }
///
/// // Make the trait object nameable.
/// nameable! {
///     struct Name['a, 'ctx];
///     impl for dyn ToNum + 'a where { }
/// }
/// 
/// // An example struct.
/// struct MyNum(i32);
///  
/// // The example struct impls the example trait.
/// impl ToNum for MyNum {
///     fn num(&self) -> i32 {
///         self.0
///     }
/// }
///
/// // Allow the example struct's trait impls to be looked up at runtime.
/// // Here the only trait that can be looked up is ToNum as its the only
/// // one in the list.
/// any_trait! {
///     impl['a, 'ctx] MyNum = [dyn ToNum + 'a];
/// }
/// ```
pub trait AnyTrait<'lt> {
    /// Upcast a borrow to the given trait object.
    ///
    /// Use the `<dyn AnyTrait>::upcast()` helper method instead, if possible.
    ///
    /// If `self` doesn't support upcasting to the requested type
    /// then `None` is returned. The returned trait object is type erased so this trait
    /// is object safe.
    fn upcast_to_id<'a>(&'a self, id: LtTypeId<'lt>) -> Option<IndirectLtAny<'a, 'lt, Ref>>
    where
        'lt: 'a;

    /// Upcast a mutable borrow to the given trait object.
    ///
    /// Use the `<dyn AnyTrait>::upcast_mut()` helper method instead, if possible.
    ///
    /// If `self` doesn't support upcasting to the requested type
    /// then `None` is returned. The returned trait object is type erased so this trait
    /// is object safe.
    fn upcast_to_id_mut<'a>(&'a mut self, id: LtTypeId<'lt>) -> Option<IndirectLtAny<'a, 'lt, Mut>>
    where
        'lt: 'a;
}

impl<'lt> dyn AnyTrait<'lt> + '_ {
    /// Upcast a borrow to the given trait object type.
    ///
    /// This should be used instead of [`upcast_to_id`][AnyTrait::upcast_to_id]
    /// as it automatically downcasts the returned [`IndirectLtAny`].
    ///
    /// If the returned [`IndirectLtAny`] is the wrong type, then a panic happens.
    pub fn upcast<'a, Trait: ?Sized + TypeNameable<'a, 'lt>>(&'a self) -> Option<&'a Trait> {
        self.upcast_to_id(LtTypeId::of::<Trait>())
            .map(|object| match object.downcast::<Trait>() {
                Ok(object) => object,
                Err(object) => panic!(
                    "Unexpected trait object. This means a bad impl of \
                    `upcast_to_id`. Expected: {:?}, Got {:?}",
                    LtTypeId::of::<Trait>(),
                    object.id()
                ),
            })
    }

    /// Upcast a mutable borrow to the given trait object type.
    ///
    /// This should be used instead of [`upcast_to_id_mut`][AnyTrait::upcast_to_id]
    /// as it automatically downcasts the returned [`IndirectLtAny`].
    ///
    /// If the returned [`IndirectLtAny`] is the wrong type, then a panic happens.
    pub fn upcast_mut<'a, Trait: ?Sized + TypeNameable<'a, 'lt>>(
        &'a mut self,
    ) -> Option<&'a mut Trait> {
        self.upcast_to_id_mut(LtTypeId::of::<Trait>())
            .map(|object| match object.downcast::<Trait>() {
                Ok(object) => object,
                Err(object) => panic!(
                    "Unexpected trait object. This means a bad impl of \
                    `upcast_to_id_mut`. Expected: {:?}, Got {:?}",
                    LtTypeId::of::<Trait>(),
                    object.id()
                ),
            })
    }
}

/// Implement [`AnyTrait`] for a type.
///
/// This allows looking up trait objects from the provided list.
/// See [`AnyTrait`] for an example.
#[doc(hidden)]
#[macro_export]
macro_rules! any_trait {
    {
        impl[$a:lifetime, $lt:lifetime $($generic:tt)*] $name:ty = [$($protocol:ty),* $(,)?];
    } => {
        impl<$lt $($generic)*> $crate::any::AnyTrait<$lt> for $name {
            #[inline]
            fn upcast_to_id<$a>(
                &$a self,
                id: $crate::any::LtTypeId<$lt>
            ) -> ::core::option::Option<$crate::any::IndirectLtAny<$a, $lt, $crate::any::Ref>>
            where
                $lt: $a
            {
                match id {
                    $(id if id == $crate::any::LtTypeId::of::<$protocol>()
                        => ::core::option::Option::Some($crate::any::IndirectLtAny::<$a, $lt, _>::new::<$protocol>(self as _)),)*
                    _ => ::core::option::Option::None
                }
            }

            #[inline]
            fn upcast_to_id_mut<$a>(
                &$a mut self,
                id: $crate::any::LtTypeId<$lt>
            ) -> ::core::option::Option<$crate::any::IndirectLtAny<$a, $lt, $crate::any::Mut>>
            where
                $lt: $a
            {
                match id {
                    $(id if id == $crate::any::LtTypeId::of::<$protocol>()
                        => ::core::option::Option::Some($crate::any::IndirectLtAny::<$a, $lt, _>::new::<$protocol>(self as _)),)*
                    _ => ::core::option::Option::None
                }
            }
        }
    };
}
#[doc(inline)]
pub use any_trait;

/// A type erased pointer like.
///
/// This acts like a `&dyn LtAny` except it is able to type erase another fat pointer.
/// This allows type erasing pointers to trait objects. A [`IndirectLtAny`] cannot
/// store an instance of itself.
///
/// The `I` generic is the flavor if pointer being used. It can be [`Ref`], [`Mut`], [`Boxed`], or
/// a custom pointer type.
#[must_use]
pub struct IndirectLtAny<'a, 'lt: 'a, I: Indirect<'a>> {
    info: fn() -> (LtTypeId<'lt>, unsafe fn(RawIndirect)),
    indirect: RawIndirect,
    _marker: PhantomData<(I::ForT<fn(&'lt ()) -> &'lt ()>, PhantomPinned, *const ())>,
}

impl<'a, 'lt, I: Indirect<'a>> Drop for IndirectLtAny<'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>> IndirectLtAny<'a, 'lt, I> {
    /// Wrap an indirection.
    ///
    /// The inner type `T` of the indirection is erased.
    pub fn new<T: ?Sized + TypeNameable<'a, 'lt>>(indirect: I::ForT<T>) -> Self {
        Self {
            info: || {
                (LtTypeId::of::<T>(), |raw| {
                    // SAFETY: This is only called in the drop impl.
                    unsafe { drop(I::from_raw::<T>(raw)) }
                })
            },
            indirect: I::into_raw(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: ?Sized + TypeNameable<'a, 'lt>>(self) -> Result<I::ForT<T>, Self> {
        let (id, _) = (self.info)();

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

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

/// A pointer like type.
///
/// For this trait the pointer like type must have the same size as a pointer.
pub unsafe trait Indirect<'a> {
    /// Get the full type for a given `T`.
    type ForT<T: ?Sized + 'a>: 'a;

    /// Convert the pointer into a raw indirection.
    fn into_raw<T: ?Sized + 'a>(value: Self::ForT<T>) -> RawIndirect;

    /// Convert a raw indirection back into the pointer.
    unsafe fn from_raw<T: ?Sized + 'a>(any: RawIndirect) -> Self::ForT<T>;
}

/// An opaque set of bytes the size of a fat pointer.
///
/// Repr wise this is exactly `MaybeUninit<[u8; { size of a fat pointer }]>`.
#[derive(Clone, Copy)]
#[repr(transparent)]
pub struct RawIndirect(MaybeUninit<[u8; INDIRECT_SIZE]>);

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

trait Helper {}

/// Marker type for [`IndirectLtAny`] for a borrow indirection (`&T`).
pub enum Ref {}

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

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

    fn into_raw<T: ?Sized + 'a>(value: Self::ForT<T>) -> RawIndirect {
        unsafe { transmute::<&'a T, RawIndirect>(value) }
    }

    unsafe fn from_raw<T: ?Sized + 'a>(any: RawIndirect) -> Self::ForT<T> {
        unsafe { transmute::<RawIndirect, &'a T>(any) }
    }
}

/// Marker type for [`IndirectLtAny`] for a mutable borrow indirection (`&mut T`).
pub enum Mut {}

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

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

    fn into_raw<T: ?Sized + 'a>(value: Self::ForT<T>) -> RawIndirect {
        unsafe { transmute::<&'a mut T, RawIndirect>(value) }
    }

    unsafe fn from_raw<T: ?Sized + 'a>(any: RawIndirect) -> Self::ForT<T> {
        unsafe { transmute::<RawIndirect, &'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;

    /// Marker type for [`IndirectLtAny`] for a box indirection (`Box<T>`).
    pub enum Boxed {}

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

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

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

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

/// # 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);

    nameable! {
        ['a, 'lt]
        X<'lt> where {'lt: 'a}
        X<'lt> where {'lt: 'a}
    }

    #[test]
    fn implementer_macro() {
        trait Z {}

        nameable! {
            ['a, 'ctx]
            dyn Z + 'a where {'ctx: 'a}
        }

        struct X<T>(T);

        impl<T: Clone> Z for X<T> {}

        any_trait! {
            impl['a, 'ctx, T: Clone] X<T> = [
                dyn Z + 'a
            ];
        }
    }

    #[test]
    fn any_trait_macro() {
        trait Z {
            fn num(&self) -> i32;
        }

        nameable! {
            ['a, 'ctx]
            dyn Z + 'a where {'ctx: 'a}
        }

        struct X(i32);

        impl Z for X {
            fn num(&self) -> i32 {
                self.0
            }
        }

        any_trait! {
            impl['a, 'ctx] X = [
                dyn Z + 'a
            ];
        }

        let x = X(42);
        let y: &dyn AnyTrait<'_> = &x;
        let z: &dyn Z = y.upcast().unwrap();
        assert_eq!(z.num(), 42);
    }
}