1 use core::fmt; 2 use core::pin::Pin; 3 use futures_core::future::Future; 4 use futures_core::ready; 5 use futures_core::stream::{Stream, FusedStream}; 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 [`take_while`](super::StreamExt::take_while) method. 13 #[must_use = "streams do nothing unless polled"] 14 pub struct TakeWhile<St: Stream, Fut, F> { 15 #[pin] 16 stream: St, 17 f: F, 18 #[pin] 19 pending_fut: Option<Fut>, 20 pending_item: Option<St::Item>, 21 done_taking: bool, 22 } 23 } 24 25 impl<St, Fut, F> fmt::Debug for TakeWhile<St, Fut, F> 26 where 27 St: Stream + fmt::Debug, 28 St::Item: fmt::Debug, 29 Fut: fmt::Debug, 30 { fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result31 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 32 f.debug_struct("TakeWhile") 33 .field("stream", &self.stream) 34 .field("pending_fut", &self.pending_fut) 35 .field("pending_item", &self.pending_item) 36 .field("done_taking", &self.done_taking) 37 .finish() 38 } 39 } 40 41 impl<St, Fut, F> TakeWhile<St, Fut, F> 42 where St: Stream, 43 F: FnMut(&St::Item) -> Fut, 44 Fut: Future<Output = bool>, 45 { new(stream: St, f: F) -> Self46 pub(super) fn new(stream: St, f: F) -> Self { 47 Self { 48 stream, 49 f, 50 pending_fut: None, 51 pending_item: None, 52 done_taking: false, 53 } 54 } 55 56 delegate_access_inner!(stream, St, ()); 57 } 58 59 impl<St, Fut, F> Stream for TakeWhile<St, Fut, F> 60 where St: Stream, 61 F: FnMut(&St::Item) -> Fut, 62 Fut: Future<Output = bool>, 63 { 64 type Item = St::Item; 65 poll_next( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll<Option<St::Item>>66 fn poll_next( 67 self: Pin<&mut Self>, 68 cx: &mut Context<'_>, 69 ) -> Poll<Option<St::Item>> { 70 if self.done_taking { 71 return Poll::Ready(None); 72 } 73 74 let mut this = self.project(); 75 76 Poll::Ready(loop { 77 if let Some(fut) = this.pending_fut.as_mut().as_pin_mut() { 78 let take = ready!(fut.poll(cx)); 79 let item = this.pending_item.take(); 80 this.pending_fut.set(None); 81 if take { 82 break item; 83 } else { 84 *this.done_taking = true; 85 break None; 86 } 87 } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) { 88 this.pending_fut.set(Some((this.f)(&item))); 89 *this.pending_item = Some(item); 90 } else { 91 break None; 92 } 93 }) 94 } 95 size_hint(&self) -> (usize, Option<usize>)96 fn size_hint(&self) -> (usize, Option<usize>) { 97 if self.done_taking { 98 return (0, Some(0)); 99 } 100 101 let pending_len = if self.pending_item.is_some() { 1 } else { 0 }; 102 let (_, upper) = self.stream.size_hint(); 103 let upper = match upper { 104 Some(x) => x.checked_add(pending_len), 105 None => None, 106 }; 107 (0, upper) // can't know a lower bound, due to the predicate 108 } 109 } 110 111 impl<St, Fut, F> FusedStream for TakeWhile<St, Fut, F> 112 where St: FusedStream, 113 F: FnMut(&St::Item) -> Fut, 114 Fut: Future<Output = bool>, 115 { is_terminated(&self) -> bool116 fn is_terminated(&self) -> bool { 117 self.done_taking || self.pending_item.is_none() && self.stream.is_terminated() 118 } 119 } 120 121 // Forwarding impl of Sink from the underlying stream 122 #[cfg(feature = "sink")] 123 impl<S, Fut, F, Item> Sink<Item> for TakeWhile<S, Fut, F> 124 where S: Stream + Sink<Item>, 125 { 126 type Error = S::Error; 127 128 delegate_sink!(stream, Item); 129 } 130