Auto merge of #137762 - thaliaarchi:io-optional-methods/write-fmt, r=<try>

Reserve before `write_fmt`

`fmt::Arguments::estimated_capacity()` is currently only used by `fmt::format()` to reserve the initial string capacity. Also use it for `impl Write` for `Vec<u8>`, `VecDeque<u8>`, `Cursor<&mut Vec<u8>>`, and `Cursor<Vec<u8>>`.

This may be worth checking perf.

Author: bors

library/alloc/src/string.rs

@@ -3186,6 +3186,16 @@ impl fmt::Write for String {
         self.push(c);
         Ok(())
     }
+
+    #[inline]
+    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result {
+        if let Some(s) = args.as_statically_known_str() {
+            self.write_str(s)
+        } else {
+            self.reserve(args.estimated_capacity());
+            fmt::write(self, args)
+        }
+    }
 }
 
 /// An iterator over the [`char`]s of a string.

library/core/src/fmt/mod.rs

@@ -710,9 +710,10 @@ impl<'a> Arguments<'a> {
     }
 
     /// Same as [`Arguments::as_str`], but will only return `Some(s)` if it can be determined at compile time.
+    #[unstable(feature = "fmt_internals", reason = "internal to standard library", issue = "none")]
     #[must_use]
     #[inline]
-    fn as_statically_known_str(&self) -> Option<&'static str> {
+    pub fn as_statically_known_str(&self) -> Option<&'static str> {
         let s = self.as_str();
         if core::intrinsics::is_val_statically_known(s.is_some()) { s } else { None }
     }

library/std/src/ffi/os_str.rs

@@ -753,10 +753,21 @@ impl Hash for OsString {
 
 #[stable(feature = "os_string_fmt_write", since = "1.64.0")]
 impl fmt::Write for OsString {
+    #[inline]
     fn write_str(&mut self, s: &str) -> fmt::Result {
         self.push(s);
         Ok(())
     }
+
+    #[inline]
+    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result {
+        if let Some(s) = args.as_statically_known_str() {
+            self.write_str(s)
+        } else {
+            self.reserve(args.estimated_capacity());
+            fmt::write(self, args)
+        }
+    }
 }
 
 impl OsStr {

library/std/src/io/cursor.rs

@@ -2,9 +2,9 @@
 mod tests;
 
 use crate::alloc::Allocator;
-use crate::cmp;
 use crate::io::prelude::*;
 use crate::io::{self, BorrowedCursor, ErrorKind, IoSlice, IoSliceMut, SeekFrom};
+use crate::{cmp, fmt};
 
 /// A `Cursor` wraps an in-memory buffer and provides it with a
 /// [`Seek`] implementation.
@@ -603,6 +603,11 @@ where
         true
     }
 
+    fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> io::Result<()> {
+        self.inner.reserve(fmt.estimated_capacity());
+        io::default_write_fmt(self, fmt)
+    }
+
     #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
@@ -627,6 +632,11 @@ where
         true
     }
 
+    fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> io::Result<()> {
+        self.inner.reserve(fmt.estimated_capacity());
+        io::default_write_fmt(self, fmt)
+    }
+
     #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())

library/std/src/io/impls.rs

@@ -495,6 +495,12 @@ impl<A: Allocator> Write for Vec<u8, A> {
         Ok(())
     }
 
+    #[inline]
+    fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> io::Result<()> {
+        self.reserve(fmt.estimated_capacity());
+        io::default_write_fmt(self, fmt)
+    }
+
     #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
@@ -638,6 +644,12 @@ impl<A: Allocator> Write for VecDeque<u8, A> {
         Ok(())
     }
 
+    #[inline]
+    fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> io::Result<()> {
+        self.reserve(fmt.estimated_capacity());
+        io::default_write_fmt(self, fmt)
+    }
+
     #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())

library/std/src/io/mod.rs

@@ -612,6 +612,47 @@ pub(crate) fn default_read_buf_exact<R: Read + ?Sized>(
     Ok(())
 }
 
+pub(crate) fn default_write_fmt<W: Write + ?Sized>(
+    this: &mut W,
+    fmt: fmt::Arguments<'_>,
+) -> Result<()> {
+    // Create a shim which translates a `Write` to a `fmt::Write` and saves off
+    // I/O errors, instead of discarding them.
+    struct Adapter<'a, T: ?Sized + 'a> {
+        inner: &'a mut T,
+        error: Result<()>,
+    }
+
+    impl<T: Write + ?Sized> fmt::Write for Adapter<'_, T> {
+        fn write_str(&mut self, s: &str) -> fmt::Result {
+            match self.inner.write_all(s.as_bytes()) {
+                Ok(()) => Ok(()),
+                Err(e) => {
+                    self.error = Err(e);
+                    Err(fmt::Error)
+                }
+            }
+        }
+    }
+
+    let mut output = Adapter { inner: this, error: Ok(()) };
+    match fmt::write(&mut output, fmt) {
+        Ok(()) => Ok(()),
+        Err(..) => {
+            // Check whether the error came from the underlying `Write`.
+            if output.error.is_err() {
+                output.error
+            } else {
+                // This shouldn't happen: the underlying stream did not error,
+                // but somehow the formatter still errored?
+                panic!(
+                    "a formatting trait implementation returned an error when the underlying stream did not"
+                );
+            }
+        }
+    }
+}
+
 /// The `Read` trait allows for reading bytes from a source.
 ///
 /// Implementors of the `Read` trait are called 'readers'.
@@ -1867,41 +1908,7 @@ pub trait Write {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> Result<()> {
-        // Create a shim which translates a Write to a fmt::Write and saves
-        // off I/O errors. instead of discarding them
-        struct Adapter<'a, T: ?Sized + 'a> {
-            inner: &'a mut T,
-            error: Result<()>,
-        }
-
-        impl<T: Write + ?Sized> fmt::Write for Adapter<'_, T> {
-            fn write_str(&mut self, s: &str) -> fmt::Result {
-                match self.inner.write_all(s.as_bytes()) {
-                    Ok(()) => Ok(()),
-                    Err(e) => {
-                        self.error = Err(e);
-                        Err(fmt::Error)
-                    }
-                }
-            }
-        }
-
-        let mut output = Adapter { inner: self, error: Ok(()) };
-        match fmt::write(&mut output, fmt) {
-            Ok(()) => Ok(()),
-            Err(..) => {
-                // check if the error came from the underlying `Write` or not
-                if output.error.is_err() {
-                    output.error
-                } else {
-                    // This shouldn't happen: the underlying stream did not error, but somehow
-                    // the formatter still errored?
-                    panic!(
-                        "a formatting trait implementation returned an error when the underlying stream did not"
-                    );
-                }
-            }
-        }
+        default_write_fmt(self, fmt)
     }
 
     /// Creates a "by reference" adapter for this instance of `Write`.