| 1 | use core::fmt; |
| 2 | use core::pin::Pin; |
| 3 | use futures_core::future::TryFuture; |
| 4 | use futures_core::ready; |
| 5 | use futures_core::stream::{FusedStream, Stream, TryStream}; |
| 6 | use futures_core::task::{Context, Poll}; |
| 7 | #[cfg (feature = "sink" )] |
| 8 | use futures_sink::Sink; |
| 9 | use pin_project_lite::pin_project; |
| 10 | |
| 11 | pin_project! { |
| 12 | /// Stream for the [`try_filter_map`](super::TryStreamExt::try_filter_map) |
| 13 | /// method. |
| 14 | #[must_use = "streams do nothing unless polled" ] |
| 15 | pub struct TryFilterMap<St, Fut, F> { |
| 16 | #[pin] |
| 17 | stream: St, |
| 18 | f: F, |
| 19 | #[pin] |
| 20 | pending: Option<Fut>, |
| 21 | } |
| 22 | } |
| 23 | |
| 24 | impl<St, Fut, F> fmt::Debug for TryFilterMap<St, Fut, F> |
| 25 | where |
| 26 | St: fmt::Debug, |
| 27 | Fut: fmt::Debug, |
| 28 | { |
| 29 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
| 30 | f.debug_struct("TryFilterMap" ) |
| 31 | .field("stream" , &self.stream) |
| 32 | .field("pending" , &self.pending) |
| 33 | .finish() |
| 34 | } |
| 35 | } |
| 36 | |
| 37 | impl<St, Fut, F> TryFilterMap<St, Fut, F> { |
| 38 | pub(super) fn new(stream: St, f: F) -> Self { |
| 39 | Self { stream, f, pending: None } |
| 40 | } |
| 41 | |
| 42 | delegate_access_inner!(stream, St, ()); |
| 43 | } |
| 44 | |
| 45 | impl<St, Fut, F, T> FusedStream for TryFilterMap<St, Fut, F> |
| 46 | where |
| 47 | St: TryStream + FusedStream, |
| 48 | Fut: TryFuture<Ok = Option<T>, Error = St::Error>, |
| 49 | F: FnMut(St::Ok) -> Fut, |
| 50 | { |
| 51 | fn is_terminated(&self) -> bool { |
| 52 | self.pending.is_none() && self.stream.is_terminated() |
| 53 | } |
| 54 | } |
| 55 | |
| 56 | impl<St, Fut, F, T> Stream for TryFilterMap<St, Fut, F> |
| 57 | where |
| 58 | St: TryStream, |
| 59 | Fut: TryFuture<Ok = Option<T>, Error = St::Error>, |
| 60 | F: FnMut(St::Ok) -> Fut, |
| 61 | { |
| 62 | type Item = Result<T, St::Error>; |
| 63 | |
| 64 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { |
| 65 | let mut this = self.project(); |
| 66 | |
| 67 | Poll::Ready(loop { |
| 68 | if let Some(p) = this.pending.as_mut().as_pin_mut() { |
| 69 | // We have an item in progress, poll that until it's done |
| 70 | let res = ready!(p.try_poll(cx)); |
| 71 | this.pending.set(None); |
| 72 | let item = res?; |
| 73 | if item.is_some() { |
| 74 | break item.map(Ok); |
| 75 | } |
| 76 | } else if let Some(item) = ready!(this.stream.as_mut().try_poll_next(cx)?) { |
| 77 | // No item in progress, but the stream is still going |
| 78 | this.pending.set(Some((this.f)(item))); |
| 79 | } else { |
| 80 | // The stream is done |
| 81 | break None; |
| 82 | } |
| 83 | }) |
| 84 | } |
| 85 | |
| 86 | fn size_hint(&self) -> (usize, Option<usize>) { |
| 87 | let pending_len = usize::from(self.pending.is_some()); |
| 88 | let (_, upper) = self.stream.size_hint(); |
| 89 | let upper = match upper { |
| 90 | Some(x) => x.checked_add(pending_len), |
| 91 | None => None, |
| 92 | }; |
| 93 | (0, upper) // can't know a lower bound, due to the predicate |
| 94 | } |
| 95 | } |
| 96 | |
| 97 | // Forwarding impl of Sink from the underlying stream |
| 98 | #[cfg (feature = "sink" )] |
| 99 | impl<S, Fut, F, Item> Sink<Item> for TryFilterMap<S, Fut, F> |
| 100 | where |
| 101 | S: Sink<Item>, |
| 102 | { |
| 103 | type Error = S::Error; |
| 104 | |
| 105 | delegate_sink!(stream, Item); |
| 106 | } |
| 107 | |