Use try_reserve in Vec's io::Write

Author: kornelski

library/std/src/io/impls.rs

@@ -377,18 +377,28 @@ impl Write for &mut [u8] {
 
 /// Write is implemented for `Vec<u8>` by appending to the vector.
 /// The vector will grow as needed.
+///
+/// # Panics
+///
+/// In case of allocation error or capacity overflow, write operations will panic.
+/// The panicking behavior is not guaranteed. In the future, it may become
+/// a regular `io::Error`.
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<A: Allocator> Write for Vec<u8, A> {
     #[inline]
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        fallible_reserve(self, buf.len());
         self.extend_from_slice(buf);
         Ok(buf.len())
     }
 
     #[inline]
     fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
-        let len = bufs.iter().map(|b| b.len()).sum();
-        self.reserve(len);
+        let len = bufs
+            .iter()
+            .try_fold(0usize, |len, b| len.checked_add(b.len()))
+            .expect("capacity overflow");
+        fallible_reserve(self, len);
         for buf in bufs {
             self.extend_from_slice(buf);
         }
@@ -402,6 +412,7 @@ impl<A: Allocator> Write for Vec<u8, A> {
 
     #[inline]
     fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+        fallible_reserve(self, buf.len());
         self.extend_from_slice(buf);
         Ok(())
     }
@@ -412,6 +423,20 @@ impl<A: Allocator> Write for Vec<u8, A> {
     }
 }
 
+#[inline]
+fn fallible_reserve<T, A: Allocator>(vec: &mut Vec<T, A>, len: usize) {
+    if len > vec.capacity().wrapping_sub(vec.len()) {
+        do_reserve_and_handle(vec, len);
+    }
+}
+
+#[cold]
+fn do_reserve_and_handle<T, A: Allocator>(vec: &mut Vec<T, A>, len: usize) {
+    if let Err(_) = vec.try_reserve(len) {
+        panic!("out of memory");
+    }
+}
+
 /// Read is implemented for `VecDeque<u8>` by consuming bytes from the front of the `VecDeque`.
 #[stable(feature = "vecdeque_read_write", since = "1.63.0")]
 impl<A: Allocator> Read for VecDeque<u8, A> {