Rollup merge of #108326 - tmiasko:read-buf, r=thomcc

Implement read_buf for a few more types

Implement read_buf for TcpStream, Stdin, StdinLock, ChildStdout,
ChildStderr (and internally for AnonPipe, Handle, Socket), so
that it skips buffer initialization.

The other provided methods like read_to_string and read_to_end are
implemented in terms of read_buf and so benefit from the optimization
as well.

This commit also implements read_vectored and is_read_vectored where
applicable.

14 files changed, 136 insertions, 24 deletions

diff --git a/library/std/src/sys/sgx/fd.rs b/library/std/src/sys/sgx/fd.rs
index e5dc5b5adaa..0c02a107691 100644
--- a/library/std/src/sys/sgx/fd.rs
+++ b/library/std/src/sys/sgx/fd.rs
@@ -1,7 +1,7 @@
 use fortanix_sgx_abi::Fd;
 
 use super::abi::usercalls;
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::mem;
 use crate::sys::{AsInner, FromInner, IntoInner};
 
@@ -30,6 +30,10 @@ impl FileDesc {
         usercalls::read(self.fd, &mut [IoSliceMut::new(buf)])
     }
 
+    pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        crate::io::default_read_buf(|b| self.read(b), buf)
+    }
+
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         usercalls::read(self.fd, bufs)
     }
diff --git a/library/std/src/sys/sgx/net.rs b/library/std/src/sys/sgx/net.rs
index 4c4cd7d1d1d..923be5eb944 100644
--- a/library/std/src/sys/sgx/net.rs
+++ b/library/std/src/sys/sgx/net.rs
@@ -1,6 +1,6 @@
 use crate::error;
 use crate::fmt;
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::net::{Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr, ToSocketAddrs};
 use crate::sync::Arc;
 use crate::sys::fd::FileDesc;
@@ -144,6 +144,10 @@ impl TcpStream {
         self.inner.inner.read(buf)
     }
 
+    pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.inner.inner.read_buf(buf)
+    }
+
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.inner.inner.read_vectored(bufs)
     }
diff --git a/library/std/src/sys/unix/fd.rs b/library/std/src/sys/unix/fd.rs
index 9874af4d3e2..ce5c048f252 100644
--- a/library/std/src/sys/unix/fd.rs
+++ b/library/std/src/sys/unix/fd.rs
@@ -469,6 +469,15 @@ impl<'a> Read for &'a FileDesc {
     fn read_buf(&mut self, cursor: BorrowedCursor<'_>) -> io::Result<()> {
         (**self).read_buf(cursor)
     }
+
+    fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+        (**self).read_vectored(bufs)
+    }
+
+    #[inline]
+    fn is_read_vectored(&self) -> bool {
+        (**self).is_read_vectored()
+    }
 }
 
 impl AsInner<OwnedFd> for FileDesc {
diff --git a/library/std/src/sys/unix/net.rs b/library/std/src/sys/unix/net.rs
index 8e05b618daa..f84240d268f 100644
--- a/library/std/src/sys/unix/net.rs
+++ b/library/std/src/sys/unix/net.rs
@@ -1,6 +1,6 @@
 use crate::cmp;
 use crate::ffi::CStr;
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedBuf, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::mem;
 use crate::net::{Shutdown, SocketAddr};
 use crate::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd};
@@ -242,19 +242,35 @@ impl Socket {
         self.0.duplicate().map(Socket)
     }
 
-    fn recv_with_flags(&self, buf: &mut [u8], flags: c_int) -> io::Result<usize> {
+    fn recv_with_flags(&self, mut buf: BorrowedCursor<'_>, flags: c_int) -> io::Result<()> {
         let ret = cvt(unsafe {
-            libc::recv(self.as_raw_fd(), buf.as_mut_ptr() as *mut c_void, buf.len(), flags)
+            libc::recv(
+                self.as_raw_fd(),
+                buf.as_mut().as_mut_ptr() as *mut c_void,
+                buf.capacity(),
+                flags,
+            )
         })?;
-        Ok(ret as usize)
+        unsafe {
+            buf.advance(ret as usize);
+        }
+        Ok(())
     }
 
     pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
-        self.recv_with_flags(buf, 0)
+        let mut buf = BorrowedBuf::from(buf);
+        self.recv_with_flags(buf.unfilled(), 0)?;
+        Ok(buf.len())
     }
 
     pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
-        self.recv_with_flags(buf, MSG_PEEK)
+        let mut buf = BorrowedBuf::from(buf);
+        self.recv_with_flags(buf.unfilled(), MSG_PEEK)?;
+        Ok(buf.len())
+    }
+
+    pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.recv_with_flags(buf, 0)
     }
 
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
diff --git a/library/std/src/sys/unix/pipe.rs b/library/std/src/sys/unix/pipe.rs
index a744d0ab640..dc17c9fac46 100644
--- a/library/std/src/sys/unix/pipe.rs
+++ b/library/std/src/sys/unix/pipe.rs
@@ -1,4 +1,4 @@
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::mem;
 use crate::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd};
 use crate::sys::fd::FileDesc;
@@ -49,6 +49,10 @@ impl AnonPipe {
         self.0.read(buf)
     }
 
+    pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.0.read_buf(buf)
+    }
+
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.0.read_vectored(bufs)
     }
diff --git a/library/std/src/sys/unix/stdio.rs b/library/std/src/sys/unix/stdio.rs
index b3626c564e8..a26f20795a1 100644
--- a/library/std/src/sys/unix/stdio.rs
+++ b/library/std/src/sys/unix/stdio.rs
@@ -1,4 +1,4 @@
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::mem::ManuallyDrop;
 use crate::os::unix::io::FromRawFd;
 use crate::sys::fd::FileDesc;
@@ -18,6 +18,10 @@ impl io::Read for Stdin {
         unsafe { ManuallyDrop::new(FileDesc::from_raw_fd(libc::STDIN_FILENO)).read(buf) }
     }
 
+    fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        unsafe { ManuallyDrop::new(FileDesc::from_raw_fd(libc::STDIN_FILENO)).read_buf(buf) }
+    }
+
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         unsafe { ManuallyDrop::new(FileDesc::from_raw_fd(libc::STDIN_FILENO)).read_vectored(bufs) }
     }
diff --git a/library/std/src/sys/unsupported/net.rs b/library/std/src/sys/unsupported/net.rs
index a5204a08453..bbc52703f96 100644
--- a/library/std/src/sys/unsupported/net.rs
+++ b/library/std/src/sys/unsupported/net.rs
@@ -1,5 +1,5 @@
 use crate::fmt;
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::net::{Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr};
 use crate::sys::unsupported;
 use crate::time::Duration;
@@ -39,6 +39,10 @@ impl TcpStream {
         self.0
     }
 
+    pub fn read_buf(&self, _buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.0
+    }
+
     pub fn read_vectored(&self, _: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.0
     }
diff --git a/library/std/src/sys/unsupported/pipe.rs b/library/std/src/sys/unsupported/pipe.rs
index 0bba673b458..d7d8f297ae5 100644
--- a/library/std/src/sys/unsupported/pipe.rs
+++ b/library/std/src/sys/unsupported/pipe.rs
@@ -1,4 +1,4 @@
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 
 pub struct AnonPipe(!);
 
@@ -7,6 +7,10 @@ impl AnonPipe {
         self.0
     }
 
+    pub fn read_buf(&self, _buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.0
+    }
+
     pub fn read_vectored(&self, _bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.0
     }
diff --git a/library/std/src/sys/wasi/fd.rs b/library/std/src/sys/wasi/fd.rs
index 0b9c8e61db8..191db4b60f7 100644
--- a/library/std/src/sys/wasi/fd.rs
+++ b/library/std/src/sys/wasi/fd.rs
@@ -2,7 +2,7 @@
 #![allow(dead_code)]
 
 use super::err2io;
-use crate::io::{self, IoSlice, IoSliceMut, SeekFrom};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut, SeekFrom};
 use crate::mem;
 use crate::net::Shutdown;
 use crate::os::wasi::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, OwnedFd, RawFd};
@@ -46,6 +46,22 @@ impl WasiFd {
         unsafe { wasi::fd_read(self.as_raw_fd() as wasi::Fd, iovec(bufs)).map_err(err2io) }
     }
 
+    pub fn read_buf(&self, mut buf: BorrowedCursor<'_>) -> io::Result<()> {
+        unsafe {
+            let bufs = [wasi::Iovec {
+                buf: buf.as_mut().as_mut_ptr() as *mut u8,
+                buf_len: buf.capacity(),
+            }];
+            match wasi::fd_read(self.as_raw_fd() as wasi::Fd, &bufs) {
+                Ok(n) => {
+                    buf.advance(n);
+                    Ok(())
+                }
+                Err(e) => Err(err2io(e)),
+            }
+        }
+    }
+
     pub fn write(&self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
         unsafe { wasi::fd_write(self.as_raw_fd() as wasi::Fd, ciovec(bufs)).map_err(err2io) }
     }
diff --git a/library/std/src/sys/wasi/fs.rs b/library/std/src/sys/wasi/fs.rs
index d4866bbc32b..3a205267e34 100644
--- a/library/std/src/sys/wasi/fs.rs
+++ b/library/std/src/sys/wasi/fs.rs
@@ -441,7 +441,7 @@ impl File {
     }
 
     pub fn read_buf(&self, cursor: BorrowedCursor<'_>) -> io::Result<()> {
-        crate::io::default_read_buf(|buf| self.read(buf), cursor)
+        self.fd.read_buf(cursor)
     }
 
     pub fn write(&self, buf: &[u8]) -> io::Result<usize> {
diff --git a/library/std/src/sys/wasi/net.rs b/library/std/src/sys/wasi/net.rs
index cf4ebba1a39..59d94a3686d 100644
--- a/library/std/src/sys/wasi/net.rs
+++ b/library/std/src/sys/wasi/net.rs
@@ -3,7 +3,7 @@
 use super::err2io;
 use super::fd::WasiFd;
 use crate::fmt;
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::net::{Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr};
 use crate::os::wasi::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd};
 use crate::sys::unsupported;
@@ -91,6 +91,10 @@ impl TcpStream {
         self.read_vectored(&mut [IoSliceMut::new(buf)])
     }
 
+    pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.socket().as_inner().read_buf(buf)
+    }
+
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.socket().as_inner().read(bufs)
     }
diff --git a/library/std/src/sys/windows/handle.rs b/library/std/src/sys/windows/handle.rs
index ae33d48c612..b290f4070e8 100644
--- a/library/std/src/sys/windows/handle.rs
+++ b/library/std/src/sys/windows/handle.rs
@@ -327,7 +327,16 @@ impl<'a> Read for &'a Handle {
         (**self).read(buf)
     }
 
+    fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        (**self).read_buf(buf)
+    }
+
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         (**self).read_vectored(bufs)
     }
+
+    #[inline]
+    fn is_read_vectored(&self) -> bool {
+        (**self).is_read_vectored()
+    }
 }
diff --git a/library/std/src/sys/windows/net.rs b/library/std/src/sys/windows/net.rs
index e0701a498fa..ee1f5482b47 100644
--- a/library/std/src/sys/windows/net.rs
+++ b/library/std/src/sys/windows/net.rs
@@ -1,7 +1,7 @@
 #![unstable(issue = "none", feature = "windows_net")]
 
 use crate::cmp;
-use crate::io::{self, IoSlice, IoSliceMut, Read};
+use crate::io::{self, BorrowedBuf, BorrowedCursor, IoSlice, IoSliceMut, Read};
 use crate::mem;
 use crate::net::{Shutdown, SocketAddr};
 use crate::os::windows::io::{
@@ -214,28 +214,38 @@ impl Socket {
         Ok(Self(self.0.try_clone()?))
     }
 
-    fn recv_with_flags(&self, buf: &mut [u8], flags: c_int) -> io::Result<usize> {
+    fn recv_with_flags(&self, mut buf: BorrowedCursor<'_>, flags: c_int) -> io::Result<()> {
         // On unix when a socket is shut down all further reads return 0, so we
         // do the same on windows to map a shut down socket to returning EOF.
-        let length = cmp::min(buf.len(), i32::MAX as usize) as i32;
-        let result =
-            unsafe { c::recv(self.as_raw_socket(), buf.as_mut_ptr() as *mut _, length, flags) };
+        let length = cmp::min(buf.capacity(), i32::MAX as usize) as i32;
+        let result = unsafe {
+            c::recv(self.as_raw_socket(), buf.as_mut().as_mut_ptr() as *mut _, length, flags)
+        };
 
         match result {
             c::SOCKET_ERROR => {
                 let error = unsafe { c::WSAGetLastError() };
 
                 if error == c::WSAESHUTDOWN {
-                    Ok(0)
+                    Ok(())
                 } else {
                     Err(io::Error::from_raw_os_error(error))
                 }
             }
-            _ => Ok(result as usize),
+            _ => {
+                unsafe { buf.advance(result as usize) };
+                Ok(())
+            }
         }
     }
 
     pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
+        let mut buf = BorrowedBuf::from(buf);
+        self.recv_with_flags(buf.unfilled(), 0)?;
+        Ok(buf.len())
+    }
+
+    pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
         self.recv_with_flags(buf, 0)
     }
 
@@ -277,7 +287,9 @@ impl Socket {
     }
 
     pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
-        self.recv_with_flags(buf, c::MSG_PEEK)
+        let mut buf = BorrowedBuf::from(buf);
+        self.recv_with_flags(buf.unfilled(), c::MSG_PEEK)?;
+        Ok(buf.len())
     }
 
     fn recv_from_with_flags(
diff --git a/library/std/src/sys/windows/pipe.rs b/library/std/src/sys/windows/pipe.rs
index 7b25edaa556..0780b29584d 100644
--- a/library/std/src/sys/windows/pipe.rs
+++ b/library/std/src/sys/windows/pipe.rs
@@ -1,7 +1,7 @@
 use crate::os::windows::prelude::*;
 
 use crate::ffi::OsStr;
-use crate::io::{self, IoSlice, IoSliceMut, Read};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut, Read};
 use crate::mem;
 use crate::path::Path;
 use crate::ptr;
@@ -252,6 +252,28 @@ impl AnonPipe {
         }
     }
 
+    pub fn read_buf(&self, mut buf: BorrowedCursor<'_>) -> io::Result<()> {
+        let result = unsafe {
+            let len = crate::cmp::min(buf.capacity(), c::DWORD::MAX as usize) as c::DWORD;
+            self.alertable_io_internal(c::ReadFileEx, buf.as_mut().as_mut_ptr() as _, len)
+        };
+
+        match result {
+            // The special treatment of BrokenPipe is to deal with Windows
+            // pipe semantics, which yields this error when *reading* from
+            // a pipe after the other end has closed; we interpret that as
+            // EOF on the pipe.
+            Err(ref e) if e.kind() == io::ErrorKind::BrokenPipe => Ok(()),
+            Err(e) => Err(e),
+            Ok(n) => {
+                unsafe {
+                    buf.advance(n);
+                }
+                Ok(())
+            }
+        }
+    }
+
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.inner.read_vectored(bufs)
     }