//! [`Well`] and [`WellMut`] traits.
#[cfg(feature = "alloc")]
use alloc::{boxed::Box, rc::Rc, sync::Arc};
use core::{
    ops::{Deref, DerefMut},
    pin::Pin,
    ptr::NonNull,
};

/// Place that can store a smart pointer.
///
/// # Safety
///
/// Types that implement this trait assert that:
///
/// 1. The `insert` and `remove` methods are true inverses of each other with no side effects.
/// 2. The pointer to the data is stable on [`Deref`] and [`DerefMut`] (if implemented), *and* this
///    pointer is the same as the one that is returned by [`remove`].
/// 3. The data may be read even when outside the well.
/// 4. If [`DerefMut`] is implemented, the data may be mutated even when outside the well.
///
/// [`insert`]: Well::insert
/// [`remove`]: Well::remove
pub unsafe trait Well: Deref<Target: Sized> {
    /// Put the data back into the well.
    ///
    /// # Safety
    ///
    /// The pointer passed to [`insert`] must previously have been received from [`remove`].
    ///
    /// [`insert`]: Well::insert
    /// [`remove`]: Well::remove
    unsafe fn insert(ptr: NonNull<Self::Target>) -> Self;

    /// Take the data out of the well.
    fn remove(self) -> NonNull<Self::Target>;
}

/// Implementation of [`WellMut`] that uses trait aliases.
#[cfg(feature = "nightly")]
macro_rules! well_mut {
    ($(#[$attr:meta])*) => {
        $(#[$attr])*
        pub trait WellMut = DerefMut + Well;
    }
}

/// Implementation of [`WellMut`] that uses traits.
#[cfg(not(feature = "nightly"))]
macro_rules! well_mut {
    ($(#[$attr:meta])*) => {
        $(#[$attr])*
        pub trait WellMut: DerefMut + Well {}
        impl<T: DerefMut + Well> WellMut for T {}
    }
}

well_mut! {
    /// Alias for [`Well`] + [`DerefMut`].
    ///
    /// Implemented as a trait alias with the `nightly` feature,
    /// and an automatically-implemented trait otherwise.
}

// SAFETY: A pointer in its almost-rawest form.
unsafe impl<T> Well for &T {
    #[inline]
    unsafe fn insert(ptr: NonNull<T>) -> Self {
        // SAFETY: Ensured by caller.
        unsafe { ptr.as_ref() }
    }

    #[inline]
    fn remove(self) -> NonNull<T> {
        NonNull::from(self)
    }
}

// SAFETY: A pointer in its almost-rawest form.
unsafe impl<T> Well for Pin<&T> {
    #[inline]
    unsafe fn insert(ptr: NonNull<T>) -> Self {
        // SAFETY: Ensured by caller.
        unsafe { Pin::new_unchecked(ptr.as_ref()) }
    }

    #[inline]
    fn remove(self) -> NonNull<T> {
        NonNull::from(self.get_ref())
    }
}

// SAFETY: A pointer in its almost-rawest form.
unsafe impl<T> Well for &mut T {
    #[inline]
    unsafe fn insert(mut ptr: NonNull<T>) -> Self {
        // SAFETY: Ensured by caller.
        unsafe { ptr.as_mut() }
    }

    #[inline]
    fn remove(self) -> NonNull<T> {
        NonNull::from(self)
    }
}

// SAFETY: A pointer in its almost-rawest form.
unsafe impl<T> Well for Pin<&mut T> {
    #[inline]
    unsafe fn insert(mut ptr: NonNull<T>) -> Self {
        // SAFETY: Ensured by caller.
        unsafe { Pin::new_unchecked(ptr.as_mut()) }
    }

    #[inline]
    fn remove(self) -> NonNull<T> {
        // SAFETY: We can only get a pinned version back out of the well.
        unsafe { NonNull::from(self.get_unchecked_mut()) }
    }
}

#[cfg(any(test, feature = "alloc"))]
// SAFETY: `Box` just wraps a pointer.
unsafe impl<T> Well for Box<T> {
    #[inline]
    unsafe fn insert(ptr: NonNull<T>) -> Self {
        // SAFETY: Ensured by caller.
        unsafe { Box::from_raw(ptr.as_ptr()) }
    }

    #[inline]
    fn remove(self) -> NonNull<T> {
        // SAFETY: Box's pointer is always non-null.
        unsafe { NonNull::new_unchecked(Box::into_raw(self)) }
    }
}

#[cfg(any(test, feature = "alloc"))]
// SAFETY: If we keep the pin, everything works out.
unsafe impl<T> Well for Pin<Box<T>> {
    #[inline]
    unsafe fn insert(ptr: NonNull<T>) -> Self {
        // SAFETY: Ensured by caller.
        unsafe { Pin::new_unchecked(Box::from_raw(ptr.as_ptr())) }
    }

    #[inline]
    fn remove(self) -> NonNull<T> {
        // SAFETY: Box's pointer is always non-null, and we're not moving anything out of the box.
        unsafe { NonNull::new_unchecked(Box::into_raw(Pin::into_inner_unchecked(self))) }
    }
}

#[cfg(feature = "alloc")]
// SAFETY: `Rc` accounts for the reference-counter offsets in `into_raw` and `from_raw`.
unsafe impl<T> Well for Rc<T> {
    #[inline]
    unsafe fn insert(ptr: NonNull<T>) -> Self {
        // SAFETY: Ensured by caller.
        unsafe { Rc::from_raw(ptr.as_ptr()) }
    }

    #[inline]
    fn remove(self) -> NonNull<T> {
        // SAFETY: Rc's pointer is always non-null.
        unsafe { NonNull::new_unchecked(Rc::into_raw(self).cast_mut()) }
    }
}

#[cfg(feature = "alloc")]
// SAFETY: If we keep the pin, everything works out.
unsafe impl<T> Well for Pin<Rc<T>> {
    #[inline]
    unsafe fn insert(ptr: NonNull<T>) -> Self {
        // SAFETY: Ensured by caller.
        unsafe { Pin::new_unchecked(Rc::from_raw(ptr.as_ptr())) }
    }

    #[inline]
    fn remove(self) -> NonNull<T> {
        // SAFETY: Rc's pointer is always non-null, and we're not moving anything out of the `Rc`.
        unsafe { NonNull::new_unchecked(Rc::into_raw(Pin::into_inner_unchecked(self)).cast_mut()) }
    }
}

#[cfg(feature = "alloc")]
// SAFETY: `Arc` accounts for the reference-counter offsets in `into_raw` and `from_raw`.
unsafe impl<T> Well for Arc<T> {
    #[inline]
    unsafe fn insert(ptr: NonNull<T>) -> Self {
        // SAFETY: Ensured by caller.
        unsafe { Arc::from_raw(ptr.as_ptr()) }
    }

    #[inline]
    fn remove(self) -> NonNull<T> {
        // SAFETY: Arc's pointer is always non-null.
        unsafe { NonNull::new_unchecked(Arc::into_raw(self).cast_mut()) }
    }
}

#[cfg(feature = "alloc")]
// SAFETY: If we keep the pin, everything works out.
unsafe impl<T> Well for Pin<Arc<T>> {
    #[inline]
    unsafe fn insert(ptr: NonNull<T>) -> Self {
        // SAFETY: Ensured by caller.
        unsafe { Pin::new_unchecked(Arc::from_raw(ptr.as_ptr())) }
    }

    #[inline]
    fn remove(self) -> NonNull<T> {
        // SAFETY: Arc's pointer is always non-null, and we're not moving anything out of the `Arc`.
        unsafe { NonNull::new_unchecked(Arc::into_raw(Pin::into_inner_unchecked(self)).cast_mut()) }
    }
}
