async_sink/ext/
unfold.rs

1use super::Sink;
2use crate::unfold_state::UnfoldState;
3use core::fmt;
4use core::future::Future;
5use core::pin::Pin;
6use core::task::{Context, Poll};
7
8/// Sink for the [`unfold`] function.
9#[must_use = "sinks do nothing unless polled"]
10pub struct Unfold<T, F, Fut> {
11    function: F,
12    state: UnfoldState<T, Fut>,
13}
14
15impl<T, F, Fut> Unfold<T, F, Fut> {
16    // Helper to get a mutable reference to the `function` field and a
17    // pinned mutable reference to the `state` field.
18    //
19    // # Safety
20    //
21    // This is `unsafe` because it returns a `Pin` to one of the fields of the
22    // struct. The caller must ensure that they don't move the struct while this
23    // `Pin` is in use.
24    unsafe fn project(self: Pin<&mut Self>) -> (&mut F, Pin<&mut UnfoldState<T, Fut>>) {
25        let this = self.get_unchecked_mut();
26        (&mut this.function, Pin::new_unchecked(&mut this.state))
27    }
28}
29
30impl<T, F, Fut> fmt::Debug for Unfold<T, F, Fut>
31where
32    T: fmt::Debug,
33    Fut: fmt::Debug,
34{
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        f.debug_struct("Unfold")
37            .field("state", &self.state)
38            .finish()
39    }
40}
41
42/// Create a sink from a function which processes one item at a time.
43///
44/// # Examples
45///
46/// ```
47/// # #[tokio::main]
48/// # async fn main() -> Result<(), core::convert::Infallible> {
49/// use core::pin::pin;
50///
51/// use crate::sink::{self, SinkExt};
52///
53/// let unfold = sink::unfold(0, |mut sum, i: i32| {
54///     async move {
55///         sum += i;
56///         eprintln!("{}", i);
57///         Ok::<_, core::convert::Infallible>(sum)
58///     }
59/// });
60/// let mut unfold = pin!(unfold);
61/// unfold.send(5).await?;
62/// # Ok(())
63/// # }
64/// ```
65pub fn unfold<T, F, Fut, Item, E>(init: T, function: F) -> Unfold<T, F, Fut>
66where
67    F: FnMut(T, Item) -> Fut,
68    Fut: Future<Output = Result<T, E>>,
69{
70    Unfold {
71        function,
72        state: UnfoldState::Value { value: init },
73    }
74}
75
76impl<T, F, Fut, Item, E> Sink<Item> for Unfold<T, F, Fut>
77where
78    F: FnMut(T, Item) -> Fut,
79    Fut: Future<Output = Result<T, E>>,
80{
81    type Error = E;
82
83    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
84        self.poll_flush(cx)
85    }
86
87    fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error> {
88        let (function, state_pin) = unsafe { self.project() };
89        let state_mut = unsafe { state_pin.get_unchecked_mut() };
90
91        let value = match state_mut {
92            UnfoldState::Value { .. } => {
93                if let UnfoldState::Value { value } = unsafe { core::ptr::read(state_mut) } {
94                    value
95                } else {
96                    unreachable!()
97                }
98            }
99            _ => panic!("start_send called without poll_ready being called first"),
100        };
101
102        let future = function(value, item);
103        unsafe { core::ptr::write(state_mut, UnfoldState::Future { future }) };
104        Ok(())
105    }
106
107    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
108        let (_, state_pin) = unsafe { self.project() };
109        let state_mut = unsafe { state_pin.get_unchecked_mut() };
110
111        if let UnfoldState::Future { future } = state_mut {
112            let result = match unsafe { Pin::new_unchecked(future) }.poll(cx) {
113                Poll::Ready(result) => result,
114                Poll::Pending => return Poll::Pending,
115            };
116
117            // The future is finished, so we can replace the state.
118            // First, destruct the old state.
119            let _old_state = unsafe { core::ptr::read(state_mut) };
120
121            match result {
122                Ok(state) => {
123                    unsafe { core::ptr::write(state_mut, UnfoldState::Value { value: state }) };
124                    Poll::Ready(Ok(()))
125                }
126                Err(err) => {
127                    unsafe { core::ptr::write(state_mut, UnfoldState::Empty) };
128                    Poll::Ready(Err(err))
129                }
130            }
131        } else {
132            Poll::Ready(Ok(()))
133        }
134    }
135
136    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
137        self.poll_flush(cx)
138    }
139}
async_sink/ext/unfold.rs

async_sink/ext/
unfold.rs
