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/// use core::pin::pin;
48/// use async_sink::SinkExt;
49/// use tokio::sync::Mutex;
50/// use std::sync::Arc;
51///
52/// #[tokio::main]
53/// async fn main() {
54/// let output: Arc<Mutex<Vec<usize>>> = Arc::new(tokio::sync::Mutex::new(Vec::new()));
55///
56/// let unfold = async_sink::unfold(0, |mut sum, i: usize| {
57///    let cb_output = output.clone();
58///    async move {
59///         sum += i;
60///         cb_output.clone().lock().await.push(sum);
61///         Ok::<_, core::convert::Infallible>(sum)
62///    }
63/// });
64/// let mut unfold = pin!(unfold);
65/// let input: [usize; 3] = [5, 15, 35];
66/// assert!(unfold.send_all(&mut tokio_stream::iter(input.iter().copied().map(|i| Ok(i)))).await.is_ok());
67/// assert_eq!(output.lock().await.as_slice(),input.iter().scan(0, |state, &x|
68///   { *state += x; Some(*state) }).collect::<Vec<usize>>().as_slice()
69/// );
70/// }
71/// ```
72pub fn unfold<T, F, Fut, Item, E>(init: T, function: F) -> Unfold<T, F, Fut>
73where
74    F: FnMut(T, Item) -> Fut,
75    Fut: Future<Output = Result<T, E>>,
76{
77    Unfold {
78        function,
79        state: UnfoldState::Value { value: init },
80    }
81}
82
83impl<T, F, Fut, Item, E> Sink<Item> for Unfold<T, F, Fut>
84where
85    F: FnMut(T, Item) -> Fut,
86    Fut: Future<Output = Result<T, E>>,
87{
88    type Error = E;
89
90    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
91        self.poll_flush(cx)
92    }
93
94    fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error> {
95        let (function, state_pin) = unsafe { self.project() };
96        let state_mut = unsafe { state_pin.get_unchecked_mut() };
97
98        let value = match state_mut {
99            UnfoldState::Value { .. } => {
100                if let UnfoldState::Value { value } = unsafe { core::ptr::read(state_mut) } {
101                    value
102                } else {
103                    unreachable!()
104                }
105            }
106            _ => panic!("start_send called without poll_ready being called first"),
107        };
108
109        let future = function(value, item);
110        unsafe { core::ptr::write(state_mut, UnfoldState::Future { future }) };
111        Ok(())
112    }
113
114    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
115        let (_, state_pin) = unsafe { self.project() };
116        let state_mut = unsafe { state_pin.get_unchecked_mut() };
117
118        if let UnfoldState::Future { future } = state_mut {
119            let result = match unsafe { Pin::new_unchecked(future) }.poll(cx) {
120                Poll::Ready(result) => result,
121                Poll::Pending => return Poll::Pending,
122            };
123
124            // The future is finished, so we can replace the state.
125            // First, destruct the old state.
126            let _old_state = unsafe { core::ptr::read(state_mut) };
127
128            match result {
129                Ok(state) => {
130                    unsafe { core::ptr::write(state_mut, UnfoldState::Value { value: state }) };
131                    Poll::Ready(Ok(()))
132                }
133                Err(err) => {
134                    unsafe { core::ptr::write(state_mut, UnfoldState::Empty) };
135                    Poll::Ready(Err(err))
136                }
137            }
138        } else {
139            Poll::Ready(Ok(()))
140        }
141    }
142
143    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
144        self.poll_flush(cx)
145    }
146}
async_sink/ext/unfold.rs

async_sink/ext/
unfold.rs
