//! 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::{bijective_higher_ranked_trait, bijective_higher_ranked_type, hkt::Invariant};
use core::marker::PhantomData;

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

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

bijective_higher_ranked_trait! {
    /// Higher ranked types with a context `'ctx` lifetime.
    ///
    /// Its recommended to name types of this class with a `Dyn` prefix.
    /// This prefix was chosen because most lower types used in treaty are trait objects.
    ///
    /// This is most important for trait objects. For example `dyn MyTrait<'ctx> + 'a` this
    /// type has a "extra" `'a` lifetime for how long its valid for. This higher ranked type class
    /// allows naming the trait object type without a `'a` being given.
    pub type class WithContextLt['ctx][]: {} [for<'a>]
}

bijective_higher_ranked_trait! {
    /// Higher ranked types that are `'static`.
    ///
    /// Types of this class can usually be sealed as they don't need to be named directly.
    ///
    /// Higher ranked types of this form have a [`TypeId`][core::any::TypeId] associated with them.
    /// This allows them to be used as a name for lifetime containing types.
    ///
    /// This type class has members in the [`WithContextLt`] higher ranked type class.
    /// To get a concrete type two lowerings need to be applied to inject two lifetimes.
    /// One for the context lifetime, and one for the lifetime of the concrete type.
    pub type class TypeName[][]: {'static} [for<'ctx> WithContextLt::MemberType<'ctx>]
}

bijective_higher_ranked_type! {
    /// Higher ranked type for borrows `&T`.
    ///
    /// The borrow gets the `'a` lifetime, not the `'ctx` lifetime.
    pub type DynRef['ctx][][T['ctx][]]: WithContextLt['ctx][]
    for<'a>
    (&'a WithContextLt::T<'a, 'ctx, T>)
    (&'a T)
    where {
        T: ?Sized
    }
}

bijective_higher_ranked_type! {
    pub type [][][T[][]]: TypeName[][]
    for<'ctx>
    (DynRef<'ctx, TypeName::T<'ctx, T>>)
    (DynRef<'ctx, T>)
    where {
        T: ?Sized
    }
}

bijective_higher_ranked_type! {
    /// Higher ranked type for mutable borrows `&mut T`.
    ///
    /// The borrow gets the `'a` lifetime, not the `'ctx` lifetime.
    pub type DynMut['ctx][][T['ctx][]]: WithContextLt['ctx][]
    for<'a>
    (&'a mut WithContextLt::T<'a, 'ctx, T>)
    (&'a mut T)
    where {
        T: ?Sized
    }
}

bijective_higher_ranked_type! {
    pub type [][][T[][]]: TypeName[][]
    for<'ctx>
    (DynMut<'ctx, TypeName::T<'ctx, T>>)
    (DynMut<'ctx, T>)
    where {
        T: ?Sized
    }
}

#[cfg(feature = "alloc")]
bijective_higher_ranked_type! {
    /// Higher ranked type for boxes `Box<T>`.
    ///
    /// A [`Box`] doesn't need either lifetime.
    pub type DynBox['ctx][][T['ctx][]]: WithContextLt['ctx][]
    for<'a>
    (Box<WithContextLt::T<'a, 'ctx, T>>)
    (Box<T>)
    where {
        T: ?Sized
    }
}

#[cfg(feature = "alloc")]
bijective_higher_ranked_type! {
    pub type [][][T[][]]: TypeName[][]
    for<'ctx>
    (DynBox<'ctx, TypeName::T<'ctx, T>>)
    (DynBox<'ctx, T>)
    where {
        T: ?Sized
    }
}

/// 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, WithContextLt, TypeName};
/// use treaty::hkt::bijective_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;
/// }
///
/// bijective_higher_ranked_type! {
///     for['ctx] type DynToNum[][]: WithContextLt['ctx][]
///     for<'a> (dyn ToNum + 'a)
/// }
///
/// // Make the trait object nameable.
/// bijective_higher_ranked_type! {
///     type [][]: TypeName[][]
///     for<'ctx> (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::LowerType<'ctx>>(
        &'a self,
    ) -> Option<&'a WithContextLt::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::LowerType<'ctx>>(
        &'a mut self,
    ) -> Option<&'a mut WithContextLt::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::WithContextLt::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::WithContextLt::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 + WithContextLt::LowerType<'a, 'ctx>>(indirect: I::ForT<T>) -> Self
    where
        WithContextLt::HigherRanked<'a, 'ctx, T>: TypeName::LowerType<'ctx>,
    {
        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 + WithContextLt::LowerType<'a, 'ctx>>(
        self,
    ) -> Result<I::ForT<T>, Self>
    where
        WithContextLt::HigherRanked<'a, 'ctx, T>: TypeName::LowerType<'ctx>,
    {
        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 crate::bijective_higher_ranked_type;

    use super::*;

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

        bijective_higher_ranked_type! {
            type DynZ['ctx][]: WithContextLt['ctx][] for<'a> (dyn Z<'ctx> + 'a)
        }

        bijective_higher_ranked_type! {
            type [][]: TypeName[][] for<'ctx> (DynZ<'ctx>)
        }

        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<'ctx>
            ]
        }

        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: &WithContextLt::T<
            'a,
            'ctx,
            TypeName::T<
                'ctx,
                TypeName::HigherRanked<'ctx, WithContextLt::HigherRanked<'a, 'ctx, T>>,
            >,
        >,
    ) where
        T: WithContextLt::LowerType<'a, 'ctx>,
        WithContextLt::HigherRanked<'a, 'ctx, T>: TypeName::LowerType<'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<
            'ctx,
            WithContextLt::HigherRanked<'a, 'ctx, WithContextLt::T<'a, 'ctx, TypeName::T<'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;
    }
}