1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153
use crate::fmt;
use crate::future::Future;
use crate::pin::Pin;
use crate::task::{Context, Poll};
/// Creates a future that wraps a function returning [`Poll`].
///
/// Polling the future delegates to the wrapped function. If the returned future is pinned, then the
/// captured environment of the wrapped function is also pinned in-place, so as long as the closure
/// does not move out of its captures it can soundly create pinned references to them.
///
/// # Examples
///
/// ```
/// # async fn run() {
/// use core::future::poll_fn;
/// use std::task::{Context, Poll};
///
/// fn read_line(_cx: &mut Context<'_>) -> Poll<String> {
/// Poll::Ready("Hello, World!".into())
/// }
///
/// let read_future = poll_fn(read_line);
/// assert_eq!(read_future.await, "Hello, World!".to_owned());
/// # }
/// ```
///
/// ## Capturing a pinned state
///
/// Example of a closure wrapping inner futures:
///
/// ```
/// # async fn run() {
/// use core::future::{self, Future};
/// use core::task::Poll;
///
/// /// Resolves to the first future that completes. In the event of a tie, `a` wins.
/// fn naive_select<T>(
/// a: impl Future<Output = T>,
/// b: impl Future<Output = T>,
/// ) -> impl Future<Output = T>
/// {
/// let (mut a, mut b) = (Box::pin(a), Box::pin(b));
/// future::poll_fn(move |cx| {
/// if let Poll::Ready(r) = a.as_mut().poll(cx) {
/// Poll::Ready(r)
/// } else if let Poll::Ready(r) = b.as_mut().poll(cx) {
/// Poll::Ready(r)
/// } else {
/// Poll::Pending
/// }
/// })
/// }
///
/// let a = async { 42 };
/// let b = future::pending();
/// let v = naive_select(a, b).await;
/// assert_eq!(v, 42);
///
/// let a = future::pending();
/// let b = async { 27 };
/// let v = naive_select(a, b).await;
/// assert_eq!(v, 27);
///
/// let a = async { 42 };
/// let b = async { 27 };
/// let v = naive_select(a, b).await;
/// assert_eq!(v, 42); // biased towards `a` in case of tie!
/// # }
/// ```
///
/// This time without [`Box::pin`]ning:
///
/// [`Box::pin`]: ../../std/boxed/struct.Box.html#method.pin
///
/// ```
/// # async fn run() {
/// use core::future::{self, Future};
/// use core::pin::pin;
/// use core::task::Poll;
///
/// /// Resolves to the first future that completes. In the event of a tie, `a` wins.
/// fn naive_select<T>(
/// a: impl Future<Output = T>,
/// b: impl Future<Output = T>,
/// ) -> impl Future<Output = T>
/// {
/// async {
/// let (mut a, mut b) = (pin!(a), pin!(b));
/// future::poll_fn(move |cx| {
/// if let Poll::Ready(r) = a.as_mut().poll(cx) {
/// Poll::Ready(r)
/// } else if let Poll::Ready(r) = b.as_mut().poll(cx) {
/// Poll::Ready(r)
/// } else {
/// Poll::Pending
/// }
/// }).await
/// }
/// }
///
/// let a = async { 42 };
/// let b = future::pending();
/// let v = naive_select(a, b).await;
/// assert_eq!(v, 42);
/// # }
/// ```
///
/// - Notice how, by virtue of being in an `async` context, we have been able to make the [`pin!`]
/// macro work, thereby avoiding any need for the `unsafe`
/// <code>[Pin::new_unchecked](&mut fut)</code> constructor.
///
/// [`pin!`]: crate::pin::pin!
#[stable(feature = "future_poll_fn", since = "1.64.0")]
pub fn poll_fn<T, F>(f: F) -> PollFn<F>
where
F: FnMut(&mut Context<'_>) -> Poll<T>,
{
PollFn { f }
}
/// A Future that wraps a function returning [`Poll`].
///
/// This `struct` is created by [`poll_fn()`]. See its
/// documentation for more.
#[must_use = "futures do nothing unless you `.await` or poll them"]
#[stable(feature = "future_poll_fn", since = "1.64.0")]
pub struct PollFn<F> {
f: F,
}
#[stable(feature = "future_poll_fn", since = "1.64.0")]
impl<F: Unpin> Unpin for PollFn<F> {}
#[stable(feature = "future_poll_fn", since = "1.64.0")]
impl<F> fmt::Debug for PollFn<F> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("PollFn").finish()
}
}
#[stable(feature = "future_poll_fn", since = "1.64.0")]
impl<T, F> Future for PollFn<F>
where
F: FnMut(&mut Context<'_>) -> Poll<T>,
{
type Output = T;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
// SAFETY: We are not moving out of the pinned field.
(unsafe { &mut self.get_unchecked_mut().f })(cx)
}
}