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        #[allow(clippy::needless_borrow)]
66        (unsafe { &mut self.get_unchecked_mut().f })(cx)
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use std::marker::PhantomPinned;
73
74    use super::*;
75
76    static_assertions::assert_impl_all!(PollFn<()>: Unpin);
77    static_assertions::assert_not_impl_all!(PollFn<PhantomPinned>: Unpin);
78
79    #[actix_rt::test]
80    async fn test_poll_fn() {
81        let res = poll_fn(|_| Poll::Ready(42)).await;
82        assert_eq!(res, 42);
83
84        let mut i = 5;
85        let res = poll_fn(|cx| {
86            i -= 1;
87
88            if i > 0 {
89                cx.waker().wake_by_ref();
90                Poll::Pending
91            } else {
92                Poll::Ready(42)
93            }
94        })
95        .await;
96        assert_eq!(res, 42);
97    }
98
99    // following soundness tests taken from https://github.com/tokio-rs/tokio/pull/5087
100
101    #[allow(dead_code)]
102    fn require_send<T: Send>(_t: &T) {}
103    #[allow(dead_code)]
104    fn require_sync<T: Sync>(_t: &T) {}
105
106    #[allow(unused)]
107    trait AmbiguousIfUnpin<A> {
108        fn some_item(&self) {}
109    }
110    impl<T: ?Sized> AmbiguousIfUnpin<()> for T {}
111    impl<T: ?Sized + Unpin> AmbiguousIfUnpin<[u8; 0]> for T {}
112
113    const _: fn() = || {
114        let pinned = std::marker::PhantomPinned;
115        let f = poll_fn(move |_| {
116            // Use `pinned` to take ownership of it.
117            let _ = &pinned;
118            std::task::Poll::Pending::<()>
119        });
120        require_send(&f);
121        require_sync(&f);
122        AmbiguousIfUnpin::some_item(&f);
123    };
124}
actix_utils/future/poll_fn.rs

actix_utils/future/
poll_fn.rs
