vector/
sink_ext.rs

1// Copyright (c) 2016 Alex Crichton
2// Copyright (c) 2017 The Tokio Authors
3
4// Permission is hereby granted, free of charge, to any
5// person obtaining a copy of this software and associated
6// documentation files (the "Software"), to deal in the
7// Software without restriction, including without
8// limitation the rights to use, copy, modify, merge,
9// publish, distribute, sublicense, and/or sell copies of
10// the Software, and to permit persons to whom the Software
11// is furnished to do so, subject to the following
12// conditions:
13
14// The above copyright notice and this permission notice
15// shall be included in all copies or substantial portions
16// of the Software.
17
18// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
19// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
20// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
21// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
22// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
23// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
24// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
25// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
26// DEALINGS IN THE SOFTWARE.
27
28use std::{
29    future::Future,
30    pin::Pin,
31    task::{ready, Context, Poll},
32};
33
34use futures::{stream::Peekable, Sink, SinkExt, Stream, StreamExt};
35
36impl<T: ?Sized, Item> VecSinkExt<Item> for T where T: Sink<Item> {}
37
38pub(crate) trait VecSinkExt<Item>: Sink<Item> {
39    /// A future that completes after the given stream has been fully processed
40    /// into the sink, including flushing.
41    /// Compare to `SinkExt::send_all` this future accept `Peekable` stream and
42    /// do not have own buffer.
43    fn send_all_peekable<'a, St>(
44        &'a mut self,
45        stream: &'a mut Peekable<St>,
46    ) -> SendAll<'a, Self, St>
47    where
48        St: Stream<Item = Item> + Sized,
49        Self: Sized,
50    {
51        SendAll { sink: self, stream }
52    }
53}
54
55/// Future for the [`send_all_peekable`](VecSinkExt::send_all_peekable) method.
56pub(crate) struct SendAll<'a, Si, St>
57where
58    St: Stream,
59{
60    sink: &'a mut Si,
61    stream: &'a mut Peekable<St>,
62}
63
64impl<Si, St, Item, Error> Future for SendAll<'_, Si, St>
65where
66    Si: Sink<Item, Error = Error> + Unpin,
67    St: Stream<Item = Item> + Unpin,
68{
69    type Output = Result<(), Error>;
70
71    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
72        loop {
73            match Pin::new(&mut *self.stream).as_mut().poll_peek(cx) {
74                Poll::Ready(Some(_)) => {
75                    ready!(self.sink.poll_ready_unpin(cx))?;
76                    let item = match self.stream.poll_next_unpin(cx) {
77                        Poll::Ready(Some(item)) => item,
78                        _ => panic!("Item should exist after poll_peek succeeds"),
79                    };
80                    self.sink.start_send_unpin(item)?;
81                }
82                Poll::Ready(None) => {
83                    ready!(self.sink.poll_flush_unpin(cx))?;
84                    return Poll::Ready(Ok(()));
85                }
86                Poll::Pending => {
87                    ready!(self.sink.poll_flush_unpin(cx))?;
88                    return Poll::Pending;
89                }
90            }
91        }
92    }
93}
vector/sink_ext.rs

vector/
sink_ext.rs
