Add a size_hint function to Display

library/alloc/src/borrow.rs

@@ -428,6 +428,12 @@ where
             Owned(ref o) => fmt::Display::fmt(o, f),
         }
     }
+    fn size_hint(&self) -> Option<usize> {
+        match *self {
+            Borrowed(ref b) => fmt::Display::size_hint(b),
+            Owned(ref o) => fmt::Display::size_hint(o),
+        }
+    }
 }
 
 #[stable(feature = "default", since = "1.11.0")]

library/alloc/src/boxed.rs

@@ -2271,6 +2271,9 @@ impl<T: fmt::Display + ?Sized, A: Allocator> fmt::Display for Box<T, A> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt::Display::fmt(&**self, f)
     }
+    fn size_hint(&self) -> Option<usize> {
+        fmt::Display::size_hint(&**self)
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

library/alloc/src/fmt.rs

@@ -649,6 +649,22 @@ use crate::string;
 pub fn format(args: Arguments<'_>) -> string::String {
     fn format_inner(args: Arguments<'_>) -> string::String {
         let capacity = args.estimated_capacity();
+        //
+        // fn debug_print(s: &str) {
+        //     unsafe extern "C" {
+        //         fn write(fd: i32, buf: *const u8, count: usize) -> isize;
+        //     }
+        //
+        //     unsafe {
+        //         write(1, s.as_ptr(), s.len()); // fd 1 = stdout
+        //     }
+        // }
+        //
+        // let str = capacity.to_string();
+        // debug_print("estimated_capacity: ");
+        // debug_print(&str);
+        // debug_print("\n");
+
         let mut output = string::String::with_capacity(capacity);
         output
             .write_fmt(args)

library/alloc/src/lib.rs

@@ -56,6 +56,7 @@
 //! [`Rc`]: rc
 //! [`RefCell`]: core::cell
 
+#![feature(split_array)]
 #![allow(unused_features)]
 #![allow(incomplete_features)]
 #![allow(unused_attributes)]

library/alloc/src/rc.rs

@@ -2826,6 +2826,9 @@ impl<T: ?Sized + fmt::Display, A: Allocator> fmt::Display for Rc<T, A> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt::Display::fmt(&**self, f)
     }
+    fn size_hint(&self) -> Option<usize> {
+        fmt::Display::size_hint(&**self)
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -3971,6 +3974,9 @@ impl<T: ?Sized + fmt::Display, A: Allocator> fmt::Display for UniqueRc<T, A> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt::Display::fmt(&**self, f)
     }
+    fn size_hint(&self) -> Option<usize> {
+        fmt::Display::size_hint(&**self)
+    }
 }
 
 #[unstable(feature = "unique_rc_arc", issue = "112566")]

library/alloc/src/string.rs

@@ -2748,6 +2748,10 @@ impl fmt::Display for String {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt::Display::fmt(&**self, f)
     }
+    #[inline]
+    fn size_hint(&self) -> Option<usize> {
+        fmt::Display::size_hint(&**self)
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

library/alloc/src/sync.rs

@@ -3762,6 +3762,9 @@ impl<T: ?Sized + fmt::Display, A: Allocator> fmt::Display for Arc<T, A> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt::Display::fmt(&**self, f)
     }
+    fn size_hint(&self) -> Option<usize> {
+        fmt::Display::size_hint(&**self)
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

library/core/src/fmt/mod.rs

@@ -743,6 +743,16 @@ impl<'a> Arguments<'a> {
     }
 }
 
+// fn debug_print(s: &str) {
+//     unsafe extern "C" {
+//         fn write(fd: i32, buf: *const u8, count: usize) -> isize;
+//     }
+//
+//     unsafe {
+//         write(1, s.as_ptr(), s.len()); // fd 1 = stdout
+//     }
+// }
+
 #[doc(hidden)]
 #[unstable(feature = "fmt_internals", issue = "none")]
 impl<'a> Arguments<'a> {
@@ -752,13 +762,16 @@ impl<'a> Arguments<'a> {
     /// when using `format!`. Note: this is neither the lower nor upper bound.
     #[inline]
     pub fn estimated_capacity(&self) -> usize {
+        // debug_print("estimated_capacity");
         if let Some(s) = self.as_str() {
             return s.len();
         }
         // Iterate over the template, counting the length of literal pieces.
         let mut length = 0usize;
-        let mut starts_with_placeholder = false;
+        let mut starts_with_placeholder_without_hint = false;
+        let mut has_placeholder_without_hint = false;
         let mut template = self.template;
+        let mut arg_index = 0;
         loop {
             // SAFETY: We can assume the template is valid.
             unsafe {
@@ -779,29 +792,42 @@ impl<'a> Arguments<'a> {
                 } else {
                     assert_unchecked(n >= 0xC0);
                     // Placeholder piece.
-                    if length == 0 {
-                        starts_with_placeholder = true;
-                    }
+                    // debug_print("placeholder\n");
                     // Skip remainder of placeholder:
                     let skip = (n & 1 != 0) as usize * 4 // flags (32 bit)
                         + (n & 2 != 0) as usize * 2  // width     (16 bit)
-                        + (n & 4 != 0) as usize * 2  // precision (16 bit)
-                        + (n & 8 != 0) as usize * 2; // arg_index (16 bit)
-                    template = template.add(skip as usize);
+                        + (n & 4 != 0) as usize * 2; // precision (16 bit)
+                    // + (n & 8 != 0) as usize * 2; // arg_index (16 bit)
+                    template = template.add(skip);
+                    let has_arg_index = (n & 8) != 0;
+                    if has_arg_index {
+                        arg_index = usize::from(u16::from_le_bytes(template.cast_array().read()));
+                        template = template.add(2);
+                    };
+                    let argument = self.args.add(arg_index).as_ref();
+                    let size_hint: Option<usize> = argument.hint();
+
+                    if let Some(hint) = size_hint {
+                        // debug_print("has hint\n");
+                        length += hint;
+                    } else {
+                        has_placeholder_without_hint = true;
+                        if length == 0 {
+                            starts_with_placeholder_without_hint = true;
+                        }
+                    }
+                    arg_index += 1;
                 }
             }
         }
 
-        if starts_with_placeholder && length < 16 {
+        if starts_with_placeholder_without_hint && length < 16 {
             // If the format string starts with a placeholder,
             // don't preallocate anything, unless length
             // of literal pieces is significant.
             0
         } else {
-            // There are some placeholders, so any additional push
-            // will reallocate the string. To avoid that,
-            // we're "pre-doubling" the capacity here.
-            length.wrapping_mul(2)
+            length.wrapping_mul(if has_placeholder_without_hint { 2 } else { 1 })
         }
     }
 }
@@ -1210,6 +1236,16 @@ pub trait Display: PointeeSized {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     fn fmt(&self, f: &mut Formatter<'_>) -> Result;
+
+    /// foobar
+    #[unstable(
+        feature = "split_array",
+        reason = "return type should have array as 2nd element",
+        issue = "90091"
+    )]
+    fn size_hint(&self) -> Option<usize> {
+        None
+    }
 }
 
 /// `o` formatting.
@@ -2894,7 +2930,27 @@ macro_rules! fmt_refs {
     }
 }
 
-fmt_refs! { Debug, Display, Octal, Binary, LowerHex, UpperHex, LowerExp, UpperExp }
+fmt_refs! { Debug, Octal, Binary, LowerHex, UpperHex, LowerExp, UpperExp }
+
+#[stable(feature = "rust1", since = "1.0.0")]
+impl<T: PointeeSized + Display> Display for &T {
+    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
+        Display::fmt(&**self, f)
+    }
+    fn size_hint(&self) -> Option<usize> {
+        Display::size_hint(&**self)
+    }
+}
+
+#[stable(feature = "rust1", since = "1.0.0")]
+impl<T: PointeeSized + Display> Display for &mut T {
+    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
+        Display::fmt(&**self, f)
+    }
+    fn size_hint(&self) -> Option<usize> {
+        Display::size_hint(&**self)
+    }
+}
 
 #[unstable(feature = "never_type", issue = "35121")]
 impl Debug for ! {
@@ -2925,6 +2981,13 @@ impl Display for bool {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         Display::fmt(if *self { "true" } else { "false" }, f)
     }
+
+    fn size_hint(&self) -> Option<usize> {
+        Some(match *self {
+            true => 4,
+            false => 5,
+        })
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -2981,6 +3044,9 @@ impl Display for str {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         f.pad(self)
     }
+    fn size_hint(&self) -> Option<usize> {
+        Some(self.len())
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -3006,6 +3072,9 @@ impl Display for char {
             f.pad(self.encode_utf8(&mut [0; char::MAX_LEN_UTF8]))
         }
     }
+    fn size_hint(&self) -> Option<usize> {
+        Some(1)
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

library/core/src/fmt/num.rs

@@ -149,6 +149,10 @@ macro_rules! impl_Display {
                     ${concat($fmt_fn, _small)}(*self as $T, true, f)
                 }
             }
+
+            fn size_hint(&self) -> Option<usize> {
+                Some((self.max(&1).ilog(10) + 1) as usize)
+            }
         }
 
         #[stable(feature = "rust1", since = "1.0.0")]
@@ -170,6 +174,10 @@ macro_rules! impl_Display {
                     return ${concat($fmt_fn, _small)}(self.unsigned_abs() as $T, *self >= 0, f);
                 }
             }
+
+            fn size_hint(&self) -> Option<usize> {
+                Some((self.abs().max(1).ilog(10) + 1) as usize + self.is_negative() as usize)
+            }
         }
 
         #[cfg(not(feature = "optimize_for_size"))]

library/core/src/fmt/rt.rs

@@ -17,6 +17,7 @@ enum ArgumentType<'a> {
         // was derived from a `&'a T`.
         value: NonNull<()>,
         formatter: unsafe fn(NonNull<()>, &mut Formatter<'_>) -> Result,
+        hint: unsafe fn(NonNull<()>) -> Option<usize>,
         _lifetime: PhantomData<&'a ()>,
     },
     Count(u16),
@@ -40,7 +41,7 @@ pub struct Argument<'a> {
 }
 
 macro_rules! argument_new {
-    ($t:ty, $x:expr, $f:expr) => {
+    ($t:ty, $x:expr, $f:expr, $hint:expr) => {
         Argument {
             // INVARIANT: this creates an `ArgumentType<'a>` from a `&'a T` and
             // a `fn(&T, ...)`, so the invariant is maintained.
@@ -78,6 +79,19 @@ macro_rules! argument_new {
                     let r = unsafe { ptr.cast::<$t>().as_ref() };
                     (func)(r, fmt)
                 },
+                #[cfg(not(any(sanitize = "cfi", sanitize = "kcfi")))]
+                hint: {
+                    let f: fn(&$t) -> Option<usize> = $hint;
+                    // SAFETY: This is only called with `value`, which has the right type.
+                    unsafe { core::mem::transmute(f) }
+                },
+                #[cfg(any(sanitize = "cfi", sanitize = "kcfi"))]
+                hint: |ptr: NonNull<()>| {
+                    let func = $hint;
+                    // SAFETY: This is the same type as the `value` field.
+                    let r = unsafe { ptr.cast::<$t>().as_ref() };
+                    (func)(r)
+                },
                 _lifetime: PhantomData,
             },
         }
@@ -87,43 +101,43 @@ macro_rules! argument_new {
 impl Argument<'_> {
     #[inline]
     pub const fn new_display<T: Display>(x: &T) -> Argument<'_> {
-        argument_new!(T, x, <T as Display>::fmt)
+        argument_new!(T, x, <T as Display>::fmt, <T as Display>::size_hint)
     }
     #[inline]
     pub const fn new_debug<T: Debug>(x: &T) -> Argument<'_> {
-        argument_new!(T, x, <T as Debug>::fmt)
+        argument_new!(T, x, <T as Debug>::fmt, |_: &T| None)
     }
     #[inline]
     pub const fn new_debug_noop<T: Debug>(x: &T) -> Argument<'_> {
-        argument_new!(T, x, |_: &T, _| Ok(()))
+        argument_new!(T, x, |_: &T, _| Ok(()), |_: &T| None)
     }
     #[inline]
     pub const fn new_octal<T: Octal>(x: &T) -> Argument<'_> {
-        argument_new!(T, x, <T as Octal>::fmt)
+        argument_new!(T, x, <T as Octal>::fmt, |_: &T| None)
     }
     #[inline]
     pub const fn new_lower_hex<T: LowerHex>(x: &T) -> Argument<'_> {
-        argument_new!(T, x, <T as LowerHex>::fmt)
+        argument_new!(T, x, <T as LowerHex>::fmt, |_: &T| None)
     }
     #[inline]
     pub const fn new_upper_hex<T: UpperHex>(x: &T) -> Argument<'_> {
-        argument_new!(T, x, <T as UpperHex>::fmt)
+        argument_new!(T, x, <T as UpperHex>::fmt, |_: &T| None)
     }
     #[inline]
     pub const fn new_pointer<T: Pointer>(x: &T) -> Argument<'_> {
-        argument_new!(T, x, <T as Pointer>::fmt)
+        argument_new!(T, x, <T as Pointer>::fmt, |_: &T| None)
     }
     #[inline]
     pub const fn new_binary<T: Binary>(x: &T) -> Argument<'_> {
-        argument_new!(T, x, <T as Binary>::fmt)
+        argument_new!(T, x, <T as Binary>::fmt, |_: &T| None)
     }
     #[inline]
     pub const fn new_lower_exp<T: LowerExp>(x: &T) -> Argument<'_> {
-        argument_new!(T, x, <T as LowerExp>::fmt)
+        argument_new!(T, x, <T as LowerExp>::fmt, |_: &T| None)
     }
     #[inline]
     pub const fn new_upper_exp<T: UpperExp>(x: &T) -> Argument<'_> {
-        argument_new!(T, x, <T as UpperExp>::fmt)
+        argument_new!(T, x, <T as UpperExp>::fmt, |_: &T| None)
     }
     #[inline]
     #[track_caller]
@@ -155,6 +169,22 @@ impl Argument<'_> {
         }
     }
 
+    #[inline]
+    pub(super) unsafe fn hint(&self) -> Option<usize> {
+        match self.ty {
+            // SAFETY:
+            // Because of the invariant that if `formatter` had the type
+            // `fn(&T, _) -> _` then `value` has type `&'b T` where `'b` is
+            // the lifetime of the `ArgumentType`, and because references
+            // and `NonNull` are ABI-compatible, this is completely equivalent
+            // to calling the original function passed to `new` with the
+            // original reference, which is sound.
+            ArgumentType::Placeholder { value, hint, .. } => unsafe { hint(value) },
+            // SAFETY: the caller promised this.
+            ArgumentType::Count(_) => None,
+        }
+    }
+
     #[inline]
     pub(super) const fn as_u16(&self) -> Option<u16> {
         match self.ty {