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)
    }
}