Optimize fmt::Arguments for the static string case.

fmt::Arguments can now store a &str without any indirection and
additional storage: There is now a fmt::Arguments::from_str() function.

Author: m-ou-se

compiler/rustc_ast_lowering/src/format.rs

@@ -401,12 +401,19 @@ fn expand_format_args<'hir>(
     let default_options = 0xE000_0020_0000_0000;
     let mut implicit_arg_index = 0;
 
-    for (i, piece) in fmt.template.iter().enumerate() {
+    let template = if fmt.template.is_empty() {
+        // Treat empty templates as a single literal piece (with an empty string),
+        // so we produce `from_str("")` for those.
+        &[FormatArgsPiece::Literal(sym::empty)][..]
+    } else {
+        &fmt.template[..]
+    };
+
+    for (i, piece) in template.iter().enumerate() {
         match piece {
             &FormatArgsPiece::Literal(sym) => {
-                assert!(!sym.is_empty());
                 // Coalesce adjacent literal pieces.
-                if let Some(FormatArgsPiece::Literal(_)) = fmt.template.get(i + 1) {
+                if let Some(FormatArgsPiece::Literal(_)) = template.get(i + 1) {
                     incomplete_lit.push_str(sym.as_str());
                     continue;
                 }
@@ -420,6 +427,26 @@ fn expand_format_args<'hir>(
                     (sym, len)
                 };
 
+                // If this is the last piece and was the only piece, that means
+                // there are no placeholders and the entire format string is just a literal.
+                //
+                // In that case, we don't need an array of `Piece`s: we can just use `from_str`.
+                if i + 1 == template.len() && pieces.is_empty() {
+                    // Generate:
+                    //     <core::fmt::Arguments>::from_str("meow")
+                    let from_str = ctx.arena.alloc(ctx.expr_lang_item_type_relative(
+                        macsp,
+                        hir::LangItem::FormatArguments,
+                        if allow_const { sym::from_str } else { sym::from_str_nonconst },
+                    ));
+                    let s = ctx.expr_str(fmt.span, sym);
+                    let args = ctx.arena.alloc_from_iter([s]);
+                    return hir::ExprKind::Call(from_str, args);
+                }
+
+                // Producing a `Piece::num(0)` would be problematic, as that is the terminator.
+                assert!(len > 0);
+
                 // ```
                 //  Piece::num(4),
                 // ```
@@ -456,7 +483,7 @@ fn expand_format_args<'hir>(
                 // placeholder is implied by two consequtive string pieces.
                 if bits == default_options + implicit_arg_index {
                     if let (Some(FormatArgsPiece::Literal(_)), Some(FormatArgsPiece::Literal(_))) =
-                        (fmt.template.get(i.wrapping_sub(1)), fmt.template.get(i + 1))
+                        (template.get(i.wrapping_sub(1)), template.get(i + 1))
                     {
                         implicit_arg_index += 1;
                         continue;
@@ -516,18 +543,6 @@ fn expand_format_args<'hir>(
 
     let arguments = fmt.arguments.all_args();
 
-    if allow_const && arguments.is_empty() && argmap.is_empty() {
-        // Generate:
-        //     <core::fmt::Arguments>::new_const(template)
-        let new = ctx.arena.alloc(ctx.expr_lang_item_type_relative(
-            macsp,
-            hir::LangItem::FormatArguments,
-            sym::new_const,
-        ));
-        let new_args = ctx.arena.alloc_from_iter([template]);
-        return hir::ExprKind::Call(new, new_args);
-    }
-
     let (let_statements, args) = if arguments.is_empty() {
         // Generate:
         //     []

compiler/rustc_const_eval/src/check_consts/ops.rs

@@ -345,7 +345,9 @@ fn build_error_for_const_call<'tcx>(
             note_trait_if_possible(&mut err, self_ty, tcx.require_lang_item(LangItem::Deref, span));
             err
         }
-        _ if tcx.opt_parent(callee) == tcx.get_diagnostic_item(sym::FmtArgumentsNew) => {
+        _ if tcx.is_diagnostic_item(sym::FmtTemplateNew, callee)
+            || tcx.is_diagnostic_item(sym::FmtArgumentsFromStrNonconst, callee) =>
+        {
             ccx.dcx().create_err(errors::NonConstFmtMacroCall {
                 span,
                 kind: ccx.const_kind(),

compiler/rustc_span/src/symbol.rs

@@ -235,7 +235,8 @@ symbols! {
         Error,
         File,
         FileType,
-        FmtArgumentsNew,
+        FmtArgumentsFromStrNonconst,
+        FmtTemplateNew,
         FmtWrite,
         Fn,
         FnMut,
@@ -1103,7 +1104,9 @@ symbols! {
         from_output,
         from_residual,
         from_size_align_unchecked,
+        from_str,
         from_str_method,
+        from_str_nonconst,
         from_u16,
         from_usize,
         from_yeet,

library/core/src/fmt/mod.rs

@@ -630,45 +630,39 @@ impl<'a> Arguments<'a> {
     /// when using `format!`. Note: this is neither the lower nor upper bound.
     #[inline]
     pub fn estimated_capacity(&self) -> usize {
+        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 template = self.template;
-        let mut has_placeholders = false;
         loop {
             // SAFETY: We can assume the template is valid.
             unsafe {
-                let n = template.next().i;
+                let n = template.next_piece().i;
                 if n == 0 {
                     // End of template.
                     break;
                 } else if n <= isize::MAX as _ {
                     // Literal string piece.
-                    if length != 0 {
-                        // More than one literal string piece means we have placeholders.
-                        has_placeholders = true;
-                    }
                     length += n as usize;
-                    let _ptr = template.next(); // Skip the string pointer.
+                    let _ptr = template.next_piece(); // Skip the string pointer.
                 } else {
                     // Placeholder piece.
                     if length == 0 {
                         starts_with_placeholder = true;
                     }
-                    has_placeholders = true;
                     // Skip remainder of placeholder:
                     #[cfg(target_pointer_width = "32")]
-                    let _ = template.next();
+                    let _ = template.next_piece();
                     #[cfg(target_pointer_width = "16")]
-                    let _ = (template.next(), template.next(), template.next());
+                    let _ = (template.next_piece(), template.next_piece(), template.next_piece());
                 }
             }
         }
 
-        if !has_placeholders {
-            // If the template has no placeholders, we know the length exactly.
-            length
-        } else if starts_with_placeholder && length < 16 {
+        if starts_with_placeholder && length < 16 {
             // If the format string starts with a placeholder,
             // don't preallocate anything, unless length
             // of literal pieces is significant.
@@ -730,27 +724,17 @@ impl<'a> Arguments<'a> {
     #[must_use]
     #[inline]
     pub const fn as_str(&self) -> Option<&'static str> {
-        let mut template = self.template;
-        // SAFETY: We can assume the template is valid.
-        let n = unsafe { template.next().i };
-        if n == 0 {
-            // The template is empty.
-            return Some("");
-        }
-        if n <= isize::MAX as _ {
-            // Template starts with a string piece.
-            // SAFETY: We can assume the template is valid.
-            unsafe {
-                let ptr = template.next().p;
-                if template.next().i == 0 {
-                    // The template has only one piece.
-                    return Some(str::from_utf8_unchecked(crate::slice::from_raw_parts(
-                        ptr, n as usize,
-                    )));
-                }
-            }
+        if let Some(len) = self.template.as_str_len() {
+            // SAFETY: This fmt::Arguments stores a &'static str.
+            Some(unsafe {
+                str::from_utf8_unchecked(crate::slice::from_raw_parts(
+                    self.args.cast().as_ptr(),
+                    len,
+                ))
+            })
+        } else {
+            None
         }
-        None
     }
 
     /// Same as [`Arguments::as_str`], but will only return `Some(s)` if it can be determined at compile time.
@@ -1493,6 +1477,10 @@ pub trait UpperExp: PointeeSized {
 /// [`write!`]: crate::write!
 #[stable(feature = "rust1", since = "1.0.0")]
 pub fn write(output: &mut dyn Write, fmt: Arguments<'_>) -> Result {
+    if let Some(s) = fmt.as_str() {
+        return output.write_str(s);
+    }
+
     let mut template = fmt.template;
     let args = fmt.args;
 
@@ -1501,7 +1489,7 @@ pub fn write(output: &mut dyn Write, fmt: Arguments<'_>) -> Result {
 
     loop {
         // SAFETY: We can assume the template is valid.
-        let n = unsafe { template.next().i };
+        let n = unsafe { template.next_piece().i };
         if n == 0 {
             // End of template.
             return Ok(());
@@ -1518,7 +1506,7 @@ pub fn write(output: &mut dyn Write, fmt: Arguments<'_>) -> Result {
                 unsafe { arg.fmt(&mut Formatter::new(output, options)) }?;
             }
             // SAFETY: We can assume the strings in the template are valid.
-            let s = unsafe { crate::str::from_raw_parts(template.next().p, n as usize) };
+            let s = unsafe { crate::str::from_raw_parts(template.next_piece().p, n as usize) };
             output.write_str(s)?;
             last_piece_was_str = true;
         } else {
@@ -1527,13 +1515,13 @@ pub fn write(output: &mut dyn Write, fmt: Arguments<'_>) -> Result {
             let (high, low) = ((n >> 32) as u32, n as u32);
             #[cfg(target_pointer_width = "32")]
             // SAFETY: We can assume the template is valid.
-            let (high, low) = (n as u32, unsafe { template.next().i } as u32);
+            let (high, low) = (n as u32, unsafe { template.next_piece().i } as u32);
             #[cfg(target_pointer_width = "16")]
             // SAFETY: We can assume the template is valid.
             let (high, low) = unsafe {
                 (
-                    (n as u32) << 16 | template.next().i as u32,
-                    (template.next().i as u32) << 16 | template.next().i as u32,
+                    (n as u32) << 16 | template.next_piece().i as u32,
+                    (template.next_piece().i as u32) << 16 | template.next_piece().i as u32,
                 )
             };
             let arg_index = (low & 0x3FF) as usize;

library/core/src/fmt/rt.rs

@@ -9,6 +9,7 @@
 use super::*;
 use crate::hint::unreachable_unchecked;
 use crate::marker::PhantomData;
+use crate::num::NonZeroUsize;
 use crate::ptr::NonNull;
 
 #[lang = "format_template"]
@@ -23,12 +24,29 @@ unsafe impl Sync for Template<'_> {}
 
 impl<'a> Template<'a> {
     #[inline]
-    pub const unsafe fn new<const N: usize>(pieces: &'a [rt::Piece; N]) -> Self {
+    #[rustc_diagnostic_item = "FmtTemplateNew"]
+    pub unsafe fn new<const N: usize>(pieces: &'a [rt::Piece; N]) -> Self {
         Self { pieces: NonNull::from_ref(pieces).cast(), lifetime: PhantomData }
     }
 
     #[inline]
-    pub const unsafe fn next(&mut self) -> Piece {
+    pub const unsafe fn new_str_len(len: usize) -> Self {
+        // SAFETY: We set the lowest bit, so it's nonzero.
+        let bits = unsafe { NonZeroUsize::new_unchecked(len << 1 | 1) };
+        Self { pieces: NonNull::without_provenance(bits), lifetime: PhantomData }
+    }
+
+    #[inline]
+    pub const fn as_str_len(self) -> Option<usize> {
+        // SAFETY: During const eval, `self.pieces` must have come from a usize, not a pointer,
+        // because `new()` above is not const.
+        // Outside const eval, transmuting a pointer to a usize is fine.
+        let bits = unsafe { crate::mem::transmute::<_, usize>(self.pieces) };
+        if bits & 1 == 1 { Some(bits >> 1) } else { None }
+    }
+
+    #[inline]
+    pub const unsafe fn next_piece(&mut self) -> Piece {
         // SAFETY: Guaranteed by caller.
         unsafe {
             let piece = *self.pieces.as_ref();
@@ -40,6 +58,7 @@ impl<'a> Template<'a> {
 
 #[lang = "format_piece"]
 #[derive(Copy, Clone)]
+#[repr(align(2))]
 pub union Piece {
     pub i: usize,
     pub p: *const u8,
@@ -224,31 +243,30 @@ impl Argument<'_> {
 
 /// Used by the format_args!() macro to create a fmt::Arguments object.
 #[doc(hidden)]
-#[rustc_diagnostic_item = "FmtArgumentsNew"]
 impl<'a> Arguments<'a> {
     #[inline]
-    pub const fn new_const(template: rt::Template<'a>) -> Arguments<'a> {
-        Arguments { template, args: NonNull::dangling() }
-    }
-
-    #[inline]
-    pub fn new<const N: usize>(
+    pub const fn new<const N: usize>(
         template: rt::Template<'a>,
         args: &'a [rt::Argument<'a>; N],
     ) -> Arguments<'a> {
         Arguments { template, args: NonNull::from_ref(args).cast() }
     }
 
-    // These two methods are used in library/core/src/panicking.rs to create a
-    // `fmt::Arguments` for a `&'static str`.
     #[inline]
-    pub(crate) const fn pieces_for_str(s: &'static str) -> [rt::Piece; 3] {
-        [rt::Piece { i: s.len() as _ }, rt::Piece::str(s), rt::Piece { i: 0 }]
+    pub const fn from_str(s: &'static str) -> Arguments<'a> {
+        Arguments {
+            // SAFETY: This is the "static str" representation of fmt::Arguments.
+            template: unsafe { rt::Template::new_str_len(s.len()) },
+            args: NonNull::from_ref(s).cast(),
+        }
     }
-    /// Safety: Only call this with the result of `pieces_for_str`.
+
+    // Same as `from_str`, but not const.
+    // Used by format_args!() expansion when arguments are inlined,
+    // e.g. format_args!("{}", 123), which is not allowed in const.
     #[inline]
-    pub(crate) const unsafe fn from_pieces(p: &'a [rt::Piece; 3]) -> Self {
-        // SAFETY: Guaranteed by caller.
-        Self::new_const(unsafe { rt::Template::new(p) })
+    #[rustc_diagnostic_item = "FmtArgumentsFromStrNonconst"]
+    pub fn from_str_nonconst(s: &'static str) -> Arguments<'a> {
+        Arguments::from_str(s)
     }
 }