multipart_write/stream/
assemble.rs

1use crate::MultipartWrite;
2
3use futures_core::future::{FusedFuture, Future};
4use futures_core::ready;
5use futures_core::stream::Stream;
6use std::fmt::{self, Debug, Formatter};
7use std::pin::Pin;
8use std::task::{Context, Poll};
9
10pin_project_lite::pin_project! {
11    /// Future for [`assemble`].
12    ///
13    /// [`assemble`]: super::MultipartStreamExt::assemble
14    #[must_use = "futures do nothing unless polled"]
15    pub struct Assemble<St: Stream, Wr> {
16        #[pin]
17        writer: Wr,
18        #[pin]
19        stream: Option<St>,
20        buffered: Option<St::Item>,
21        is_terminated: bool,
22    }
23}
24
25impl<St: Stream, Wr> Assemble<St, Wr> {
26    pub(super) fn new(stream: St, writer: Wr) -> Self {
27        Self {
28            writer,
29            stream: Some(stream),
30            buffered: None,
31            is_terminated: false,
32        }
33    }
34}
35
36impl<St, Wr> FusedFuture for Assemble<St, Wr>
37where
38    St: Stream,
39    Wr: MultipartWrite<St::Item>,
40{
41    fn is_terminated(&self) -> bool {
42        self.is_terminated
43    }
44}
45
46impl<St, Wr> Future for Assemble<St, Wr>
47where
48    St: Stream,
49    Wr: MultipartWrite<St::Item>,
50{
51    type Output = Result<Wr::Output, Wr::Error>;
52
53    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
54        let mut this = self.project();
55
56        loop {
57            if this.buffered.is_some() {
58                match this.writer.as_mut().poll_ready(cx)? {
59                    Poll::Pending => return Poll::Pending,
60                    Poll::Ready(()) => {
61                        let _ = this
62                            .writer
63                            .as_mut()
64                            .start_send(this.buffered.take().unwrap())?;
65                    }
66                }
67            }
68
69            let Some(st) = this.stream.as_mut().as_pin_mut() else {
70                let out = ready!(this.writer.as_mut().poll_complete(cx));
71                *this.is_terminated = true;
72                return Poll::Ready(out);
73            };
74
75            match ready!(st.poll_next(cx)) {
76                Some(it) => *this.buffered = Some(it),
77                None => this.stream.set(None),
78            }
79        }
80    }
81}
82
83impl<St, Wr> Debug for Assemble<St, Wr>
84where
85    St: Stream + Debug,
86    St::Item: Debug,
87    Wr: Debug,
88{
89    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
90        f.debug_struct("Assemble")
91            .field("writer", &self.writer)
92            .field("stream", &self.stream)
93            .field("buffered", &self.buffered)
94            .field("is_terminated", &self.is_terminated)
95            .finish()
96    }
97}