diff --git a/tokio/src/stream/chain.rs b/tokio/src/stream/chain.rs index 5f0324a4b..6124c91e4 100644 --- a/tokio/src/stream/chain.rs +++ b/tokio/src/stream/chain.rs @@ -44,14 +44,6 @@ where } fn size_hint(&self) -> (usize, Option) { - let (a_lower, a_upper) = self.a.size_hint(); - let (b_lower, b_upper) = self.b.size_hint(); - - let upper = match (a_upper, b_upper) { - (Some(a_upper), Some(b_upper)) => Some(a_upper + b_upper), - _ => None, - }; - - (a_lower + b_lower, upper) + super::merge_size_hints(self.a.size_hint(), self.b.size_hint()) } } diff --git a/tokio/src/stream/merge.rs b/tokio/src/stream/merge.rs index 4850cd40c..50ba518ce 100644 --- a/tokio/src/stream/merge.rs +++ b/tokio/src/stream/merge.rs @@ -52,15 +52,7 @@ where } fn size_hint(&self) -> (usize, Option) { - let (a_lower, a_upper) = self.a.size_hint(); - let (b_lower, b_upper) = self.b.size_hint(); - - let upper = match (a_upper, b_upper) { - (Some(a_upper), Some(b_upper)) => Some(a_upper + b_upper), - _ => None, - }; - - (a_lower + b_lower, upper) + super::merge_size_hints(self.a.size_hint(), self.b.size_hint()) } } diff --git a/tokio/src/stream/mod.rs b/tokio/src/stream/mod.rs index 307ead5fb..a59bdfcd4 100644 --- a/tokio/src/stream/mod.rs +++ b/tokio/src/stream/mod.rs @@ -817,3 +817,16 @@ pub trait StreamExt: Stream { } impl StreamExt for St where St: Stream {} + +/// Merge the size hints from two streams. +fn merge_size_hints( + (left_low, left_high): (usize, Option), + (right_low, right_hign): (usize, Option), +) -> (usize, Option) { + let low = left_low.saturating_add(right_low); + let high = match (left_high, right_hign) { + (Some(h1), Some(h2)) => h1.checked_add(h2), + _ => None, + }; + (low, high) +} diff --git a/tokio/tests/stream_chain.rs b/tokio/tests/stream_chain.rs index 0e14618b4..98461a8cc 100644 --- a/tokio/tests/stream_chain.rs +++ b/tokio/tests/stream_chain.rs @@ -69,3 +69,27 @@ async fn pending_first() { assert_eq!(stream.size_hint(), (0, None)); assert_eq!(None, assert_ready!(stream.poll_next())); } + +#[test] +fn size_overflow() { + struct Monster; + + impl tokio::stream::Stream for Monster { + type Item = (); + fn poll_next( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + panic!() + } + + fn size_hint(&self) -> (usize, Option) { + (usize::max_value(), Some(usize::max_value())) + } + } + + let m1 = Monster; + let m2 = Monster; + let m = m1.chain(m2); + assert_eq!(m.size_hint(), (usize::max_value(), None)); +} diff --git a/tokio/tests/stream_merge.rs b/tokio/tests/stream_merge.rs index f0168d72e..45ecdcb66 100644 --- a/tokio/tests/stream_merge.rs +++ b/tokio/tests/stream_merge.rs @@ -52,3 +52,27 @@ async fn merge_async_streams() { assert!(rx.is_woken()); assert_eq!(None, assert_ready!(rx.poll_next())); } + +#[test] +fn size_overflow() { + struct Monster; + + impl tokio::stream::Stream for Monster { + type Item = (); + fn poll_next( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + panic!() + } + + fn size_hint(&self) -> (usize, Option) { + (usize::max_value(), Some(usize::max_value())) + } + } + + let m1 = Monster; + let m2 = Monster; + let m = m1.merge(m2); + assert_eq!(m.size_hint(), (usize::max_value(), None)); +}