diff options
author | hozan23 <hozan23@proton.me> | 2023-11-08 13:03:27 +0300 |
---|---|---|
committer | hozan23 <hozan23@proton.me> | 2023-11-08 13:03:27 +0300 |
commit | 4fe665fc8bc6265baf5bfba6b6a5f3ee2dba63dc (patch) | |
tree | 77c7c40c9725539546e53b00f424deafe5ec81a8 /karyons_core/src/async_utils/select.rs |
first commit
Diffstat (limited to 'karyons_core/src/async_utils/select.rs')
-rw-r--r-- | karyons_core/src/async_utils/select.rs | 99 |
1 files changed, 99 insertions, 0 deletions
diff --git a/karyons_core/src/async_utils/select.rs b/karyons_core/src/async_utils/select.rs new file mode 100644 index 0000000..d61b355 --- /dev/null +++ b/karyons_core/src/async_utils/select.rs @@ -0,0 +1,99 @@ +use std::pin::Pin; +use std::task::{Context, Poll}; + +use pin_project_lite::pin_project; +use smol::future::Future; + +/// Returns the result of the future that completes first, preferring future1 +/// if both are ready. +/// +/// # Examples +/// +/// ``` +/// use std::future; +/// +/// use karyons_core::async_utils::{select, Either}; +/// +/// async { +/// let fut1 = future::pending::<String>(); +/// let fut2 = future::ready(0); +/// let res = select(fut1, fut2).await; +/// assert!(matches!(res, Either::Right(0))); +/// // .... +/// }; +/// +/// ``` +/// +pub fn select<T1, T2, F1, F2>(future1: F1, future2: F2) -> Select<F1, F2> +where + F1: Future<Output = T1>, + F2: Future<Output = T2>, +{ + Select { future1, future2 } +} + +pin_project! { + #[derive(Debug)] + pub struct Select<F1, F2> { + #[pin] + future1: F1, + #[pin] + future2: F2, + } +} + +/// The return value from the `select` function, indicating which future +/// completed first. +#[derive(Debug)] +pub enum Either<T1, T2> { + Left(T1), + Right(T2), +} + +// Implement the Future trait for the Select struct. +impl<T1, T2, F1, F2> Future for Select<F1, F2> +where + F1: Future<Output = T1>, + F2: Future<Output = T2>, +{ + type Output = Either<T1, T2>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = self.project(); + + if let Poll::Ready(t) = this.future1.poll(cx) { + return Poll::Ready(Either::Left(t)); + } + + if let Poll::Ready(t) = this.future2.poll(cx) { + return Poll::Ready(Either::Right(t)); + } + + Poll::Pending + } +} + +#[cfg(test)] +mod tests { + use super::{select, Either}; + use smol::Timer; + use std::future; + + #[test] + fn test_async_select() { + smol::block_on(async move { + let fut = select(Timer::never(), future::ready(0 as u32)).await; + assert!(matches!(fut, Either::Right(0))); + + let fut1 = future::pending::<String>(); + let fut2 = future::ready(0); + let res = select(fut1, fut2).await; + assert!(matches!(res, Either::Right(0))); + + let fut1 = future::ready(0); + let fut2 = future::pending::<String>(); + let res = select(fut1, fut2).await; + assert!(matches!(res, Either::Left(_))); + }); + } +} |