// SelfRef - Pain-free self-referential pinned types
// Copyright (C) 2022 Soni L.
// This software is made with love by a queer trans person.
// With help from quinedot
//
// SPDX-License-Identifier: MIT OR Apache-2.0

#![feature(generic_associated_types)]
#![no_std]
//! Pain-free self-referential pinned types.
//!
//! This crate (currently) requires nightly, as it relies on GATs.
//!
//! # Example
//!
//! ```rust
//! #![feature(generic_associated_types)]
//! #![feature(pin_macro)]
//!
//! use core::cell::Cell;
//! use core::marker::PhantomData;
//! use core::pin::Pin;
//! use core::pin::pin;
//! 
//! use selfref::Holder;
//! use selfref::NewWith;
//! use selfref::OperateIn;
//! use selfref::opaque;
//!
//! struct Foo<'a, 'b: 'a> {
//!     foo: Cell<Option<&'a Foo<'a, 'b>>>,
//!     t: &'b str,
//! }
//!
//! struct FooKey<'b>(PhantomData<&'b str>);
//! opaque! {
//!     impl['b] Opaque for FooKey<'b> {
//!         type Kind<'a> = Foo<'a, 'b>;
//!     }
//! }
//!
//! fn main() {
//!     struct FooBuilder<'b>(&'b str);
//!     impl<'k, 'b: 'k> NewWith<'k, FooKey<'b>> for FooBuilder<'b> {
//!         fn new_with<'a>(self) -> Foo<'a, 'b> where 'k: 'a {
//!             Foo {
//!                 foo: Default::default(),
//!                 t: self.0,
//!             }
//!         }
//!     }
//!     let stack_array: [u8; 5] = *b"hello";
//!     // a non-'static &str
//!     let stack_str = core::str::from_utf8(&stack_array).unwrap();
//!     let holder = pin!(Holder::new_with(FooBuilder(stack_str)));
//!     // Actually making our Foo refer to itself.
//!     struct SetFooRef;
//!     impl<'k, 'b: 'k> OperateIn<'k, FooKey<'b>> for SetFooRef {
//!         type Out = ();
//!         fn operate_in<'a>(self, foo: Pin<&'a Foo<'a, 'b>>)
//!         where 'k: 'a {
//!             foo.foo.set(Some(foo.get_ref()));
//!         }
//!     }
//!     holder.as_ref().operate_in(SetFooRef);
//! }
//! ```

use core::marker::PhantomPinned;
use core::pin::Pin;
use core::mem;

/// An opaqueified self-referential struct "key".
///
/// # Safety
///
/// This is unsafe because there are a bunch of soundness invariants that need
/// to be upheld. The following list is non-exhaustive:
///
/// - `Kind` must not have a `Drop` impl in any "path" that may trace back to
///     the original self-referential type, if said `Drop` impl can observe
///     the self-referential type.
/// - We assume `Kind` has the same layout for any `'a`. This is true as of the
///     time of writing this, and relies on Rust not having lifetime
///     specialization.
///
/// # Examples
///
/// ```rust
/// #![feature(generic_associated_types)]
///
/// use core::cell::Cell;
/// 
/// use selfref::Opaque;
///
/// struct Foo<'a> {
///     foo: Cell<Option<&'a Foo<'a>>>,
/// }
///
/// struct FooKey;
/// // SAFETY: Foo has no Drop impl and has the same layout for any 'a.
/// unsafe impl Opaque for FooKey {
///     type Kind<'a> = Foo<'a>;
/// }
/// ```
pub unsafe trait Opaque {
    /// The actual self-referential struct.
    type Kind<'a>: ?Sized where Self: 'a;
    #[doc(hidden)]
    fn ub_check() {
    }
}

/// Creates an opaqueified self-referential struct "key".
///
/// Safe wrapper around [`Opaque`] that checks the soundness requirements at
/// compile-time.
///
/// There are 2 forms of this macro. The second form accepts type parameters.
///
/// Note that where bounds go after the impl block.
///
/// # Examples
///
/// Simple example:
///
/// ```rust
/// #![feature(generic_associated_types)]
///
/// use core::cell::Cell;
/// 
/// use selfref::opaque;
///
/// struct Foo<'a> {
///     foo: Cell<Option<&'a Foo<'a>>>,
/// }
///
/// struct FooKey;
/// opaque! {
///     impl Opaque for FooKey {
///         type Kind<'a> = Foo<'a>;
///     }
/// }
/// ```
///
/// Type parameters and where bounds:
///
/// ```rust
/// #![feature(generic_associated_types)]
///
/// use core::cell::Cell;
/// use core::fmt::Display;
/// use core::marker::PhantomData;
/// 
/// use selfref::opaque;
///
/// struct Foo<'a, T: Display> {
///     foo: Cell<Option<&'a Foo<'a, T>>>,
///     t: T,
/// }
///
/// struct FooKey<T>(PhantomData<T>);
/// opaque! {
///     impl[T] Opaque for FooKey<T> {
///         type Kind<'a> = Foo<'a, T>;
///     } where T: Display
/// }
/// ```
#[macro_export]
macro_rules! opaque {
    (
        impl Opaque for $key:ty {
            type Kind<$l:lifetime> = $kind:ty;
        } $(where $($bounds:tt)*)?
    ) => {
        unsafe impl $crate::Opaque for $key $(where $($bounds)*)? {
            type Kind<$l> = $kind where Self: $l;
            fn ub_check() {
                fn ub_detect_helper(
                    _f: impl ::core::ops::Fn(
                        for<$l> fn([&$l (); 0]) -> $kind,
                        for<$l> fn(&$l $kind)
                    )
                ) $(where $($bounds)*)? {
                }
                ub_detect_helper(|f, g| {
                    let _foo: Self::Kind<'_> = f([]);
                    g(&_foo);
                });
            }
        }
    };
    (
        impl[$($params:tt)+] Opaque for $key:ty {
            type Kind<$l:lifetime> = $kind:ty;
        } $(where $($bounds:tt)*)?
    ) => {
        unsafe impl<$($params)+> $crate::Opaque for $key
        $(where $($bounds)*)? {
            type Kind<$l> = $kind where Self: $l;
            fn ub_check() {
                fn ub_detect_helper<$($params)+>(
                    _f: impl ::core::ops::Fn(
                        for<$l> fn([&$l (); 0]) -> $kind,
                        for<$l> fn(&$l $kind)
                    )
                ) $(where $($bounds)*)? {
                }
                ub_detect_helper(|f, g| {
                    let _foo: Self::Kind<'_> = f([]);
                    g(&_foo);
                });
            }
        }
    };
}

/// Holds an "opaqueified" `T::Kind`.
///
/// Note the lifetime, `'k`. This can be anything, as long as `T` outlives it.
/// 
/// # Examples
///
/// ```rust
/// #![feature(generic_associated_types)]
///
/// use core::cell::Cell;
/// 
/// use selfref::Holder;
/// use selfref::new_with_closure;
/// use selfref::opaque;
///
/// #[derive(Default)]
/// struct Foo<'a> {
///     foo: Cell<Option<&'a Foo<'a>>>,
/// }
///
/// struct FooKey;
/// opaque! {
///     impl Opaque for FooKey {
///         type Kind<'a> = Foo<'a>;
///     }
/// }
///
/// fn main() {
///     // We can use a closure here, but we need to give the compiler a hint.
///     let holder = Holder::<'_, FooKey>::new_with(
///         new_with_closure::<FooKey, _>(|foo| Foo::default())
///     );
/// }
/// ```
pub struct Holder<'k, T> where T: Opaque + 'k {
  _pinned: PhantomPinned,
  inner: <T as Opaque>::Kind<'k>,
}

/// Helper trait for creating a [`Holder`].
///
/// This is necessary because closures don't work properly here.
///
/// See [`Holder::new_with`] for examples.
pub trait NewWith<'k, T: Opaque + 'k> {
    /// Creates a new `T::Kind` with an unbound lifetime.
    fn new_with<'a>(self) -> T::Kind<'a> where 'k: 'a, T::Kind<'a>: Sized;
}

impl<'k, T: Opaque + 'k, F> NewWith<'k, T> for F
where
    F: for<'a> FnOnce([&'a (); 0]) -> <T as Opaque>::Kind<'a>,
{
    fn new_with<'a>(self) -> T::Kind<'a> where 'k: 'a, T::Kind<'a>: Sized {
        self([])
    }
}

/// Helper for creating a [`Holder`] using a closure.
///
/// This only works if `K` is `'static`.
pub fn new_with_closure<K: Opaque, F>(f: F) -> F
where
    F: for<'a> FnOnce([&'a (); 0]) -> <K as Opaque>::Kind<'a>,
    for<'a> K::Kind<'a>: Sized,
{
    f
}

impl<'k, T> Holder<'k, T> where T: Opaque {
    /// Creates a new holder.
    ///
    /// # Examples
    ///
    /// Simple example:
    ///
    /// ```rust
    /// #![feature(generic_associated_types)]
    ///
    /// use core::cell::Cell;
    /// 
    /// use selfref::Holder;
    /// use selfref::new_with_closure;
    /// use selfref::opaque;
    ///
    /// #[derive(Default)]
    /// struct Foo<'a> {
    ///     foo: Cell<Option<&'a Foo<'a>>>,
    /// }
    ///
    /// struct FooKey;
    /// opaque! {
    ///     impl Opaque for FooKey {
    ///         type Kind<'a> = Foo<'a>;
    ///     }
    /// }
    ///
    /// fn main() {
    ///     // We can use a closure here, but we need to help the compiler.
    ///     let holder = Holder::<'_, FooKey>::new_with(
    ///         new_with_closure::<FooKey, _>(|foo| Foo::default())
    ///     );
    /// }
    /// ```
    ///
    /// Lifetime parameters:
    ///
    /// ```rust
    /// #![feature(generic_associated_types)]
    ///
    /// use core::cell::Cell;
    /// use core::marker::PhantomData;
    /// 
    /// use selfref::Holder;
    /// use selfref::NewWith;
    /// use selfref::opaque;
    ///
    /// struct Foo<'a, 'b: 'a> {
    ///     foo: Cell<Option<&'a Foo<'a, 'b>>>,
    ///     t: &'b str,
    /// }
    ///
    /// struct FooKey<'b>(PhantomData<&'b str>);
    /// opaque! {
    ///     impl['b] Opaque for FooKey<'b> {
    ///         type Kind<'a> = Foo<'a, 'b>;
    ///     }
    /// }
    ///
    /// fn main() {
    ///     struct FooBuilder<'b>(&'b str);
    ///     impl<'k, 'b: 'k> NewWith<'k, FooKey<'b>> for FooBuilder<'b> {
    ///         fn new_with<'a>(self) -> Foo<'a, 'b> where 'k: 'a {
    ///             Foo {
    ///                 foo: Default::default(),
    ///                 t: self.0,
    ///             }
    ///         }
    ///     }
    ///     let stack_array: [u8; 5] = *b"hello";
    ///     // a non-'static &str
    ///     let stack_str = core::str::from_utf8(&stack_array).unwrap();
    ///     let holder = Holder::new_with(FooBuilder(stack_str));
    /// }
    /// ```
    pub fn new_with<F>(f: F) -> Self
    where
        F: NewWith<'k, T>,  T::Kind<'k>: Sized
    {
        Self {
            // it is important that the constructor cannot observe 'k!
            inner: f.new_with(),
            _pinned: PhantomPinned
        }
    }
}

/// Helper trait for working with a [`Holder`].
///
/// This is necessary because closures don't work properly here.
///
/// See [`Holder::operate_in`] for examples.
pub trait OperateIn<'k, T: Opaque + 'k> {
    /// The value returned by this operation.
    type Out;
    /// Do this operation.
    fn operate_in<'a>(self, x: Pin<&'a T::Kind<'a>>) -> Self::Out where 'k: 'a;
}

impl<'k, T: Opaque + 'k, F, R> OperateIn<'k, T> for F
where F: for<'a> FnOnce(Pin<&'a T::Kind<'a>>) -> R {
    type Out = R;
    fn operate_in<'a>(self, x: Pin<&'a T::Kind<'a>>) -> R where 'k: 'a {
        self(x)
    }
}

/// Helper for working with a [`Holder`] using a closure.
///
/// This only works if `K` is `'static`.
pub fn operate_in_closure<K: Opaque, F, R>(f: F) -> F 
where F: for<'a> FnOnce(Pin<&'a <K as Opaque>::Kind<'a>>) -> R {
    f
}

impl<'k, T> Holder<'k, T> where T: Opaque {
    /// Operates in this (pinned) holder.
    ///
    /// This "unwraps" the value in this holder, and binds its lifetime to a
    /// new stack frame.
    ///
    /// # Examples
    ///
    /// Simple example:
    ///
    /// ```rust
    /// #![feature(generic_associated_types)]
    /// #![feature(pin_macro)]
    ///
    /// use core::cell::Cell;
    /// use core::pin::pin;
    /// 
    /// use selfref::Holder;
    /// use selfref::new_with_closure;
    /// use selfref::opaque;
    /// use selfref::operate_in_closure;
    ///
    /// #[derive(Default)]
    /// struct Foo<'a> {
    ///     foo: Cell<Option<&'a Foo<'a>>>,
    /// }
    ///
    /// struct FooKey;
    /// opaque! {
    ///     impl Opaque for FooKey {
    ///         type Kind<'a> = Foo<'a>;
    ///     }
    /// }
    ///
    /// fn main() {
    ///     let holder = pin!(Holder::<'_, FooKey>::new_with(
    ///         new_with_closure::<FooKey, _>(|foo| Foo::default())
    ///     ));
    ///     // Actually making our Foo refer to itself.
    ///     holder.as_ref().operate_in(
    ///         operate_in_closure::<FooKey, _, _>(|foo| {
    ///             foo.foo.set(Some(foo.get_ref()));
    ///         })
    ///     );
    /// }
    /// ```
    ///
    /// With lifetime parameters:
    ///
    /// ```rust
    /// #![feature(generic_associated_types)]
    /// #![feature(pin_macro)]
    ///
    /// use core::cell::Cell;
    /// use core::marker::PhantomData;
    /// use core::pin::Pin;
    /// use core::pin::pin;
    /// 
    /// use selfref::Holder;
    /// use selfref::NewWith;
    /// use selfref::OperateIn;
    /// use selfref::opaque;
    ///
    /// struct Foo<'a, 'b: 'a> {
    ///     foo: Cell<Option<&'a Foo<'a, 'b>>>,
    ///     t: &'b str,
    /// }
    ///
    /// struct FooKey<'b>(PhantomData<&'b str>);
    /// opaque! {
    ///     impl['b] Opaque for FooKey<'b> {
    ///         type Kind<'a> = Foo<'a, 'b>;
    ///     }
    /// }
    ///
    /// fn main() {
    ///     struct FooBuilder<'b>(&'b str);
    ///     impl<'k, 'b: 'k> NewWith<'k, FooKey<'b>> for FooBuilder<'b> {
    ///         fn new_with<'a>(self) -> Foo<'a, 'b> where 'k: 'a {
    ///             Foo {
    ///                 foo: Default::default(),
    ///                 t: self.0,
    ///             }
    ///         }
    ///     }
    ///     let stack_array: [u8; 5] = *b"hello";
    ///     // a non-'static &str
    ///     let stack_str = core::str::from_utf8(&stack_array).unwrap();
    ///     let holder = pin!(Holder::new_with(FooBuilder(stack_str)));
    ///     // Actually making our Foo refer to itself.
    ///     struct SetFooRef;
    ///     impl<'k, 'b: 'k> OperateIn<'k, FooKey<'b>> for SetFooRef {
    ///         type Out = ();
    ///         fn operate_in<'a>(self, foo: Pin<&'a Foo<'a, 'b>>)
    ///         where 'k: 'a {
    ///             foo.foo.set(Some(foo.get_ref()));
    ///         }
    ///     }
    ///     holder.as_ref().operate_in(SetFooRef);
    /// }
    /// ```
    pub fn operate_in<'i, F, R>(self: Pin<&'i Self>, f: F) -> R
    where 
        F: OperateIn<'k, T, Out=R>
    {
        /// Converts Pin<&'a T::Kind<'k>> to Pin<&'b T::Kind<'b>>
        unsafe fn upcast_dangling<'a, 'b, 'c, T: Opaque + 'c>(
            x: Pin<&'a T::Kind<'c>>,
        ) -> Pin<&'b T::Kind<'b>>
        where T::Kind<'c>: 'a {
            mem::transmute(x)
        }

        f.operate_in(unsafe {
            upcast_dangling::<'i, 'k, '_, T>
            (self.map_unchecked(|self_ref| {
                &self_ref.inner
            }))
        })
    }
}