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

Reserve before `write_fmt` for owned buffers

`fmt::Arguments::estimated_capacity()` is currently only used by `fmt::format()` to reserve the initial string capacity. Also use it in `fmt::Write::write_fmt` for `String` and `OsString`; and `io::Write::write_fmt` 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

@@ -3172,6 +3172,15 @@ impl fmt::Write for String {
         self.push(c);
         Ok(())
     }
+
+    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

@@ -636,6 +636,9 @@ impl<'a> Arguments<'a> {
     /// when using `format!`. Note: this is neither the lower nor upper bound.
     #[inline]
     pub fn estimated_capacity(&self) -> usize {
+        // Since the format args are constructed from a single string literal in
+        // `format_args!`, the total length of the pieces is under `isize::MAX`
+        // and this sum cannot overflow.
         let pieces_length: usize = self.pieces.iter().map(|x| x.len()).sum();
 
         if self.args.is_empty() {
@@ -649,7 +652,8 @@ impl<'a> Arguments<'a> {
             // There are some arguments, so any additional push
             // will reallocate the string. To avoid that,
             // we're "pre-doubling" the capacity here.
-            pieces_length.checked_mul(2).unwrap_or(0)
+            // It cannot overflow, because the maximum length is `isize::MAX`.
+            pieces_length * 2
         }
     }
 }

library/std/src/ffi/os_str.rs

@@ -799,10 +799,20 @@ 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(())
     }
+
+    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.
@@ -647,6 +647,15 @@ where
         Ok(())
     }
 
+    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> io::Result<()> {
+        if let Some(s) = args.as_statically_known_str() {
+            self.write_all(s.as_bytes())
+        } else {
+            self.inner.reserve(args.estimated_capacity());
+            io::default_write_fmt(self, args)
+        }
+    }
+
     #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
@@ -681,6 +690,15 @@ where
         Ok(())
     }
 
+    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> io::Result<()> {
+        if let Some(s) = args.as_statically_known_str() {
+            self.write_all(s.as_bytes())
+        } else {
+            self.inner.reserve(args.estimated_capacity());
+            io::default_write_fmt(self, args)
+        }
+    }
+
     #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())

library/std/src/io/impls.rs

@@ -511,6 +511,15 @@ impl<A: Allocator> Write for Vec<u8, A> {
         Ok(())
     }
 
+    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> io::Result<()> {
+        if let Some(s) = args.as_statically_known_str() {
+            self.write_all(s.as_bytes())
+        } else {
+            self.reserve(args.estimated_capacity());
+            io::default_write_fmt(self, args)
+        }
+    }
+
     #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())
@@ -662,6 +671,15 @@ impl<A: Allocator> Write for VecDeque<u8, A> {
         Ok(())
     }
 
+    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> io::Result<()> {
+        if let Some(s) = args.as_statically_known_str() {
+            self.write_all(s.as_bytes())
+        } else {
+            self.reserve(args.estimated_capacity());
+            io::default_write_fmt(self, args)
+        }
+    }
+
     #[inline]
     fn flush(&mut self) -> io::Result<()> {
         Ok(())