rec_cell 0.1.0

use core::{
    fmt::Debug,
    iter,
    marker::PhantomData,
    mem::{ManuallyDrop, MaybeUninit},
    ptr,
};

use crate::{RecCell, RecToken, utils::Invariant};

/// A builder used to compose cyclic initializations.
///
/// It is provided in the closure given to [`RecToken::with_builder`]. The
/// closure must return the builder to indicate success of all initializations,
/// such that access to the `RecToken` can be resumed.
pub struct RecBuilder<'b, 't> {
    // The extra `'b` lifetime keeps it from leaving the builder scope.
    _phantom: PhantomData<(Invariant<'t>, Invariant<'b>)>,
}

impl Debug for RecBuilder<'_, '_> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("RecBuilder").finish_non_exhaustive()
    }
}

impl<'t> RecToken<'t> {
    /// Temporarily replaces this token with a [`RecBuilder`] for composing
    /// several cyclic initializations.
    ///
    /// # Examples
    ///
    /// This builds two nodes pointing to each other in sequence. Each `try_add`
    /// receives the next builder, so the second node can point to the first;
    /// the outer continuation then initializes the first node pointing to the
    /// second. The same construction can usually be written more compactly as
    /// one [`RecToken::make_cyclic`] call over a tuple of slots (shown
    /// below).
    ///
    /// ```
    /// use core::mem::MaybeUninit;
    /// use rec_cell::{RecCell, RecToken};
    ///
    /// struct Node<'a, 't> { next: &'a RecCell<'t, Node<'a, 't>> }
    ///
    /// RecToken::new(|token| {
    ///     let (mut first, mut second) = (MaybeUninit::uninit(), MaybeUninit::uninit());
    ///     let ((first_cell, second_cell), token) = token.with_builder(|builder| {
    ///         builder.try_add(&mut first, |builder, first_cell| {
    ///             let (cells, builder) = builder.try_add(&mut second, move |builder, second_cell| {
    ///                 let cells = (first_cell, second_cell);
    ///                 Ok::<_, ()>((cells, builder, Node { next: first_cell }))
    ///             })?;
    ///             let (_, second_cell) = cells;
    ///             Ok::<_, ()>((cells, builder, Node { next: second_cell }))
    ///         })
    ///     }).unwrap();
    ///     assert!(core::ptr::eq(first_cell.borrow(&token).next, second_cell));
    ///     assert!(core::ptr::eq(second_cell.borrow(&token).next, first_cell));
    /// });
    /// ```
    ///
    /// This can also be done more compactly with a tuple input to
    /// [`RecToken::make_cyclic`]:
    ///
    /// ```
    /// use core::mem::MaybeUninit;
    /// use rec_cell::{RecCell, RecToken};
    /// RecToken::new(|token| {
    ///     let (mut left, mut right) = (MaybeUninit::uninit(), MaybeUninit::uninit());
    ///     struct Node<'a, 't> { next: &'a RecCell<'t, Node<'a, 't>> }
    ///     let ((left, right), token) = token.make_cyclic(
    ///         (&mut left, &mut right),
    ///         |(left, right)| ((left, right), (Node { next: right }, Node { next: left })),
    ///     );
    ///     assert!(core::ptr::eq(left.borrow(&token).next, right));
    ///     assert!(core::ptr::eq(right.borrow(&token).next, left));
    /// });
    /// ```
    pub fn with_builder<R, E>(
        self,
        scope: impl for<'b> FnOnce(RecBuilder<'b, 't>) -> Result<(R, RecBuilder<'b, 't>), E>,
    ) -> Result<(R, Self), E> {
        scope(RecBuilder {
            _phantom: PhantomData,
        })
        .map(|(result, _)| (result, self))
    }
}

impl<'t> RecBuilder<'_, 't> {
    /// Adds one or more cyclic slots to this builder.
    ///
    /// The continuation receives the next builder and references to the
    /// future cells, then returns their values and the builder. This sequencing
    /// keeps the token unavailable until every added slot is initialized.
    ///
    /// # Examples
    ///
    /// ```
    /// use core::mem::MaybeUninit;
    /// use rec_cell::RecToken;
    /// RecToken::new(|token| {
    ///     let mut slot = MaybeUninit::uninit();
    ///     let (value, _) = token.with_builder(|builder| {
    ///         builder.try_add(&mut slot, |builder, _| Ok::<_, ()>((1, builder, 2)))
    ///     }).unwrap();
    ///     assert_eq!(value, 1);
    /// });
    /// ```
    pub fn try_add<'a, U, R, E>(
        self,
        uninit: U,
        scope: impl for<'c> FnOnce(
            RecBuilder<'c, 't>,
            U::ScopeInput<'t>,
        ) -> Result<(R, RecBuilder<'c, 't>, U::ScopeOutput<'t>), E>,
    ) -> Result<(R, Self), E>
    where
        U: CyclicBuildable<'a>,
        't: 'a,
    {
        // SAFETY: the linear builder keeps the matching token unavailable.
        let (input, output_fn) = unsafe { U::__prepare(uninit) };
        let (result, _, output) = scope(
            RecBuilder {
                _phantom: PhantomData,
            },
            input,
        )?;
        output_fn(output);
        Ok((result, self))
    }
}

/// Internal sealing trait; downstream crates cannot implement
/// [`CyclicBuildable`].
pub trait Sealed {}

/// Storage where cyclic initialization can be performed.
#[allow(clippy::missing_safety_doc, reason = "sealed trait")]
pub unsafe trait CyclicBuildable<'a>: Sized + Sealed {
    /// References made available to the initialization closure.
    type ScopeInput<'t>
    where
        't: 'a;
    /// Values the initialization closure must return for the storage.
    type ScopeOutput<'t>
    where
        't: 'a;

    #[doc(hidden)]
    // `__prepare` creates an uninitialized structure and a closure
    // used to complete its initialization.
    //
    // The caller of this function guarantees that the corresponding
    // `RecToken` is not available to any other code while the structure is
    // uninitialized.
    unsafe fn __prepare<'t: 'a>(
        uninit: Self,
    ) -> (Self::ScopeInput<'t>, impl FnOnce(Self::ScopeOutput<'t>));
}

impl<T> Sealed for &'_ mut MaybeUninit<T> {}

// SAFETY: this __prepare properly follows the initialization protocol.
unsafe impl<'a, T> CyclicBuildable<'a> for &'a mut MaybeUninit<T> {
    type ScopeInput<'t>
        = &'a RecCell<'t, T>
    where
        't: 'a;
    type ScopeOutput<'t>
        = T
    where
        't: 'a;

    unsafe fn __prepare<'t: 'a>(
        uninit: Self,
    ) -> (Self::ScopeInput<'t>, impl FnOnce(Self::ScopeOutput<'t>)) {
        // SAFETY: the caller has removed the matching token from circulation.
        let cell = unsafe { RecCell::from_uninit(uninit) };
        (cell, move |value: T| {
            // SAFETY: this completion closure is called exactly once for its slot.
            unsafe { cell.write_to_uninit(value) }
        })
    }
}

// Show one representative tuple implementation in rustdoc. The remaining
// concrete arities are implementation details of the fake variadic impl.
#[cfg(docsrs)]
macro_rules! maybe_tuple_doc {
    ($type:ident @ $item:item) => {
        #[doc(fake_variadic)]
        #[doc = "This trait is implemented for tuples up to twelve items long."]
        $item
    };
    ($($type:ident)* @ $item:item) => {
        #[doc(hidden)]
        $item
    };
}
#[cfg(not(docsrs))]
macro_rules! maybe_tuple_doc {
    ($($type:ident)* @ $item:item) => {
        $item
    };
}

macro_rules! impl_tuple_cyclic_buildable {
    // This branch is redundant, but it leads to many kinds of
    // unused code warnings in the other branch, so we write this
    // separately.
    () => {
        impl Sealed for () {}

        maybe_tuple_doc! {
            @
            // SAFETY: the empty tuple needs no initialization protocol.
            unsafe impl<'a> CyclicBuildable<'a> for ()
            {
                type ScopeInput<'t>
                    = ()
                where
                    't: 'a;
                type ScopeOutput<'t>
                    = ()
                where
                    't: 'a;

                unsafe fn __prepare<'t: 'a>(
                    _: Self,
                ) -> (Self::ScopeInput<'t>, impl FnOnce(Self::ScopeOutput<'t>)) {
                    ((), |()| {})
                }
            }
        }
    };
    ($($type:ident : $index:tt),* $(,)?) => {
        impl<$($type),*> Sealed for ($($type,)*) {}

        maybe_tuple_doc! {
            $($type)* @
            // SAFETY: this __prepare properly follows the initialization protocol
            // given that each element's __prepare does.
            unsafe impl<'a, $($type),*> CyclicBuildable<'a> for ($($type,)*)
            where
                $($type: CyclicBuildable<'a>,)*
            {
                type ScopeInput<'t>
                    = ($($type::ScopeInput<'t>,)*)
                where
                    't: 'a;
                type ScopeOutput<'t>
                    = ($($type::ScopeOutput<'t>,)*)
                where
                    't: 'a;

                unsafe fn __prepare<'t: 'a>(
                    uninit: Self,
                ) -> (Self::ScopeInput<'t>, impl FnOnce(Self::ScopeOutput<'t>)) {
                    // SAFETY: the caller provides the token-exclusion invariant to each component.
                    let prepared = unsafe {
                        ($(CyclicBuildable::__prepare(uninit.$index),)*)
                    };
                    (($ (prepared.$index.0,)*), move |output| {
                        $(prepared.$index.1(output.$index);)*
                    })
                }
            }
        }
    };
}

macro_rules! impl_tuple_cyclic_buildables {
    ($($type:ident : $index:tt),* $(,)?) => {
        impl_tuple_cyclic_buildable!();
        impl_tuple_cyclic_buildables!(@impl (); $(($type : $index))*);
    };
    (@impl ($($done:tt)*);) => {};
    (@impl ($($done:tt)*); ($type:ident : $index:tt) $($rest:tt)*) => {
        impl_tuple_cyclic_buildable!($($done)* $type : $index);
        impl_tuple_cyclic_buildables!(
            @impl ($($done)* $type : $index,); $($rest)*
        );
    };
}

// Rust's standard library commonly provides tuple trait implementations through
// arity 12.
//
// First element named T rather than T0 such that fake_variadic looks nicer.
impl_tuple_cyclic_buildables!(
    T: 0, T1: 1, T2: 2, T3: 3, T4: 4, T5: 5, T6: 6, T7: 7, T8: 8, T9: 9, T10: 10,
    T11: 11
);

impl<T, const N: usize> Sealed for [T; N] {}

// SAFETY: this __prepare properly follows the initialization protocol
// given that each element's __prepare does.
unsafe impl<'a, T, const N: usize> CyclicBuildable<'a> for [T; N]
where
    T: CyclicBuildable<'a>,
{
    type ScopeInput<'t>
        = [T::ScopeInput<'t>; N]
    where
        't: 'a;
    type ScopeOutput<'t>
        = [T::ScopeOutput<'t>; N]
    where
        't: 'a;

    unsafe fn __prepare<'t: 'a>(
        uninit: Self,
    ) -> (Self::ScopeInput<'t>, impl FnOnce(Self::ScopeOutput<'t>)) {
        let uninit = ManuallyDrop::new(uninit);
        let mut inputs = [const { MaybeUninit::uninit() }; N];
        let mut finishes = [const { MaybeUninit::uninit() }; N];
        for i in 0..N {
            // SAFETY: each element is read exactly once from the manually dropped array.
            let uninit = unsafe { ptr::read(&raw const uninit[i]) };
            // SAFETY: the builder's caller has removed the matching token.
            let (input, finish) = unsafe { T::__prepare(uninit) };
            inputs[i].write(input);
            finishes[i].write(finish);
        }

        // SAFETY: every element was initialized in the loop above.
        let inputs = unsafe { MaybeUninit::from(inputs).assume_init() };
        // SAFETY: every element was initialized in the loop above.
        let finishes = unsafe { MaybeUninit::from(finishes).assume_init() };
        (inputs, move |outputs| {
            for (finish, output) in iter::zip(finishes, outputs) {
                finish(output);
            }
        })
    }
}