flaque 0.2.0

// Copyright 2026 Ivan Enderlin <ivan@mnt.io>. Licensed under the EUPL, version
// 1.2. See the `LICENSE` file.

use std::{
    pin::Pin,
    task::{Context, Poll},
};

#[cfg(doc)]
use super::Pool;
use super::{
    Arc, Async, AsyncBuilder, Object, PoolInner, Position, Sync, SyncBuilder,
    TryAcquireError, hint,
};

/// Error returned by [`Get`] via [`Pool::get`].
#[derive(Debug, PartialEq)]
pub enum GetError<B> {
    /// The pool has been closed: it's impossible to acquire a new object.
    Closed,

    /// A slot has been acquired successfully but the builder has failed to
    /// build an object. The slot is released before this error is returned,
    /// i.e. it can acquired later.
    Builder(B),
}

/// The future returned by [`Pool::get`].
///
/// It resolves to `Ok(`[`Object`]`)` once an object is available in the pool.
/// It resolves to `Err(`[`GetError`]`)` if the pool is either closed with
/// [`Pool::close`], or if the builder failed to build a new object.
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Get<T, B> {
    /// The inner value of the pool, so that it has access to
    /// [`PoolInner::try_acquire`] and [`PoolInner::get_wakers_sender`].
    ///
    /// It should ideally be using `MaybeUninit`. However, if `<Get as
    /// Future>::poll` is called after the future has been resolved, it will
    /// create an Undefined Behaviour, which is incorrect because `poll` is
    /// not marked as `unsafe`. So, in order to panic properly, we are
    /// using `Option`. (Alternatively, we could fuse the future and use
    /// `MaybeUninit`, but it would add an overhead).
    pool_inner: Option<Arc<PoolInner<T, B>>>,
}

impl<T, B> Get<T, B> {
    pub(super) fn new(inner: &Arc<PoolInner<T, B>>) -> Self {
        Self { pool_inner: Some(inner.clone()) }
    }

    /// Mark the future as ready with an `Ok` value.
    ///
    /// This is an internal API used by the various `poll` methods.
    ///
    /// # Safety
    ///
    /// `pool_inner` **MUST BE** `Some`.
    unsafe fn ready_ok<E>(
        pool_inner: &mut Option<Arc<PoolInner<T, B>>>,
        position: Position,
        object: T,
    ) -> Poll<Result<Object<T, B>, E>> {
        Poll::Ready(Ok(Object::new(
            position,
            object,
            // SAFETY: `pool_inner` is guaranteed to be `Some`, that's the
            // contract of the function.
            unsafe { pool_inner.take().unwrap_unchecked() },
        )))
    }

    /// Mark the future as ready with a builder error.
    ///
    /// This is an internal API used by the various `poll` methods.
    fn ready_builder_error<E>(
        builder_error: E,
    ) -> Poll<Result<Object<T, B>, GetError<E>>> {
        Poll::Ready(Err(GetError::Builder(builder_error)))
    }

    /// Mark the future as ready with a “poll is closed” error.
    ///
    /// This is an internal API used by the various `poll` methods.
    fn ready_closed_error<E>() -> Poll<Result<Object<T, B>, GetError<E>>> {
        Poll::Ready(Err(GetError::Closed))
    }

    /// Mark the future as pending.
    ///
    /// This is an internal API used by the various `poll` methods.
    fn pending<P>(
        pool_inner: &Arc<PoolInner<T, B>>,
        context: &Context,
    ) -> Poll<P> {
        let waker = context.waker().clone();
        let _ = pool_inner.get_wakers_sender.send(waker);

        Poll::Pending
    }
}

impl<T, B> Future for Get<T, Sync<B>>
where
    B: SyncBuilder<T>,
{
    type Output = Result<Object<T, Sync<B>>, GetError<B::Error>>;

    fn poll(
        mut self: Pin<&mut Self>,
        context: &mut Context<'_>,
    ) -> Poll<Self::Output> {
        if self.pool_inner.is_none() {
            hint::cold_path();

            unreachable!(
                "The `Get` future has already been resolved: it is violating the contract of `Future`"
            );
        }

        // SAFETY: The `Option` is guaranteed to be `Some`, as checked at the
        // beginning of this method.
        unsafe {
            let pool_inner = self.pool_inner.as_ref().unwrap_unchecked();

            // Try to acquire a slot.
            match pool_inner.try_acquire() {
                Ok((position, maybe_object_value)) => {
                    match pool_inner.build_object_value_if_necessary(
                        &position,
                        maybe_object_value,
                    ) {
                        Ok(object_value) => Get::ready_ok(
                            &mut self.pool_inner,
                            position,
                            object_value,
                        ),
                        Err(builder_error) => {
                            Get::ready_builder_error(builder_error)
                        }
                    }
                }

                Err(TryAcquireError::Closed) => Get::ready_closed_error(),

                Err(TryAcquireError::AllSlotsOccupied) => {
                    Get::pending(pool_inner, context)
                }
            }
        }
    }
}

// Map a `Get<T, Async<B>>` to an `AsyncGet`.
//
// This `IntoFuture` implementation helps to keep the same `Pool::get`
// implementation for both `SyncBuilder` and `AsyncBuilder` cases.
impl<T, B> IntoFuture for Get<T, Async<B>>
where
    T: 'static,
    B: AsyncBuilder<T> + 'static,
{
    type Output = Result<Object<T, Async<B>>, GetError<B::Error>>;
    type IntoFuture = AsyncGet<T, Async<B>>;

    fn into_future(self) -> Self::IntoFuture {
        AsyncGet { pool_inner: self.pool_inner, builder_state: None }
    }
}

/// Async-flavour of [`Get`], built automatically by [`<Get<T, Async<B>>
/// as IntoFuture>`][Get#impl-IntoFuture-for-Get<T,+Async<B>>].
///
/// This type is transparent for the user: with [`SyncBuilder`], calling
/// `.await` on `Get` will automatically poll the future, or, with
/// [`AsyncBuilder`], it will automatically transforms [`Get`] into [`AsyncGet`]
/// and then poll the future.
//
// Basically, this type adds the `builder_state` field.
pub struct AsyncGet<T, B>
where
    B: AsyncBuilder<T>,
{
    pool_inner: Option<Arc<PoolInner<T, B>>>,
    #[allow(clippy::type_complexity)]
    builder_state:
        Option<(Position, Pin<Box<dyn Future<Output = Result<T, B::Error>>>>)>,
}

impl<T, B> AsyncGet<T, B>
where
    B: AsyncBuilder<T>,
{
    /// Try to resolve the future assuming the object builder future exists.
    ///
    /// # Safety
    ///
    /// `Self::pool_inner` **MUST BE** `Some`.
    #[allow(clippy::type_complexity)]
    unsafe fn poll_with_object_builder(
        &mut self,
        position: Position,
        mut object_builder_future: Pin<
            Box<dyn Future<Output = Result<T, B::Error>>>,
        >,
        context: &mut Context,
    ) -> Poll<Result<Object<T, B>, GetError<B::Error>>> {
        unsafe {
            // Optimistically poll the object builder future.
            match object_builder_future.as_mut().poll(context) {
                // Yeah! We have our object value!
                Poll::Ready(Ok(object)) => {
                    Get::ready_ok(&mut self.pool_inner, position, object)
                }

                // Oh no, poor us, the builder failed to build an object… `Get`
                // must give up by forwarding the error.
                Poll::Ready(Err(builder_error)) => {
                    Get::ready_builder_error(builder_error)
                }

                // Builder is taking its time. Save it for later poll.
                Poll::Pending => {
                    self.builder_state
                        .replace((position, object_builder_future));

                    Get::pending(
                        self.pool_inner.as_ref().unwrap_unchecked(),
                        context,
                    )
                }
            }
        }
    }
}

impl<T, B> Future for AsyncGet<T, Async<B>>
where
    T: 'static,
    B: AsyncBuilder<T> + 'static,
{
    type Output = Result<Object<T, Async<B>>, GetError<B::Error>>;

    fn poll(
        mut self: Pin<&mut Self>,
        context: &mut Context<'_>,
    ) -> Poll<Self::Output> {
        if self.pool_inner.is_none() {
            hint::cold_path();

            unreachable!(
                "The `Get` future has already been resolved: it is violating the contract of `Future`"
            );
        }

        // SAFETY: The `Option` is guaranteed to be `Some`, as checked at the
        // beginning of this method.
        unsafe {
            // If there is a `builder_state`, it means the future has been
            // pooled once, and the object builder future was pending. Now, it's
            // time to poll it again.
            if let Some((position, object_builder_future)) =
                self.builder_state.take()
            {
                return self.poll_with_object_builder(
                    position,
                    object_builder_future,
                    context,
                );
            }

            let pool_inner = self.pool_inner.as_ref().unwrap_unchecked();

            match pool_inner.try_acquire() {
                Ok((position, maybe_object_value)) => {
                    let object_builder_future = Box::pin({
                        let pool_inner = pool_inner.clone();

                        async move {
                            pool_inner
                                .build_object_value_if_necessary(
                                    position,
                                    maybe_object_value,
                                )
                                .await
                        }
                    });

                    self.poll_with_object_builder(
                        position,
                        object_builder_future,
                        context,
                    )
                }

                Err(TryAcquireError::Closed) => Get::ready_closed_error(),

                Err(TryAcquireError::AllSlotsOccupied) => {
                    Get::pending(pool_inner, context)
                }
            }
        }
    }
}