//! Generic over borrow mutability.

use core::mem::{ManuallyDrop, MaybeUninit};

/// A higher kinded type.
pub trait HigherKinded<'a> {
    /// The type with `T` applied.
    type ForT<T: ?Sized + 'a>: Unpin + 'a;
}

/// A pointer like type.
///
/// Allows being generic over `&T` and `&mut T`.
///
/// Implemented by marker types [`Ref`] and [`Mut`].
pub trait Indirect<'a>: HigherKinded<'a> + sealed::Sealed<'a> {}

pub(super) mod sealed {
    use crate::hkt::{Invariant, Marker};

    use super::*;

    /// Sealed trait to prevent users from implementing Indirect and provides conversion
    /// to and from a raw indirect.
    pub trait Sealed<'a>: HigherKinded<'a> + Sized {
        /// Convert the pointer into a raw indirection.
        fn into_raw<T: ?Sized + 'a>(value: Self::ForT<T>) -> MaybeUninit<[u8; INDIRECT_SIZE]>;

        /// Convert a raw indirection back into the pointer.
        ///
        /// # Safety
        /// `value` must have been created by `Self::into_raw`.
        /// This function must not be called twice for the same `value`.
        /// The type `T` must be the one given to `Self::into_raw`.
        #[allow(unsafe_code)]
        unsafe fn from_raw<T: ?Sized + 'a>(
            value: MaybeUninit<[u8; INDIRECT_SIZE]>,
        ) -> 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 }]>`.
    #[repr(transparent)]
    pub struct RawIndirect<'a, I> {
        indirect: MaybeUninit<[u8; INDIRECT_SIZE]>,
        _lifetime: Invariant<'a>,
        _generic: Marker<I>,
    }

    impl<'a, I: Indirect<'a>> RawIndirect<'a, I> {
        pub fn new<T: ?Sized>(indirect: I::ForT<T>) -> Self {
            Self {
                indirect: I::into_raw(indirect),
                _lifetime: Default::default(),
                _generic: Default::default(),
            }
        }

        /// # Safety
        /// The type `T` must be the same one used in `Self::new`.
        #[allow(unsafe_code)]
        pub unsafe fn into_inner<T: ?Sized>(self) -> I::ForT<T> {
            // SAFETY: indirect was created with this I's into_raw in Self::new.
            // This function cannot be called twice because we take ownership of self and
            // this type is not clone or copy.
            // The caller makes sure that the T is the same.
            // Also the lifetime 'a must be the same because Self is invariant over it.
            unsafe { I::from_raw(self.indirect) }
        }
    }

    /// Size of a fat pointer.
    const INDIRECT_SIZE: usize = core::mem::size_of::<usize>() * 2;

    /// Helper trait for double checking the sizes of fat pointers.
    trait _Helper {
        fn run(&self);
    }

    const _: () = assert!(
        core::mem::size_of::<&dyn _Helper>() <= core::mem::size_of::<RawIndirect<'_, Ref>>()
    );

    // Borrow doesn't need to be dropped.
    const _: () = assert!(!core::mem::needs_drop::<&i32>());

    const _: () = assert!(
        core::mem::size_of::<&mut dyn _Helper>()
            <= core::mem::size_of::<sealed::RawIndirect<'_, Mut>>()
    );

    // Mutable borrow doesn't need to be dropped.
    const _: () = assert!(!core::mem::needs_drop::<&mut i32>());

    impl<'a> Sealed<'a> for Ref {
        fn into_raw<T: ?Sized + 'a>(value: Self::ForT<T>) -> MaybeUninit<[u8; INDIRECT_SIZE]> {
            // SAFETY: A possibly fat borrow can be stores in a 2 usize wide maybe uninit array.
            #[allow(unsafe_code)]
            unsafe {
                transmute::<&'a T, MaybeUninit<[u8; INDIRECT_SIZE]>>(value)
            }
        }

        #[allow(unsafe_code)]
        unsafe fn from_raw<T: ?Sized + 'a>(any: MaybeUninit<[u8; INDIRECT_SIZE]>) -> Self::ForT<T> {
            // SAFETY: We know the value is from Self::into_raw because of the caller invariants.
            unsafe { transmute::<MaybeUninit<[u8; INDIRECT_SIZE]>, &'a T>(any) }
        }
    }

    impl<'a> Sealed<'a> for Mut {
        fn into_raw<T: ?Sized + 'a>(value: Self::ForT<T>) -> MaybeUninit<[u8; INDIRECT_SIZE]> {
            // SAFETY: A possibly fat borrow can be stores in a 2 usize wide maybe uninit array.
            #[allow(unsafe_code)]
            unsafe {
                transmute::<&'a mut T, MaybeUninit<[u8; INDIRECT_SIZE]>>(value)
            }
        }

        #[allow(unsafe_code)]
        unsafe fn from_raw<T: ?Sized + 'a>(any: MaybeUninit<[u8; INDIRECT_SIZE]>) -> Self::ForT<T> {
            // SAFETY: We know the value is from Self::into_raw because of the caller invariants.
            unsafe { transmute::<MaybeUninit<[u8; INDIRECT_SIZE]>, &'a mut T>(any) }
        }
    }
}

/// Marker type for [`AnyTraitObject`](super::AnyTraitObject) for a borrow `&T`.
pub enum Ref {}

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

impl<'a> Indirect<'a> for Ref {}

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

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

impl<'a> Indirect<'a> for Mut {}

/// # Safety
/// Same rules as [`core::mem::transmute()`].
#[allow(unsafe_code)]
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::*;

    use sealed::RawIndirect;

    #[test]
    fn can_type_erase_borrow() {
        let x = 42;
        let y: &i32 = &x;
        let z = RawIndirect::<Ref>::new(y);

        // SAFETY: Same type as y which we made it from.
        #[allow(unsafe_code)]
        let w: &i32 = unsafe { z.into_inner() };

        assert_eq!(w, y);
    }

    trait AsInt {
        fn as_int(&self) -> i32;
    }

    impl AsInt for i32 {
        fn as_int(&self) -> i32 {
            *self
        }
    }

    #[test]
    fn can_type_erase_unsized_borrow() {
        let x = 42;
        let y: &dyn AsInt = &x;
        let z = RawIndirect::<Ref>::new(y);

        // SAFETY: Same type as y which we made it from.
        #[allow(unsafe_code)]
        let w: &i32 = unsafe { z.into_inner() };

        assert_eq!(w.as_int(), x);
    }

    #[test]
    fn can_type_erase_mut_borrow() {
        let mut x = 42;
        let y: &mut i32 = &mut x;
        let z = RawIndirect::<Mut>::new(y);

        // SAFETY: Same type as y which we made it from.
        #[allow(unsafe_code)]
        let w: &mut i32 = unsafe { z.into_inner() };

        *w += 1;
        assert_eq!(x, 43);
    }

    trait AddOne {
        fn add_one(&mut self);
    }

    impl AddOne for i32 {
        fn add_one(&mut self) {
            *self += 1;
        }
    }

    #[test]
    fn can_type_erase_unsized_mut_borrow() {
        let mut x = 42;
        let y: &mut dyn AddOne = &mut x;
        let z = RawIndirect::<Mut>::new(y);

        // SAFETY: Same type as y which we made it from.
        #[allow(unsafe_code)]
        let w: &mut dyn AddOne = unsafe { z.into_inner() };

        w.add_one();

        assert_eq!(x, 43);
    }
}