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

pub use std::{future::Future, pin::Pin, task::Poll};

pub use futures_util_fork::{maybe_done, poll_fn, MaybeDone, PollFn};

pub use crate::utils::thread_rng_n;

mod futures_util_fork {
    use core::{fmt, mem, pin::Pin};
    use std::{
        future::Future,
        task::{Context, Poll},
    };

    /// Future for the [`poll_fn`] function.
    #[must_use = "futures do nothing unless you `.await` or poll them"]
    pub struct PollFn<F> {
        f: F,
    }

    impl<F> Unpin for PollFn<F> {}

    /// Creates a new future wrapping around a function returning [`Poll`].
    ///
    /// Polling the returned future delegates to the wrapped function.
    ///
    /// # Examples
    ///
    /// ```
    /// # futures::executor::block_on(async {
    /// use futures::{
    ///     future::poll_fn,
    ///     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());
    /// # });
    /// ```
    pub fn poll_fn<T, F>(f: F) -> PollFn<F>
    where
        F: FnMut(&mut Context<'_>) -> Poll<T>,
    {
        PollFn { f }
    }

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

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

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
            (self.f)(cx)
        }
    }

    /// A future that may have completed.
    ///
    /// This is created by the [`maybe_done()`] function.
    #[derive(Debug)]
    pub enum MaybeDone<Fut: Future> {
        /// A not-yet-completed future
        Future(/* #[pin] */ Fut),
        /// The output of the completed future
        Done(Fut::Output),
        /// The empty variant after the result of a [`MaybeDone`] has been
        /// taken using the [`take_output`](MaybeDone::take_output) method.
        Gone,
    }

    impl<Fut: Future + Unpin> Unpin for MaybeDone<Fut> {}

    /// Wraps a future into a `MaybeDone`
    ///
    /// # Examples
    ///
    /// ```
    /// # futures::executor::block_on(async {
    /// use futures::{future, pin_mut};
    ///
    /// let future = future::maybe_done(async { 5 });
    /// pin_mut!(future);
    /// assert_eq!(future.as_mut().take_output(), None);
    /// let () = future.as_mut().await;
    /// assert_eq!(future.as_mut().take_output(), Some(5));
    /// assert_eq!(future.as_mut().take_output(), None);
    /// # });
    /// ```
    pub fn maybe_done<Fut: Future>(future: Fut) -> MaybeDone<Fut> {
        MaybeDone::Future(future)
    }

    impl<Fut: Future> MaybeDone<Fut> {
        /// Returns an [`Option`] containing a mutable reference to the output
        /// of the future. The output of this method will be [`Some`] if
        /// and only if the inner future has been completed and
        /// [`take_output`](MaybeDone::take_output) has not yet been
        /// called.
        #[inline]
        pub fn output_mut(self: Pin<&mut Self>) -> Option<&mut Fut::Output> {
            unsafe {
                match self.get_unchecked_mut() {
                    MaybeDone::Done(res) => Some(res),
                    _ => None,
                }
            }
        }

        /// Attempt to take the output of a `MaybeDone` without driving it
        /// towards completion.
        #[inline]
        pub fn take_output(self: Pin<&mut Self>) -> Option<Fut::Output> {
            match &*self {
                Self::Done(_) => {}
                Self::Future(_) | Self::Gone => return None,
            }
            unsafe {
                match mem::replace(self.get_unchecked_mut(), Self::Gone) {
                    MaybeDone::Done(output) => Some(output),
                    _ => unreachable!(),
                }
            }
        }
    }

    impl<Fut: Future> Future for MaybeDone<Fut> {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            unsafe {
                match self.as_mut().get_unchecked_mut() {
                    MaybeDone::Future(f) => {
                        let res = ready!(Pin::new_unchecked(f).poll(cx));
                        self.set(Self::Done(res));
                    }
                    MaybeDone::Done(_) => {}
                    MaybeDone::Gone => panic!("MaybeDone polled after value taken"),
                }
            }
            Poll::Ready(())
        }
    }
}