async-safe-defer 0.2.0

//! Allocator-free asynchronous cleanup with fixed-capacity registries.

use core::{
    error::Error,
    fmt,
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};

type FixedFuture<'a> = Pin<&'a mut (dyn Future<Output = ()> + 'a)>;
type FixedSendFuture<'a> = Pin<&'a mut (dyn Future<Output = ()> + Send + 'a)>;

fn erase_future<'a, F>(future: Pin<&'a mut F>) -> FixedFuture<'a>
where
    F: Future<Output = ()> + 'a,
{
    future
}

fn erase_send_future<'a, F>(future: Pin<&'a mut F>) -> FixedSendFuture<'a>
where
    F: Future<Output = ()> + Send + 'a,
{
    future
}

/// Preserves a task rejected by a full fixed-capacity scope.
#[must_use]
pub struct CapacityError<T> {
    task: T,
}

impl<T> CapacityError<T> {
    /// Recovers the rejected task.
    #[inline]
    pub fn into_inner(self) -> T {
        self.task
    }

    /// Borrows the rejected task.
    #[inline]
    pub const fn task(&self) -> &T {
        &self.task
    }
}

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

impl<T> fmt::Display for CapacityError<T> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("the fixed async scope is at capacity")
    }
}

impl<T> Error for CapacityError<T> {}

trait CleanupFuture {
    fn poll_cleanup(&mut self, context: &mut Context<'_>) -> Poll<()>;
}

impl<F> CleanupFuture for Pin<&mut F>
where
    F: Future<Output = ()> + ?Sized,
{
    fn poll_cleanup(&mut self, context: &mut Context<'_>) -> Poll<()> {
        self.as_mut().poll(context)
    }
}

struct Registry<T, const N: usize> {
    tasks: [Option<T>; N],
    len: usize,
}

impl<T, const N: usize> Registry<T, N> {
    const fn new() -> Self {
        Self {
            tasks: [const { None }; N],
            len: 0,
        }
    }

    const fn len(&self) -> usize {
        self.len
    }

    const fn is_empty(&self) -> bool {
        self.len == 0
    }

    const fn is_full(&self) -> bool {
        self.len == N
    }

    fn push(&mut self, task: T) {
        debug_assert!(!self.is_full());
        self.tasks[self.len] = Some(task);
        self.len += 1;
    }

    fn clear(&mut self) {
        while self.len > 0 {
            self.len -= 1;
            self.tasks[self.len] = None;
        }
    }

    fn poll_all(&mut self, context: &mut Context<'_>) -> Poll<()>
    where
        T: CleanupFuture,
    {
        loop {
            let Some(index) = self.len.checked_sub(1) else {
                return Poll::Ready(());
            };
            let mut future = self.tasks[index]
                .take()
                .expect("an occupied fixed scope slot contains a future");
            self.len = index;

            match future.poll_cleanup(context) {
                Poll::Pending => {
                    self.tasks[index] = Some(future);
                    self.len = index + 1;
                    return Poll::Pending;
                }
                Poll::Ready(()) => {}
            }
        }
    }
}

fn try_store<T, U, const N: usize>(
    registry: &mut Registry<T, N>,
    task: U,
    erase: impl FnOnce(U) -> T,
) -> Result<(), CapacityError<U>> {
    if registry.is_full() {
        return Err(CapacityError { task });
    }

    registry.push(erase(task));
    Ok(())
}

/// An allocator-free LIFO stack of caller-owned cleanup futures.
///
/// The scope stores mutable pinned references in an inline array, allowing
/// heterogeneous future types and borrowed local state without allocation.
/// Use [`FixedSendAsyncScope`] when the scope or its execution futures must be [`Send`].
///
/// Dropping a pending [`Run`] keeps its current future registered. A later run
/// resumes it after any newer LIFO entries.
#[must_use = "call run().await or finish().await to execute registered cleanup futures"]
pub struct FixedAsyncScope<'a, const N: usize> {
    registry: Registry<FixedFuture<'a>, N>,
}

/// Drains a [`FixedAsyncScope`] without consuming it.
#[must_use = "futures do nothing unless polled or awaited"]
pub struct Run<'scope, 'task, const N: usize> {
    scope: &'scope mut FixedAsyncScope<'task, N>,
}

impl<const N: usize> Future for Run<'_, '_, N> {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
        self.scope.registry.poll_all(context)
    }
}

impl<'a, const N: usize> FixedAsyncScope<'a, N> {
    /// Creates an empty scope with capacity `N`.
    #[inline]
    pub const fn new() -> Self {
        Self {
            registry: Registry::new(),
        }
    }

    /// Stores a caller-owned pinned future without polling it.
    ///
    /// The original pinned reference is returned inside [`CapacityError`] when
    /// the scope is full.
    #[inline]
    pub fn try_defer<F>(
        &mut self,
        future: Pin<&'a mut F>,
    ) -> Result<(), CapacityError<Pin<&'a mut F>>>
    where
        F: Future<Output = ()> + 'a,
    {
        try_store(&mut self.registry, future, erase_future::<F>)
    }

    /// Returns the number of cleanup futures that have not completed.
    #[inline]
    pub const fn len(&self) -> usize {
        self.registry.len()
    }

    /// Returns the maximum number of futures the scope can hold.
    #[inline]
    pub const fn capacity(&self) -> usize {
        N
    }

    /// Returns `true` when no cleanup futures remain.
    #[inline]
    pub const fn is_empty(&self) -> bool {
        self.registry.is_empty()
    }

    /// Returns `true` when no more cleanup futures can be registered.
    #[inline]
    pub const fn is_full(&self) -> bool {
        self.registry.is_full()
    }

    /// Unregisters all futures without polling them further.
    pub fn clear(&mut self) {
        self.registry.clear();
    }

    /// Borrows the scope and drains its futures in LIFO order.
    #[inline]
    pub fn run(&mut self) -> Run<'_, 'a, N> {
        Run { scope: self }
    }

    /// Consumes the scope and drains its futures in LIFO order.
    ///
    /// Dropping the returned future unregisters every remaining future.
    #[inline]
    pub async fn finish(mut self) {
        self.run().await;
    }
}

impl<const N: usize> Default for FixedAsyncScope<'_, N> {
    fn default() -> Self {
        Self::new()
    }
}

impl<const N: usize> fmt::Debug for FixedAsyncScope<'_, N> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("FixedAsyncScope")
            .field("pending", &self.len())
            .field("capacity", &N)
            .finish()
    }
}

/// An allocator-free LIFO stack of caller-owned `Send` cleanup futures.
///
/// Execution and cancellation match [`FixedAsyncScope`]. The scope,
/// [`FixedSendRun`], and future returned by [`Self::finish`] implement [`Send`].
#[must_use = "call run().await or finish().await to execute registered cleanup futures"]
pub struct FixedSendAsyncScope<'a, const N: usize> {
    registry: Registry<FixedSendFuture<'a>, N>,
}

/// A `Send` future that drains a [`FixedSendAsyncScope`] without consuming it.
#[must_use = "futures do nothing unless polled or awaited"]
pub struct FixedSendRun<'scope, 'task, const N: usize> {
    scope: &'scope mut FixedSendAsyncScope<'task, N>,
}

impl<const N: usize> Future for FixedSendRun<'_, '_, N> {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
        self.scope.registry.poll_all(context)
    }
}

impl<'a, const N: usize> FixedSendAsyncScope<'a, N> {
    /// Creates an empty scope with capacity `N`.
    #[inline]
    pub const fn new() -> Self {
        Self {
            registry: Registry::new(),
        }
    }

    /// Stores a caller-owned pinned `Send` future without polling it.
    ///
    /// The original pinned reference is returned inside [`CapacityError`] when
    /// the scope is full.
    #[inline]
    pub fn try_defer<F>(
        &mut self,
        future: Pin<&'a mut F>,
    ) -> Result<(), CapacityError<Pin<&'a mut F>>>
    where
        F: Future<Output = ()> + Send + 'a,
    {
        try_store(&mut self.registry, future, erase_send_future::<F>)
    }

    /// Returns the number of cleanup futures that have not completed.
    #[inline]
    pub const fn len(&self) -> usize {
        self.registry.len()
    }

    /// Returns the maximum number of futures the scope can hold.
    #[inline]
    pub const fn capacity(&self) -> usize {
        N
    }

    /// Returns `true` when no cleanup futures remain.
    #[inline]
    pub const fn is_empty(&self) -> bool {
        self.registry.is_empty()
    }

    /// Returns `true` when no more cleanup futures can be registered.
    #[inline]
    pub const fn is_full(&self) -> bool {
        self.registry.is_full()
    }

    /// Unregisters all futures without polling them further.
    pub fn clear(&mut self) {
        self.registry.clear();
    }

    /// Borrows the scope and drains its futures in LIFO order.
    #[inline]
    pub fn run(&mut self) -> FixedSendRun<'_, 'a, N> {
        FixedSendRun { scope: self }
    }

    /// Consumes the scope and drains its futures in LIFO order.
    ///
    /// Dropping the returned future unregisters every remaining future.
    #[inline]
    pub async fn finish(mut self) {
        self.run().await;
    }
}

impl<const N: usize> Default for FixedSendAsyncScope<'_, N> {
    fn default() -> Self {
        Self::new()
    }
}

impl<const N: usize> fmt::Debug for FixedSendAsyncScope<'_, N> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("FixedSendAsyncScope")
            .field("pending", &self.len())
            .field("capacity", &N)
            .finish()
    }
}