#
actix_utils/future/
poll_fn.rs

    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

    //! Simple "poll function" future and factory.

    use core::{
        fmt,
        future::Future,
        pin::Pin,
        task::{Context, Poll},
    };

    /// Creates a future driven by the provided function that receives a task context.
    ///
    /// # Examples
    /// ```
    /// # use std::task::Poll;
    /// # use actix_utils::future::poll_fn;
    /// # async fn test_poll_fn() {
    /// let res = poll_fn(|_| Poll::Ready(42)).await;
    /// assert_eq!(res, 42);
    ///
    /// let mut i = 5;
    /// let res = poll_fn(|cx| {
    ///     i -= 1;
    ///
    ///     if i > 0 {
    ///         cx.waker().wake_by_ref();
    ///         Poll::Pending
    ///     } else {
    ///         Poll::Ready(42)
    ///     }
    /// })
    /// .await;
    /// assert_eq!(res, 42);
    /// # }
    /// # actix_rt::Runtime::new().unwrap().block_on(test_poll_fn());
    /// ```
    #[inline]
    pub fn poll_fn<F, T>(f: F) -> PollFn<F>
    where
        F: FnMut(&mut Context<'_>) -> Poll<T>,
    {
        PollFn { f }
    }

    /// Future for the [`poll_fn`] function.
    pub struct PollFn<F> {
        f: F,
    }

    impl<F> fmt::Debug for PollFn<F> {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("PollFn").finish()
        }
    }

    impl<F, T> Future for PollFn<F>
    where
        F: FnMut(&mut Context<'_>) -> Poll<T>,
    {
        type Output = T;

        #[inline]
        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            // SAFETY: we are not moving out of the pinned field
            // see https://github.com/rust-lang/rust/pull/102737
            #[allow(clippy::needless_borrow)]
            (unsafe { &mut self.get_unchecked_mut().f })(cx)
        }
    }

    #[cfg(test)]
    mod tests {
        use std::marker::PhantomPinned;

        use super::*;

        static_assertions::assert_impl_all!(PollFn<()>: Unpin);
        static_assertions::assert_not_impl_all!(PollFn<PhantomPinned>: Unpin);

        #[actix_rt::test]
        async fn test_poll_fn() {
            let res = poll_fn(|_| Poll::Ready(42)).await;
            assert_eq!(res, 42);

            let mut i = 5;
            let res = poll_fn(|cx| {
                i -= 1;

                if i > 0 {
                    cx.waker().wake_by_ref();
                    Poll::Pending
                } else {
                    Poll::Ready(42)
                }
            })
            .await;
            assert_eq!(res, 42);
        }

        // following soundness tests taken from https://github.com/tokio-rs/tokio/pull/5087

        #[allow(dead_code)]
        fn require_send<T: Send>(_t: &T) {}
        #[allow(dead_code)]
        fn require_sync<T: Sync>(_t: &T) {}

        #[allow(unused)]
        trait AmbiguousIfUnpin<A> {
            fn some_item(&self) {}
        }
        impl<T: ?Sized> AmbiguousIfUnpin<()> for T {}
        impl<T: ?Sized + Unpin> AmbiguousIfUnpin<[u8; 0]> for T {}

        const _: fn() = || {
            let pinned = std::marker::PhantomPinned;
            let f = poll_fn(move |_| {
                // Use `pinned` to take ownership of it.
                let _ = &pinned;
                std::task::Poll::Pending::<()>
            });
            require_send(&f);
            require_sync(&f);
            AmbiguousIfUnpin::some_item(&f);
        };
    }