Specialize io::Bytes::size_hint for more types

rust-lang · Jun 10, 2021 · 2cbd5d1 · 2cbd5d1
1 parent 0279cb1
commit 2cbd5d1
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 3 deletions.
diff --git a/library/std/src/io/buffered/bufreader.rs b/library/std/src/io/buffered/bufreader.rs
@@ -438,7 +438,13 @@ impl<R: Seek> Seek for BufReader<R> {
 }
 
 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
@@ -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 @@ impl<T: BufRead, U: BufRead> BufRead for Chain<T, U> {
 }
 
 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 @@ impl<T: BufRead> BufRead for Take<T> {
     }
 }
 
+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 @@ trait SizeHint {
 }
 
 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
@@ -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
@@ -83,6 +83,7 @@ impl fmt::Debug for Empty {
 }
 
 impl SizeHint for Empty {
+    #[inline]
     fn upper_bound(&self) -> Option<usize> {
         Some(0)
     }
@@ -147,6 +148,18 @@ impl Read for Repeat {
     }
 }
 
+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 {