// 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

#![no_std]
#![cfg_attr(feature="nightly", feature(dropck_eyepatch))]

//! Pain-free self-referential/recursively-referential pinned types.
//!
//! Because this crate optimizes for recursive references, it's much nicer to
//! use for porting/copying C code than alternatives. However, the cost is that
//! it requires generic associated types (GATs). Further, this is the only
//! crate for self-referential types which supports additional lifetimes than
//! the self-referential one (with the small ergonomic caveat that closures
//! don't support where bounds for now, so a proper trait has to be used).
//!
//! When using this crate, all you have to do is define your struct like you
//! would in C, and then tell `selfref` about it. However, do watch out; the
//! following will not be usable, despite compiling:
//!
//! ```rust
//! use selfref::opaque;
//!
//! // This is the actual self-referential struct.
//! struct MySelfRefStruct<'this> {
//!     this: &'this MySelfRefStruct<'this>,
//! }
//!
//! // This is a "marker" type, used to tell `selfref` about the type.
//! // This is actually the "type family" GAT pattern.
//! struct MySelfRefStructKey;
//!
//! // This tells `selfref` about our type.
//! opaque! {
//!     impl Opaque for MySelfRefStructKey {
//!         type Kind<'this> = MySelfRefStruct<'this>;
//!     }
//! }
//! ```
//!
//! That's because there's no way to create the reference
//! `&'this MySelfRefStruct<'this>` before constructing the struct! That is,
//! you cannot do this:
//!
//! ```rust compile_fail
//! struct MySelfRefStruct<'this> {
//!     this: &'this MySelfRefStruct<'this>,
//! }
//!
//! fn main() {
//!     let x = MySelfRefStruct { this: &x };
//! }
//! ```
//!
//! But then, you cannot do this in C either. Instead, to make it usable, we
//! use an `Option` and, to be able to set it, a `Cell`, like so:
//!
//! ```rust
//! use std::cell::Cell;
//!
//! use selfref::opaque;
//!
//! struct MySelfRefStruct<'this> {
//!     this: Cell<Option<&'this MySelfRefStruct<'this>>>,
//! }
//!
//! struct MySelfRefStructKey;
//!
//! opaque! {
//!     impl Opaque for MySelfRefStructKey {
//!         type Kind<'this> = MySelfRefStruct<'this>;
//!     }
//! }
//! ```
//!
//! and then we can use it:
//!
//! ```rust
//! // lines from the above example have been omitted
//!
//! # use std::cell::Cell;
//! #
//! use selfref::Holder;
//! # use selfref::opaque;
//! #
//! # struct MySelfRefStruct<'this> {
//! #     this: Cell<Option<&'this MySelfRefStruct<'this>>>,
//! # }
//! #
//! # struct MySelfRefStructKey;
//! #
//! # opaque! {
//! #     impl Opaque for MySelfRefStructKey {
//! #         type Kind<'this> = MySelfRefStruct<'this>;
//! #     }
//! # }
//!
//! fn main() {
//!     // first, construct the struct
//!     let holder = Box::pin(Holder::<'_, MySelfRefStructKey>::new_with(
//!         |foo| foo.build({
//!             MySelfRefStruct {
//!                 this: Cell::new(None)
//!             }
//!         })
//!     ));
//!
//!     // then, build the self-reference
//!     holder.as_ref().operate_in(
//!         |this| {
//!             this.this.set(Some(this.get_ref()));
//!         }
//!     );
//! }
//! ```
//!
//! # Examples
//!
//! This is a more complex example borrowing from an external lifetime:
//!
//! ```rust
//! #![feature(pin_macro)]
//!
//! use core::cell::Cell;
//! use core::marker::PhantomData;
//! use core::pin::Pin;
//! use core::pin::pin;
//! 
//! use selfref::Holder;
//! 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() {
//!     // a non-'static &str
//!     let stack_array: [u8; 5] = *b"hello";
//!     let stack_str = core::str::from_utf8(&stack_array).unwrap();
//!
//!     // construct the struct
//!     let holder = pin!(Holder::<'_, FooKey<'_>>::new_with(|foo| {
//!         foo.build(Foo {
//!             foo: Default::default(),
//!             t: stack_str,
//!         })
//!     }));
//!
//!     holder.as_ref().operate_in(|foo| {
//!         foo.foo.set(Some(foo.get_ref()));
//!     });
//! }
//! ```
//!
//! # Features
//!
//! Due to [PhantomData is unsound](https://github.com/rust-lang/rust/issues/102810)
//! we currently require the following features for `T: ?Sized` support in
//! `selfref::opaque!`:
//!
//! - `alloc` - `selfref::opaque!` for `T: ?Sized` is provided by `Box`.
//! - `nightly` - `selfref::opaque!` for `T: ?Sized` is provided by a *wrapper*
//!     around `PhantomData`, which works around the above issue. we call this
//!     "PhantomDrop".
//!
//! When enabling both features, `nightly` takes over and we use the wrapper
//! always. This doesn't make a practical difference since the generated UB
//! check is dead code anyway, but `PhantomDrop` should be lighter on compile
//! times.
//!
//! If not using either feature, `T: ?Sized` support requires `unsafe`ly
//! implementing `Opaque`.
//!
//! Note that we do **not** enable any features by default! We assume most
//! folks aren't coming to this crate for its `T: ?Sized` support, so these are
//! the best defaults for crates to depend on. If they do need the `?Sized`
//! support they can just enable one of these (probably `alloc`).

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

// there's no sound way to dropck T: ?Sized without either alloc or nightly.
//
// so we just have the user opt-in to alloc or nightly as desired.
//
// when using alloc, we use Box<T> for UBCheck.
//
// when using nightly, we use our custom PhantomDrop<T> for UBCheck.
//
// when using neither, we just error on T: ?Sized and require a manual unsafe
// impl of Opaque.

#[cfg(all(feature="alloc", not(feature="nightly")))]
extern crate alloc;

#[cfg(all(feature="alloc", not(feature="nightly")))]
#[doc(hidden)]
pub use alloc::boxed::Box as UBCheck;

#[cfg(feature="nightly")]
#[doc(hidden)]
pub struct UBCheck<T: ?Sized>(core::marker::PhantomData<T>);

#[cfg(all(not(feature="alloc"), not(feature="nightly")))]
#[doc(hidden)]
pub struct UBCheck<T>(T); // use feature "alloc" or "nightly" for T: ?Sized

#[cfg(feature="nightly")]
// SAFETY: dropck's like a Box<T>, but is no-alloc friendly.
unsafe impl<#[may_dangle] T: ?Sized> Drop for UBCheck<T> {
    fn drop(&mut self) {}
}

/// 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.
///
/// It's recommended to use the `selfref::opaque!` macro instead, which
/// enforces these invariants. For example, this doesn't compile:
///
/// ```rust compile_fail
/// use std::cell::Cell;
/// use selfref::opaque;
///
/// struct Foo<'a> {
///     foo: Cell<Option<&'a Foo<'a>>>,
/// }
///
/// impl<'a> Drop for Foo<'a> {
///     fn drop(&mut self) {
///     }
/// }
///
/// struct FooKey;
/// opaque! {
///     impl Opaque for FooKey {
///         type Kind<'a> = Foo<'a>;
///     }
/// }
/// ```
///
/// But by removing the `Drop` impl, it compiles:
///
/// ```rust
/// use std::cell::Cell;
/// use selfref::opaque;
///
/// struct Foo<'a> {
///     foo: Cell<Option<&'a Foo<'a>>>,
/// }
///
/// //impl<'a> Drop for Foo<'a> {
/// //    fn drop(&mut self) {
/// //    }
/// //}
///
/// struct FooKey;
/// opaque! {
///     impl Opaque for FooKey {
///         type Kind<'a> = Foo<'a>;
///     }
/// }
/// ```
///
/// # Examples
///
/// ```rust
/// 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
/// 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
/// 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]) -> $crate::UBCheck<$kind>,
                        for<$l> fn(&$l $crate::UBCheck<$kind>)
                    )
                ) $(where $($bounds)*)? {
                }
                ub_detect_helper(|f, g| {
                    let x: $crate::UBCheck<Self::Kind<'_>> = f([]);
                    g(&x);
                });
            }
        }
    };
    (
        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]) -> $crate::UBCheck<$kind>,
                        for<$l> fn(&$l $crate::UBCheck<$kind>)
                    )
                ) $(where $($bounds)*)? {
                }
                ub_detect_helper(|f, g| {
                    let x: $crate::UBCheck<Self::Kind<'_>> = f([]);
                    g(&x);
                });
            }
        }
    };
}

/// Holds an "opaqueified" `T::Kind`.
///
/// Note the lifetime, `'k`. This can be anything, as long as `T` outlives it.
/// 
/// # Examples
///
/// ```rust
/// use core::cell::Cell;
/// 
/// use selfref::Holder;
/// 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(
///         |foo| foo.build(Foo::default())
///     );
/// }
/// ```
pub struct Holder<'k, T> where T: Opaque + 'k {
  _pinned: PhantomPinned,
  inner: <T as Opaque>::Kind<'k>,
}

/// Helper for creating a [`Holder`].
///
/// This is necessary because closures don't work properly here.
///
/// See [`Holder::new_with`] for examples.
pub struct Builder<'k, T: Opaque + 'k> where T::Kind<'k>: Sized {
    inner: Option<T::Kind<'k>>,
}

impl<'k, T: Opaque + 'k> Builder<'k, T> where T::Kind<'k>: Sized {
    /// Builds the [`Holder`].
    #[inline]
    pub fn build(&mut self, t: T::Kind<'k>) {
        self.inner = Some(t);
    }
}

impl<'k, T> Holder<'k, T> where T: Opaque {
    /// Creates a new holder.
    ///
    /// # Examples
    ///
    /// Simple example:
    ///
    /// ```rust
    /// use core::cell::Cell;
    /// 
    /// use selfref::Holder;
    /// 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(
    ///         |foo| foo.build(Foo::default())
    ///     );
    /// }
    /// ```
    ///
    /// Lifetime parameters:
    ///
    /// ```rust
    /// use core::cell::Cell;
    /// use core::marker::PhantomData;
    /// 
    /// use selfref::Holder;
    /// 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() {
    ///     let stack_array: [u8; 5] = *b"hello";
    ///     // a non-'static &str
    ///     let stack_str = core::str::from_utf8(&stack_array).unwrap();
    ///     let holder = Holder::<'_, FooKey<'_>>::new_with(|foo| {
    ///         foo.build(Foo {
    ///             foo: Default::default(),
    ///             t: stack_str,
    ///         });
    ///     });
    /// }
    /// ```
    pub fn new_with<F>(f: F) -> Self
    where
        F: for<'a> FnOnce(&mut Builder<'a, T>),
        T::Kind<'k>: Sized,
    {
        let mut builder = Builder { inner: None };
        f(&mut builder);
        Self {
            // it is important that the constructor cannot observe 'k!
            inner: builder.inner.unwrap(),
            _pinned: PhantomPinned
        }
    }
}

/// Wrapper around a `Pin<&'k T::Kind<'k>>` for implied bounds.
///
/// Derefs to `Pin<&'k T::Kind<'k>>`.
pub struct OperateIn<'k, T> where T: Opaque + 'k {
    inner: Pin<&'k T::Kind<'k>>,
}

impl<'k, T> core::ops::Deref for OperateIn<'k, T> where T: Opaque {
    type Target = Pin<&'k T::Kind<'k>>;

    fn deref(&self) -> &Pin<&'k T::Kind<'k>> {
        &self.inner
    }
}

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(pin_macro)]
    ///
    /// use core::cell::Cell;
    /// use core::pin::pin;
    /// 
    /// use selfref::Holder;
    /// 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() {
    ///     let holder = pin!(Holder::<'_, FooKey>::new_with(
    ///         |foo| foo.build(Foo::default())
    ///     ));
    ///     // Actually making our Foo refer to itself.
    ///     holder.as_ref().operate_in(
    ///         |foo| {
    ///             foo.foo.set(Some(foo.get_ref()));
    ///         }
    ///     );
    /// }
    /// ```
    ///
    /// With lifetime parameters:
    ///
    /// ```rust
    /// #![feature(pin_macro)]
    ///
    /// use core::cell::Cell;
    /// use core::marker::PhantomData;
    /// use core::pin::Pin;
    /// use core::pin::pin;
    /// 
    /// use selfref::Holder;
    /// 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() {
    ///     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::<'_, FooKey<'_>>::new_with(|foo| {
    ///         foo.build(Foo {
    ///             foo: Default::default(),
    ///             t: stack_str,
    ///         });
    ///     }));
    ///     // Actually making our Foo refer to itself.
    ///     holder.as_ref().operate_in(|foo| {
    ///         foo.foo.set(Some(foo.get_ref()));
    ///     });
    /// }
    /// ```
    pub fn operate_in<'i, F, R>(self: Pin<&'i Self>, f: F) -> R
    where 
        F: for<'x> FnOnce(OperateIn<'x, T>) -> R
    {
        /// Converts `Pin<&'a T::Kind<'k>>` to `Pin<&'b T::Kind<'b>>`.
        ///
        /// Not sure why this is called "upcast_dangling" since none of these
        /// are actually dangling. But anyway.
        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(OperateIn {
            inner: unsafe {
                upcast_dangling::<'i, 'k, '_, T>
                (self.map_unchecked(|self_ref| {
                    &self_ref.inner
                }))
            }
        })
    }
}