From 4e27ed3af19e604d7b65e130145fcecdc69fba7a Mon Sep 17 00:00:00 2001
From: Ben Kimock <kimockb@gmail.com>
Date: Sat, 19 Dec 2020 15:02:01 -0500
Subject: [PATCH] Add benchmark and fast path for BufReader::read_exact

---
 library/std/src/io/buffered/bufreader.rs | 14 +++++++++
 library/std/src/io/buffered/tests.rs     | 12 ++++++++
 library/std/src/io/mod.rs                | 38 +++++++++++++-----------
 3 files changed, 47 insertions(+), 17 deletions(-)

diff --git a/library/std/src/io/buffered/bufreader.rs b/library/std/src/io/buffered/bufreader.rs
index 16c18d6e14645..8bae3da1273eb 100644
--- a/library/std/src/io/buffered/bufreader.rs
+++ b/library/std/src/io/buffered/bufreader.rs
@@ -271,6 +271,20 @@ impl<R: Read> Read for BufReader<R> {
         Ok(nread)
     }
 
+    // Small read_exacts from a BufReader are extremely common when used with a deserializer.
+    // The default implementation calls read in a loop, which results in surprisingly poor code
+    // generation for the common path where the buffer has enough bytes to fill the passed-in
+    // buffer.
+    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
+        if self.buffer().len() >= buf.len() {
+            buf.copy_from_slice(&self.buffer()[..buf.len()]);
+            self.consume(buf.len());
+            return Ok(());
+        }
+
+        crate::io::default_read_exact(self, buf)
+    }
+
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         let total_len = bufs.iter().map(|b| b.len()).sum::<usize>();
         if self.pos == self.cap && total_len >= self.buf.len() {
diff --git a/library/std/src/io/buffered/tests.rs b/library/std/src/io/buffered/tests.rs
index 66a64f667baa4..f6c2b499567ab 100644
--- a/library/std/src/io/buffered/tests.rs
+++ b/library/std/src/io/buffered/tests.rs
@@ -443,6 +443,18 @@ fn bench_buffered_reader(b: &mut test::Bencher) {
     b.iter(|| BufReader::new(io::empty()));
 }
 
+#[bench]
+fn bench_buffered_reader_small_reads(b: &mut test::Bencher) {
+    let data = (0..u8::MAX).cycle().take(1024 * 4).collect::<Vec<_>>();
+    b.iter(|| {
+        let mut reader = BufReader::new(&data[..]);
+        let mut buf = [0u8; 4];
+        for _ in 0..1024 {
+            reader.read_exact(&mut buf).unwrap();
+        }
+    });
+}
+
 #[bench]
 fn bench_buffered_writer(b: &mut test::Bencher) {
     b.iter(|| BufWriter::new(io::sink()));
diff --git a/library/std/src/io/mod.rs b/library/std/src/io/mod.rs
index 3f5b7c0b29be6..c87a56586c65e 100644
--- a/library/std/src/io/mod.rs
+++ b/library/std/src/io/mod.rs
@@ -416,6 +416,25 @@ where
     write(buf)
 }
 
+pub(crate) fn default_read_exact<R: Read + ?Sized>(this: &mut R, mut buf: &mut [u8]) -> Result<()> {
+    while !buf.is_empty() {
+        match this.read(buf) {
+            Ok(0) => break,
+            Ok(n) => {
+                let tmp = buf;
+                buf = &mut tmp[n..];
+            }
+            Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
+            Err(e) => return Err(e),
+        }
+    }
+    if !buf.is_empty() {
+        Err(Error::new(ErrorKind::UnexpectedEof, "failed to fill whole buffer"))
+    } else {
+        Ok(())
+    }
+}
+
 /// The `Read` trait allows for reading bytes from a source.
 ///
 /// Implementors of the `Read` trait are called 'readers'.
@@ -766,23 +785,8 @@ pub trait Read {
     /// }
     /// ```
     #[stable(feature = "read_exact", since = "1.6.0")]
-    fn read_exact(&mut self, mut buf: &mut [u8]) -> Result<()> {
-        while !buf.is_empty() {
-            match self.read(buf) {
-                Ok(0) => break,
-                Ok(n) => {
-                    let tmp = buf;
-                    buf = &mut tmp[n..];
-                }
-                Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
-                Err(e) => return Err(e),
-            }
-        }
-        if !buf.is_empty() {
-            Err(Error::new(ErrorKind::UnexpectedEof, "failed to fill whole buffer"))
-        } else {
-            Ok(())
-        }
+    fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
+        default_read_exact(self, buf)
     }
 
     /// Creates a "by reference" adaptor for this instance of `Read`.