//! Extended type erasure support.
//!
//! The `AnyTrait` trait provides dynamic upcasting to trait objects.

pub mod indirect;
mod static_wrapper;
mod type_name_id;

use crate::{
    higher_ranked_trait, higher_ranked_type,
    hkt::{Invariant, Marker},
};
use core::marker::PhantomData;

pub use static_wrapper::*;
pub use type_name_id::*;

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

higher_ranked_trait! {
    pub type class TypeName for<'a, 'ctx> {
        type Bound = &'a &'ctx ();

        type T: { Send + Sync + 'a } where {
            'ctx: 'a
        };

        type HigherRanked: { Send + Sync + 'static };
    }
}

pub struct RefHrt<T: ?Sized>(Marker<T>);

higher_ranked_type! {
    impl TypeName {
        impl['a, 'ctx, T] type T['a, 'ctx] for RefHrt<T> =
            &'a TypeName::T<'a, 'ctx, T>
        where {
            T: ?Sized + TypeName::LowerForLt<'a, 'ctx, TypeName::Bound<'a, 'ctx>>,
            TypeName::T<'a, 'ctx, T>: 'a
        };

        impl['a, 'ctx, T] type HigherRanked['a, 'ctx] for &'a T =
            RefHrt<TypeName::HigherRanked<'a, 'ctx, T>>
        where {
            T: ?Sized + TypeName::LowerType<'a, 'ctx>
        };
    }
}

pub struct MutHrt<T: ?Sized>(Marker<T>);

higher_ranked_type! {
    impl TypeName {
        impl['a, 'ctx, T] type T['a, 'ctx] for MutHrt<T> =
            &'a mut TypeName::T<'a, 'ctx, T>
        where {
            T: ?Sized + TypeName::LowerForLt<'a, 'ctx, TypeName::Bound<'a, 'ctx>>,
            TypeName::T<'a, 'ctx, T>: 'a
        };

        impl['a, 'ctx, T] type HigherRanked['a, 'ctx] for &'a mut T =
            MutHrt<TypeName::HigherRanked<'a, 'ctx, T>>
        where {
            T: ?Sized + TypeName::LowerType<'a, 'ctx>
        };
    }
}

#[cfg(feature = "alloc")]
higher_ranked_type! {
    impl TypeName {
        impl['a, 'ctx, T] type T['a, 'ctx] for Box<T> =
            Box<TypeName::T<'a, 'ctx, T>>
        where {
            T: ?Sized + TypeName::LowerForLt<'a, 'ctx, TypeName::Bound<'a, 'ctx>>,
        };

        impl['a, 'ctx, T] type HigherRanked['a, 'ctx] for Box<T> =
            Box<TypeName::HigherRanked<'a, 'ctx, T>>
        where {
            T: ?Sized + TypeName::LowerType<'a, 'ctx>
        };
    }
}

/// 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, TypeName};
/// use treaty::hkt::higher_ranked_type;
///
/// // 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<'_> + Send) = &my_num;
///
/// // We can still upcast to an impl of ToNum.
/// let to_num_object: &dyn ToNum = anything.upcast::<DynToNum>().unwrap();
///
/// assert_eq!(to_num_object.num(), 42);
///
/// // === Type Setup ===
///
/// // An example trait.
/// trait ToNum {
///     fn num(&self) -> i32;
/// }
///
/// enum DynToNum {}
///
/// higher_ranked_type! {
///     impl TypeName {
///         impl['a, 'ctx] type T['a, 'ctx] for DynToNum =
///             dyn ToNum + Send + Sync + 'a;
///
///         impl['a, 'ctx] type HigherRanked['a, 'ctx] for dyn ToNum + Send + Sync + 'a =
///             DynToNum;
///     }
/// }
///
/// // 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['ctx] MyNum = [DynToNum]
/// }
/// ```
pub trait AnyTrait<'ctx> {
    /// 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: TypeNameId,
    ) -> Option<AnyTraitObject<'a, 'ctx, indirect::Ref>>
    where
        'ctx: '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: TypeNameId,
    ) -> Option<AnyTraitObject<'a, 'ctx, indirect::Mut>>
    where
        'ctx: 'a;
}

impl<'b, 'ctx: 'b> dyn AnyTrait<'ctx> + Send + 'b {
    /// 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 [`AnyTraitObject`].
    ///
    /// If the returned [`AnyTraitObject`] is the wrong type, then a panic happens.
    pub fn upcast<'a, Trait: ?Sized + TypeName::MemberType>(
        &'a self,
    ) -> Option<&'a TypeName::T<'a, 'ctx, Trait>> {
        self.upcast_to_id(TypeNameId::of::<Trait>())
            .map(|object| match object.downcast() {
                Ok(object) => object,
                Err(object) => panic!(
                    "Unexpected trait object. This means a bad impl of \
                    `upcast_to_id`. Expected: {:?}, Got {:?}",
                    TypeNameId::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 [`AnyTraitObject`].
    ///
    /// If the returned [`AnyTraitObject`] is the wrong type, then a panic happens.
    pub fn upcast_mut<'a, Trait: ?Sized + TypeName::MemberType>(
        &'a mut self,
    ) -> Option<&'a mut TypeName::T<'a, 'ctx, Trait>> {
        self.upcast_to_id_mut(TypeNameId::of::<Trait>())
            .map(|object| match object.downcast() {
                Ok(object) => object,
                Err(object) => panic!(
                    "Unexpected trait object. This means a bad impl of \
                    `upcast_to_id_mut`. Expected: {:?}, Got {:?}",
                    TypeNameId::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[$lt:lifetime $($generic:tt)*] $name:ty = [$($protocol:ty),* $(,)?]
        else {
            let $id:ident;
            $($fallback:tt)*
        }
        $(where $($bound:tt)*)?
    } => {
        impl<$lt $($generic)*> $crate::any::AnyTrait<$lt> for $name
        $(where $($bound)*)?
        {
            #[inline]
            fn upcast_to_id<'__>(
                &'__ self,
                id: $crate::any::TypeNameId
            ) -> ::core::option::Option<$crate::any::AnyTraitObject<'__, $lt, $crate::any::indirect::Ref>>
            where
                $lt: '__
            {
                // This match should be optimized well by llvm.
                match id {
                    $(id if id == $crate::any::TypeNameId::of::<$protocol>()
                        => ::core::option::Option::Some($crate::any::AnyTraitObject::<'__, $lt, _>::new::<
                            $crate::any::TypeName::T<'__, $lt, $protocol>
                        >(self as _)),)*
                    $id => {
                        $($fallback)*
                    }
                }
            }

            #[inline]
            fn upcast_to_id_mut<'__>(
                &'__ mut self,
                id: $crate::any::TypeNameId
            ) -> ::core::option::Option<$crate::any::AnyTraitObject<'__, $lt, $crate::any::indirect::Mut>>
            where
                $lt: '__
            {
                // This match should be optimized well by llvm.
                match id {
                    $(id if id == $crate::any::TypeNameId::of::<$protocol>()
                        => ::core::option::Option::Some($crate::any::AnyTraitObject::<'__, $lt, _>::new::<
                            $crate::any::TypeName::T<'__, $lt, $protocol>
                        >(self as _)),)*
                    $id => {
                        $($fallback)*
                    }
                }
            }
        }
    };
    {
        impl[$lt:lifetime $($generic:tt)*] $name:ty = [$($protocol:ty),* $(,)?]
        $(where $($bound:tt)*)?
    } => {
        $crate::any::any_trait! {
            impl[$lt $($generic)*] $name = [$($protocol),*]
            else {
                // Always answer no in the fallback branch if no fallback was given.
                let _id;
                ::core::option::Option::None
            } $(where $($bound)*)?
        }
    }
}
#[doc(inline)]
pub use any_trait;

use self::indirect::{sealed::RawIndirect, Indirect};

/// A double fat pointer.
///
/// This struct wraps a possibly fat pointer of type described by `I`, and adds
/// an additional vtable to allow downcasting to a specific fat pointer.
/// This type is similar to if `&dyn Any` was allowed to itself store unsized types like trait
/// objects.
///
/// The `'a` lifetime is the lifetime of the pointer, and the `'ctx` lifetime is a context lifetime
/// the inner type can use. This type is always invariant over both lifetimes.
///
/// `&'a dyn MyTrait<ctx>` becomes `AnyTraitObject<'a, 'ctx, Ref>`.
///
/// The `I` generic is the flavor if pointer being used. It can be [`Ref`][indirect::Ref] or [`Mut`][indirect::Mut].
#[must_use]
pub struct AnyTraitObject<'a, 'ctx: 'a, I: Indirect<'a>> {
    /// The extra vtable pointer.
    ///
    /// The TypeNameId gives the TypeId of the T's type name.
    /// This is unique per T minus the 'a and 'ctx lifetimes.
    /// which means a `dyn Trait<'ctx> + 'a` can have a TypeNameId.
    ///
    /// The unsafe function is the drop impl for the pointer.
    /// It must only be called once per RawIndirect value, and the value must not be used after the
    /// call. Only a RawIndirect of the correct I type must be passed.
    info: fn() -> TypeNameId,

    /// The inner pointer value.
    ///
    /// This is some form of &T where T may be sized or not.
    indirect: RawIndirect<'a, I>,

    _lifetime: Invariant<'ctx>,
    _not_send_sync: PhantomData<*const ()>,
}

impl<'a, 'ctx, I: Indirect<'a>> AnyTraitObject<'a, 'ctx, I> {
    /// Type erase a pointer.
    ///
    /// `T` doesn't need to be [`Sized`]. As such, a fat pointer can be passed to this function.
    pub fn new<T: ?Sized + TypeName::LowerType<'a, 'ctx>>(indirect: I::ForT<T>) -> Self {
        Self {
            info: TypeNameId::of_lower::<T>,
            indirect: RawIndirect::new(indirect),
            _lifetime: Default::default(),
            _not_send_sync: 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 + TypeName::LowerType<'a, 'ctx>>(self) -> Result<I::ForT<T>, Self> {
        if self.id() == TypeNameId::of_lower::<T>() {
            // SAFETY: We know that the type name type is unique per T because it is bijective.
            // A self is only made in Self::new where the info is taken from T.
            // If the check above passes then we know T must be the same minus the lifetimes.
            // We know the lifetime 'ctx is correct because Self is invariant over it.
            // RawIndirect makes sure that the 'a is correct by being invariant over it.
            //
            // See the tests at the bottom of the file for a proof that the type name is bijective
            // to T.
            Ok(unsafe { self.indirect.into_inner::<T>() })
        } else {
            Err(self)
        }
    }

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

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

    #[test]
    fn any_trait_macro_implements_the_trait() {
        trait Z<'ctx> {
            fn get(&self) -> i32;
        }

        struct DynZ;

        higher_ranked_type! {
            impl TypeName {
                impl['a, 'ctx] type T['a, 'ctx] for DynZ =
                    dyn Z<'ctx> + Send + Sync + 'a;

                impl['a, 'ctx] type HigherRanked['a, 'ctx] for dyn Z<'ctx> + Send + Sync + 'a =
                    DynZ;
            }
        }

        struct X<'ctx>(&'ctx i32);

        impl<'ctx> Z<'ctx> for X<'ctx> {
            fn get(&self) -> i32 {
                *self.0
            }
        }

        any_trait! {
            impl['ctx] X<'ctx> = [
                DynZ
            ]
        }

        let z = 42;
        let x = X(&z);
        let y = (&x as &(dyn AnyTrait<'_> + Send)).upcast::<DynZ>().unwrap();
        assert_eq!(y.get(), 42);
    }

    // The following proves that the higher ranked types are bijective using the type system.
    //
    // We have the type tower: T<'a, 'ctx> <-> DynT<'ctx> <-> NameT
    // We want every T, DynT, NameT set to be unique.
    //
    // Assume there was a U that tried to use NameT in it's type tower:
    // U<'a, 'ctx> <-> DynU<'ctx> <-> NameT
    //
    // If we traverse the type tower in this order: T -r-> A -r-> B -l-> C -l-> D
    // where -r-> is a raise and -l-> is a lower, then if D is always T we know that no sequence
    // U -r-> A2 -r-> B -l-> C2 -l-> D can exist exept where T == U. This is because B cannot
    // have information about where it came from and still be the same B in the type system.
    // The following makes sure that a U could never become a T by the raise then lower process.

    // This proves that the bijective type names are really bijective.
    fn _is_bijective_raise<'a, 'ctx: 'a, T>(
        x: &TypeName::T<'a, 'ctx, TypeName::HigherRanked<'a, 'ctx, T>>,
    ) where
        T: TypeName::LowerType<'a, 'ctx>,
    {
        // If C -> B -> A -> B -> C (shown by this assignment), then C and A must be bijective.
        let _y: &T = x;
    }

    // This proves that the bijective type names are really bijective.
    fn _is_bijective_lower<'a, 'ctx: 'a, U>(
        x: &TypeName::HigherRanked<'a, 'ctx, TypeName::T<'a, 'ctx, U>>,
    ) where
        U: TypeName::MemberType,
    {
        // If A -> B -> C -> B -> A (shown by this assignment), then A and C must be bijective.
        let _y: &U = x;
    }
}