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, Waker},
};

use super::{Pool, hint};

/// The future returned by [`Pool::close`].
///
/// It owns the [`Pool`]. It transfers ownership back to the caller once
/// resolved. See [`Pool::close`] for an example.
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Close<T, B> {
    /// The pool. This struct owns the [`Pool`], so that no one else can use
    /// it.
    ///
    /// It should ideally be using `MaybeUninit`. However, if `<Close 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: Option<Pool<T, B>>,
}

impl<T, B> Close<T, B> {
    /// Close the pool and create a new [`Self`].
    ///
    /// This method is used by [`Pool::close`].
    pub(super) fn new(pool: Pool<T, B>) -> Self {
        pool.inner.close();

        Self { pool: Some(pool) }
    }
}

impl<T, B> Future for Close<T, B> {
    type Output = Pool<T, B>;

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

            unreachable!(
                "The `Close` 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 mut closing_state = loop {
                if let Some(closing_state) = self
                    .pool
                    .as_ref()
                    .unwrap_unchecked()
                    .inner
                    .closing_state
                    .try_write_lock()
                {
                    break closing_state;
                }

                hint::cold_path();
                hint::spin_loop();
            };

            let all_vacants = self
                .pool
                .as_ref()
                .unwrap_unchecked()
                .inner
                .blocks
                .iter()
                .all(|slot| slot.is_all_vacants());

            if all_vacants {
                // At this step, we know that all slots are vacant. We also know
                // that, when `Close` has been created, the pool was closed,
                // which means that it's impossible for a slot
                // to be occupied once all slots have been
                // vacant. Finally, because `Pool::close` takes
                // ownership of the `Pool`, it's impossible to reopen it. So we
                // are good!

                *closing_state = ClosingState::Closed;
                drop(closing_state);

                Poll::Ready(self.pool.take().unwrap_unchecked())
            } else {
                // At this step, not all slots are vacant.

                // We can register the waker and mark the future as pending.
                let waker = context.waker().clone();

                *closing_state = ClosingState::Closing(Some(waker));
                drop(closing_state);

                Poll::Pending
            }
        }
    }
}

pub(super) enum ClosingState {
    /// Pool is opened.
    Opened,

    /// [`Close`] has been built (if `None`) or is being resolved (if `Some`).
    Closing(Option<Waker>),

    /// [`Close`] has resolved.
    Closed,
}