From 42f77208ab95cfa1468b72cc8d5bf8ffdc252830 Mon Sep 17 00:00:00 2001
From: Nicholas Nethercote <n.nethercote@gmail.com>
Date: Fri, 13 Oct 2023 08:26:49 +1100
Subject: [PATCH] Specialize `Bytes<R>::next` when `R` is a `BufReader`.

This reduces the runtime for a simple program using `Bytes::next` to
iterate through a file from 220ms to 70ms on my Linux box.
---
 library/std/src/io/buffered/bufreader.rs | 24 +++++++++++---
 library/std/src/io/mod.rs                | 41 ++++++++++++++++++++----
 2 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/library/std/src/io/buffered/bufreader.rs b/library/std/src/io/buffered/bufreader.rs
index 7097dfef88d4e..17255efc12b17 100644
--- a/library/std/src/io/buffered/bufreader.rs
+++ b/library/std/src/io/buffered/bufreader.rs
@@ -2,7 +2,8 @@ mod buffer;
 
 use crate::fmt;
 use crate::io::{
-    self, BorrowedCursor, BufRead, IoSliceMut, Read, Seek, SeekFrom, SizeHint, DEFAULT_BUF_SIZE,
+    self, BorrowedCursor, BufRead, IoSliceMut, Read, Seek, SeekFrom, SizeHint, SpecReadByte,
+    DEFAULT_BUF_SIZE,
 };
 use buffer::Buffer;
 
@@ -259,6 +260,21 @@ impl<R: ?Sized + Seek> BufReader<R> {
     }
 }
 
+impl<R> SpecReadByte for BufReader<R>
+where
+    Self: Read,
+{
+    #[inline]
+    fn spec_read_byte(&mut self) -> Option<io::Result<u8>> {
+        let mut byte = 0;
+        if self.buf.consume_with(1, |claimed| byte = claimed[0]) {
+            return Some(Ok(byte));
+        }
+
+        self.slow_read_byte()
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<R: ?Sized + Read> Read for BufReader<R> {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
@@ -269,10 +285,8 @@ impl<R: ?Sized + Read> Read for BufReader<R> {
             self.discard_buffer();
             return self.inner.read(buf);
         }
-        let nread = {
-            let mut rem = self.fill_buf()?;
-            rem.read(buf)?
-        };
+        let mut rem = self.fill_buf()?;
+        let nread = rem.read(buf)?;
         self.consume(nread);
         Ok(nread)
     }
diff --git a/library/std/src/io/mod.rs b/library/std/src/io/mod.rs
index c0a729481121a..07b464e094931 100644
--- a/library/std/src/io/mod.rs
+++ b/library/std/src/io/mod.rs
@@ -2777,11 +2777,43 @@ pub struct Bytes<R> {
 impl<R: Read> Iterator for Bytes<R> {
     type Item = Result<u8>;
 
-    #[inline]
+    // Not `#[inline]`. This function gets inlined even without it, but having
+    // the inline annotation can result in worse code generation. See #116785.
     fn next(&mut self) -> Option<Result<u8>> {
+        SpecReadByte::spec_read_byte(&mut self.inner)
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        SizeHint::size_hint(&self.inner)
+    }
+}
+
+// Trait for specialization of `Bytes::next`.
+trait SpecReadByte {
+    // Specializable method for reading a single byte.
+    fn spec_read_byte(&mut self) -> Option<Result<u8>>;
+
+    // Non-specializable method for reading a single byte. This is used by the
+    // default `spec_read_byte` and can also be called on a slow path of a
+    // specialized `spec_read_byte` method.
+    fn slow_read_byte(&mut self) -> Option<Result<u8>>;
+}
+
+impl<R> SpecReadByte for R
+where
+    Self: Read,
+{
+    #[inline]
+    default fn spec_read_byte(&mut self) -> Option<Result<u8>> {
+        self.slow_read_byte()
+    }
+
+    #[inline(never)]
+    fn slow_read_byte(&mut self) -> Option<Result<u8>> {
         let mut byte = 0;
         loop {
-            return match self.inner.read(slice::from_mut(&mut byte)) {
+            return match self.read(slice::from_mut(&mut byte)) {
                 Ok(0) => None,
                 Ok(..) => Some(Ok(byte)),
                 Err(ref e) if e.is_interrupted() => continue,
@@ -2789,11 +2821,6 @@ impl<R: Read> Iterator for Bytes<R> {
             };
         }
     }
-
-    #[inline]
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        SizeHint::size_hint(&self.inner)
-    }
 }
 
 trait SizeHint {