From 2cbd5d1df54400b8bd718b7e0dadc4c38c6f9932 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20du=20Garreau?= <bdgdlm@outlook.com>
Date: Thu, 10 Jun 2021 19:09:04 +0200
Subject: [PATCH] Specialize `io::Bytes::size_hint` for more types

---
 library/std/src/io/buffered/bufreader.rs |  8 +++-
 library/std/src/io/mod.rs                | 58 +++++++++++++++++++++++-
 library/std/src/io/tests.rs              | 20 +++++++-
 library/std/src/io/util.rs               | 13 ++++++
 4 files changed, 96 insertions(+), 3 deletions(-)

diff --git a/library/std/src/io/buffered/bufreader.rs b/library/std/src/io/buffered/bufreader.rs
index d8021d3e99a..32d194d9616 100644
--- a/library/std/src/io/buffered/bufreader.rs
+++ b/library/std/src/io/buffered/bufreader.rs
@@ -438,7 +438,13 @@ fn stream_position(&mut self) -> io::Result<u64> {
 }
 
 impl<T> SizeHint for BufReader<T> {
+    #[inline]
     fn lower_bound(&self) -> usize {
-        self.buffer().len()
+        SizeHint::lower_bound(self.get_ref()) + self.buffer().len()
+    }
+
+    #[inline]
+    fn upper_bound(&self) -> Option<usize> {
+        SizeHint::upper_bound(self.get_ref()).and_then(|up| self.buffer().len().checked_add(up))
     }
 }
diff --git a/library/std/src/io/mod.rs b/library/std/src/io/mod.rs
index 4c154dbe01a..95aef0bd0fc 100644
--- a/library/std/src/io/mod.rs
+++ b/library/std/src/io/mod.rs
@@ -252,6 +252,7 @@
 mod tests;
 
 use crate::cmp;
+use crate::convert::TryInto;
 use crate::fmt;
 use crate::ops::{Deref, DerefMut};
 use crate::ptr;
@@ -2291,13 +2292,15 @@ fn consume(&mut self, amt: usize) {
 }
 
 impl<T, U> SizeHint for Chain<T, U> {
+    #[inline]
     fn lower_bound(&self) -> usize {
         SizeHint::lower_bound(&self.first) + SizeHint::lower_bound(&self.second)
     }
 
+    #[inline]
     fn upper_bound(&self) -> Option<usize> {
         match (SizeHint::upper_bound(&self.first), SizeHint::upper_bound(&self.second)) {
-            (Some(first), Some(second)) => Some(first + second),
+            (Some(first), Some(second)) => first.checked_add(second),
             _ => None,
         }
     }
@@ -2502,6 +2505,21 @@ fn consume(&mut self, amt: usize) {
     }
 }
 
+impl<T> SizeHint for Take<T> {
+    #[inline]
+    fn lower_bound(&self) -> usize {
+        cmp::min(SizeHint::lower_bound(&self.inner) as u64, self.limit) as usize
+    }
+
+    #[inline]
+    fn upper_bound(&self) -> Option<usize> {
+        match SizeHint::upper_bound(&self.inner) {
+            Some(upper_bound) => Some(cmp::min(upper_bound as u64, self.limit) as usize),
+            None => self.limit.try_into().ok(),
+        }
+    }
+}
+
 /// An iterator over `u8` values of a reader.
 ///
 /// This struct is generally created by calling [`bytes`] on a reader.
@@ -2546,15 +2564,53 @@ fn size_hint(&self) -> (usize, Option<usize>) {
 }
 
 impl<T> SizeHint for T {
+    #[inline]
     default fn lower_bound(&self) -> usize {
         0
     }
 
+    #[inline]
     default fn upper_bound(&self) -> Option<usize> {
         None
     }
 }
 
+impl<T> SizeHint for &mut T {
+    #[inline]
+    fn lower_bound(&self) -> usize {
+        SizeHint::lower_bound(*self)
+    }
+
+    #[inline]
+    fn upper_bound(&self) -> Option<usize> {
+        SizeHint::upper_bound(*self)
+    }
+}
+
+impl<T> SizeHint for Box<T> {
+    #[inline]
+    fn lower_bound(&self) -> usize {
+        SizeHint::lower_bound(&**self)
+    }
+
+    #[inline]
+    fn upper_bound(&self) -> Option<usize> {
+        SizeHint::upper_bound(&**self)
+    }
+}
+
+impl SizeHint for &[u8] {
+    #[inline]
+    fn lower_bound(&self) -> usize {
+        self.len()
+    }
+
+    #[inline]
+    fn upper_bound(&self) -> Option<usize> {
+        Some(self.len())
+    }
+}
+
 /// An iterator over the contents of an instance of `BufRead` split on a
 /// particular byte.
 ///
diff --git a/library/std/src/io/tests.rs b/library/std/src/io/tests.rs
index 2b14e161503..a483847fb26 100644
--- a/library/std/src/io/tests.rs
+++ b/library/std/src/io/tests.rs
@@ -224,6 +224,24 @@ fn empty_size_hint() {
     assert_eq!(size_hint, (0, Some(0)));
 }
 
+#[test]
+fn slice_size_hint() {
+    let size_hint = (&[1, 2, 3]).bytes().size_hint();
+    assert_eq!(size_hint, (3, Some(3)));
+}
+
+#[test]
+fn take_size_hint() {
+    let size_hint = (&[1, 2, 3]).take(2).bytes().size_hint();
+    assert_eq!(size_hint, (2, Some(2)));
+
+    let size_hint = (&[1, 2, 3]).take(4).bytes().size_hint();
+    assert_eq!(size_hint, (3, Some(3)));
+
+    let size_hint = io::repeat(0).take(3).bytes().size_hint();
+    assert_eq!(size_hint, (3, Some(3)));
+}
+
 #[test]
 fn chain_empty_size_hint() {
     let chain = io::empty().chain(io::empty());
@@ -242,7 +260,7 @@ fn chain_size_hint() {
 
     let chain = buf_reader_1.chain(buf_reader_2);
     let size_hint = chain.bytes().size_hint();
-    assert_eq!(size_hint, (testdata.len(), None));
+    assert_eq!(size_hint, (testdata.len(), Some(testdata.len())));
 }
 
 #[test]
diff --git a/library/std/src/io/util.rs b/library/std/src/io/util.rs
index 73f2f3eb3f5..f3bff391fb3 100644
--- a/library/std/src/io/util.rs
+++ b/library/std/src/io/util.rs
@@ -83,6 +83,7 @@ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 }
 
 impl SizeHint for Empty {
+    #[inline]
     fn upper_bound(&self) -> Option<usize> {
         Some(0)
     }
@@ -147,6 +148,18 @@ unsafe fn initializer(&self) -> Initializer {
     }
 }
 
+impl SizeHint for Repeat {
+    #[inline]
+    fn lower_bound(&self) -> usize {
+        usize::MAX
+    }
+
+    #[inline]
+    fn upper_bound(&self) -> Option<usize> {
+        None
+    }
+}
+
 #[stable(feature = "std_debug", since = "1.16.0")]
 impl fmt::Debug for Repeat {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {