actix_utils/future/
poll_fn.rs

1//! Simple "poll function" future and factory.
2
3use core::{
4    fmt,
5    future::Future,
6    pin::Pin,
7    task::{Context, Poll},
8};
9
10/// Creates a future driven by the provided function that receives a task context.
11///
12/// # Examples
13/// ```
14/// # use std::task::Poll;
15/// # use actix_utils::future::poll_fn;
16/// # async fn test_poll_fn() {
17/// let res = poll_fn(|_| Poll::Ready(42)).await;
18/// assert_eq!(res, 42);
19///
20/// let mut i = 5;
21/// let res = poll_fn(|cx| {
22///     i -= 1;
23///
24///     if i > 0 {
25///         cx.waker().wake_by_ref();
26///         Poll::Pending
27///     } else {
28///         Poll::Ready(42)
29///     }
30/// })
31/// .await;
32/// assert_eq!(res, 42);
33/// # }
34/// # actix_rt::Runtime::new().unwrap().block_on(test_poll_fn());
35/// ```
36#[inline]
37pub fn poll_fn<F, T>(f: F) -> PollFn<F>
38where
39    F: FnMut(&mut Context<'_>) -> Poll<T>,
40{
41    PollFn { f }
42}
43
44/// Future for the [`poll_fn`] function.
45pub struct PollFn<F> {
46    f: F,
47}
48
49impl<F> fmt::Debug for PollFn<F> {
50    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
51        f.debug_struct("PollFn").finish()
52    }
53}
54
55impl<F, T> Future for PollFn<F>
56where
57    F: FnMut(&mut Context<'_>) -> Poll<T>,
58{
59    type Output = T;
60
61    #[inline]
62    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
63        // SAFETY: we are not moving out of the pinned field
64        // see https://github.com/rust-lang/rust/pull/102737
65        (unsafe { &mut self.get_unchecked_mut().f })(cx)
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use std::marker::PhantomPinned;
72
73    use super::*;
74
75    static_assertions::assert_impl_all!(PollFn<()>: Unpin);
76    static_assertions::assert_not_impl_all!(PollFn<PhantomPinned>: Unpin);
77
78    #[actix_rt::test]
79    async fn test_poll_fn() {
80        let res = poll_fn(|_| Poll::Ready(42)).await;
81        assert_eq!(res, 42);
82
83        let mut i = 5;
84        let res = poll_fn(|cx| {
85            i -= 1;
86
87            if i > 0 {
88                cx.waker().wake_by_ref();
89                Poll::Pending
90            } else {
91                Poll::Ready(42)
92            }
93        })
94        .await;
95        assert_eq!(res, 42);
96    }
97
98    // following soundness tests taken from https://github.com/tokio-rs/tokio/pull/5087
99
100    #[allow(dead_code)]
101    fn require_send<T: Send>(_t: &T) {}
102    #[allow(dead_code)]
103    fn require_sync<T: Sync>(_t: &T) {}
104
105    trait AmbiguousIfUnpin<A> {
106        fn some_item(&self) {}
107    }
108    impl<T: ?Sized> AmbiguousIfUnpin<()> for T {}
109    impl<T: ?Sized + Unpin> AmbiguousIfUnpin<[u8; 0]> for T {}
110
111    const _: fn() = || {
112        let pinned = std::marker::PhantomPinned;
113        let f = poll_fn(move |_| {
114            // Use `pinned` to take ownership of it.
115            let _ = &pinned;
116            std::task::Poll::Pending::<()>
117        });
118        require_send(&f);
119        require_sync(&f);
120        AmbiguousIfUnpin::some_item(&f);
121    };
122}
actix_utils/future/poll_fn.rs

actix_utils/future/
poll_fn.rs
