ntex-util 4.2.1

//! Future and stream utilities.
use std::future::{Future, poll_fn};
use std::task::{Context, Poll, ready};
use std::{mem, pin::Pin};

pub use futures_core::{Stream, TryFuture};

mod either;
mod join;
mod lazy;
mod on_drop;
mod select;

pub use self::either::Either;
pub use self::join::{join, join_all};
pub use self::lazy::{Lazy, lazy};
pub use self::on_drop::{OnDropFn, OnDropFuture, OnDropFutureExt};
pub use self::select::select;

/// An owned, dynamically dispatched future.
///
/// This is useful when a concrete future type cannot be named or when type
/// erasure is required.
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;

/// Waits for and returns the next item from a stream.
pub async fn stream_recv<S>(stream: &mut S) -> Option<S::Item>
where
    S: Stream + Unpin,
{
    poll_fn(|cx| Pin::new(&mut *stream).poll_next(cx)).await
}

enum MaybeDone<F>
where
    F: Future,
{
    Pending(F),
    Done(F::Output),
    Gone,
}

impl<F: Future> MaybeDone<F> {
    fn take_output(self: Pin<&mut Self>) -> Option<F::Output> {
        match &*self {
            Self::Done(_) => {}
            Self::Pending(_) | Self::Gone => return None,
        }
        unsafe {
            match mem::replace(self.get_unchecked_mut(), Self::Gone) {
                MaybeDone::Done(output) => Some(output),
                _ => unreachable!(),
            }
        }
    }
}

impl<F: Future> Future for MaybeDone<F> {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        unsafe {
            match self.as_mut().get_unchecked_mut() {
                MaybeDone::Pending(f) => {
                    let res = ready!(Pin::new_unchecked(f).poll(cx));
                    self.set(Self::Done(res));
                }
                MaybeDone::Done(_) => {}
                MaybeDone::Gone => panic!("MaybeDone polled after value taken"),
            }
        }
        Poll::Ready(())
    }
}

#[cfg(test)]
mod tests {
    use std::{future::pending, future::ready, pin::pin};

    use super::*;

    #[ntex::test]
    async fn maybe_done_take_output() {
        let mut fut = pin!(MaybeDone::Pending(pending::<()>()));
        assert!(fut.as_mut().take_output().is_none());

        let mut fut = pin!(MaybeDone::Pending(ready(1)));
        fut.as_mut().await;
        assert_eq!(fut.as_mut().take_output(), Some(1));
        // the output is taken once
        assert!(fut.as_mut().take_output().is_none());
    }

    #[ntex::test]
    #[should_panic(expected = "MaybeDone polled after value taken")]
    async fn maybe_done_poll_after_take() {
        let mut fut = pin!(MaybeDone::Pending(ready(1)));
        fut.as_mut().await;
        assert_eq!(fut.as_mut().take_output(), Some(1));
        fut.as_mut().await;
    }
}