From 6bf02aa1dfd1fd7951080ac010732461fb676b0b Mon Sep 17 00:00:00 2001 From: Asa Paparo Date: Sat, 19 Jul 2025 01:44:18 -0400 Subject: [PATCH] add poll_fn --- futures-combinators/src/join.rs | 69 +++++++++++---------------------- futures-util/src/lib.rs | 2 + futures-util/src/maybe_done.rs | 1 + futures-util/src/poll_fn.rs | 55 ++++++++++++++++++++++++++ 4 files changed, 81 insertions(+), 46 deletions(-) create mode 100644 futures-util/src/poll_fn.rs diff --git a/futures-combinators/src/join.rs b/futures-combinators/src/join.rs index f20637f..597c1a0 100644 --- a/futures-combinators/src/join.rs +++ b/futures-combinators/src/join.rs @@ -1,10 +1,8 @@ use futures_core::{ScopedFuture, Wake}; use futures_util::{MaybeDone, MaybeDoneState, maybe_done}; -use std::cell::{Cell, UnsafeCell}; -use std::sync::atomic::Ordering; -use std::{sync::atomic::AtomicBool, task::Poll}; +use std::{cell::Cell, task::Poll}; -/// from yoshuawuyts/futures-concurrency +/// from [futures-concurrency](https://github.com/yoshuawuyts/futures-concurrency/tree/main) /// Wait for all futures to complete. /// /// Awaits multiple futures simultaneously, returning the output of the futures @@ -12,59 +10,22 @@ use std::{sync::atomic::AtomicBool, task::Poll}; pub trait Join<'scope> { /// The resulting output type. type Output; + /// The [`ScopedFuture`] implementation returned by this method. type Future: ScopedFuture<'scope, Output = Self::Output>; + /// Waits for multiple futures to complete. /// /// Awaits multiple futures simultaneously, returning the output of the futures /// in the same container type they we're created once all complete. /// - /// # Examples - /// - /// Awaiting multiple futures of the same type can be done using either a vector - /// or an array. - /// ```rust - /// # futures::executor::block_on(async { - /// use futures_concurrency::prelude::*; - /// - /// // all futures passed here are of the same type - /// let fut1 = core::future::ready(1); - /// let fut2 = core::future::ready(2); - /// let fut3 = core::future::ready(3); - /// - /// let outputs = [fut1, fut2, fut3].join().await; - /// assert_eq!(outputs, [1, 2, 3]); - /// # }) - /// ``` - /// - /// In practice however, it's common to want to await multiple futures of - /// different types. For example if you have two different `async {}` blocks, - /// you want to `.await`. To do that, you can call `.join` on tuples of futures. - /// ```rust - /// # futures::executor::block_on(async { - /// use futures_concurrency::prelude::*; - /// - /// async fn some_async_fn() -> usize { 3 } - /// - /// // the futures passed here are of different types - /// let fut1 = core::future::ready(1); - /// let fut2 = async { 2 }; - /// let fut3 = some_async_fn(); - /// // ^ NOTE: no `.await` here! - /// - /// let outputs = (fut1, fut2, fut3).join().await; - /// assert_eq!(outputs, (1, 2, 3)); - /// # }) - /// ``` - /// - ///

/// This function returns a new future which polls all futures concurrently. fn join(self) -> Self::Future; } struct WakeStore<'scope> { parent: Cell>>, - ready: AtomicBool, + ready: Cell, } impl<'scope> WakeStore<'scope> { @@ -75,13 +36,13 @@ impl<'scope> WakeStore<'scope> { } } fn take_ready(&self) -> bool { - self.ready.swap(false, Ordering::SeqCst) + self.ready.replace(false) } } impl<'scope> Wake<'scope> for WakeStore<'scope> { fn wake(&self) { - self.ready.swap(true, Ordering::SeqCst); + self.ready.replace(true); if let Some(parent) = &self.parent.get() { parent.wake(); } @@ -107,6 +68,7 @@ macro_rules! impl_join_tuple { } #[allow(non_snake_case)] + #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct $StructName<'scope, $($F: ScopedFuture<'scope>),+> { $($F: MaybeDone<'scope, $F>,)* wakers: $namespace::Wakers<'scope>, @@ -190,3 +152,18 @@ impl_join_tuple!(join9 Join9 A B C D E F G H I); impl_join_tuple!(join10 Join10 A B C D E F G H I J); impl_join_tuple!(join11 Join11 A B C D E F G H I J K); impl_join_tuple!(join12 Join12 A B C D E F G H I J K L); + +#[cfg(test)] +mod tests { + use futures_util::poll_fn; + + use super::*; + + #[test] + fn basic() { + let f1 = poll_fn(|_| Poll::Ready(1)); + let f2 = poll_fn(|_| Poll::Ready(2)); + let dummy_waker = WakeStore::new(); + assert_eq!((f1, f2).join().poll(&dummy_waker), Poll::Ready((1, 2))); + } +} diff --git a/futures-util/src/lib.rs b/futures-util/src/lib.rs index 7c166fa..060525b 100644 --- a/futures-util/src/lib.rs +++ b/futures-util/src/lib.rs @@ -1,7 +1,9 @@ mod maybe_done; +mod poll_fn; use futures_core::ScopedFuture; pub use maybe_done::*; +pub use poll_fn::poll_fn; // Just a helper function to ensure the futures we're returning all have the // right implementations. diff --git a/futures-util/src/maybe_done.rs b/futures-util/src/maybe_done.rs index c65d612..5af76ce 100644 --- a/futures-util/src/maybe_done.rs +++ b/futures-util/src/maybe_done.rs @@ -8,6 +8,7 @@ use std::{ task::{Poll, ready}, }; +#[must_use = "futures do nothing unless you `.await` or poll them"] pub struct MaybeDone<'scope, Fut: ScopedFuture<'scope>> { state: UnsafeCell>, } diff --git a/futures-util/src/poll_fn.rs b/futures-util/src/poll_fn.rs new file mode 100644 index 0000000..1e2b116 --- /dev/null +++ b/futures-util/src/poll_fn.rs @@ -0,0 +1,55 @@ +use core::fmt; +use std::task::Poll; + +use futures_core::{ScopedFuture, Wake}; + +use crate::assert_future; + +/// Future for the [`poll_fn`] function. +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct PollFn { + f: F, +} + +/// Creates a new future wrapping around a function returning [`Poll`]. +/// +/// Polling the returned future delegates to the wrapped function. +/// +/// # Examples +/// +/// ``` +/// # futures::executor::block_on(async { +/// use futures::future::poll_fn; +/// use futures::task::{Context, Poll}; +/// +/// fn read_line(_cx: &mut Context<'_>) -> Poll { +/// Poll::Ready("Hello, World!".into()) +/// } +/// +/// let read_future = poll_fn(read_line); +/// assert_eq!(read_future.await, "Hello, World!".to_owned()); +/// # }); +/// ``` +pub fn poll_fn<'scope, T, F>(f: F) -> PollFn +where + F: Fn(&'scope dyn Wake) -> Poll, +{ + assert_future::(PollFn { f }) +} + +impl fmt::Debug for PollFn { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PollFn").finish() + } +} + +impl<'scope, T, F> ScopedFuture<'scope> for PollFn +where + F: Fn(&'scope dyn Wake) -> Poll, +{ + type Output = T; + + fn poll(&self, wake: &'scope dyn Wake) -> Poll { + (&self.f)(wake) + } +} -- 2.51.2