Auto merge of #106824 - m-ou-se:format-args-flatten, r=oli-obk

Flatten/inline format_args!() and (string and int) literal arguments into format_args!()

Implements https://github.com/rust-lang/rust/issues/78356

Gated behind `-Zflatten-format-args=yes`.

Part of #99012

This change inlines string literals, integer literals and nested format_args!() into format_args!() during ast lowering, making all of the following pairs result in equivalent hir:

```rust
println!("Hello, {}!", "World");
println!("Hello, World!");
```

```rust
println!("[info] {}", format_args!("error"));
println!("[info] error");
```

```rust
println!("[{}] {}", status, format_args!("error: {}", msg));
println!("[{}] error: {}", status, msg);
```

```rust
println!("{} + {} = {}", 1, 2, 1 + 2);
println!("1 + 2 = {}", 1 + 2);
```

And so on.

This is useful for macros. E.g. a `log::info!()` macro could just pass the tokens from the user directly into a `format_args!()` that gets efficiently flattened/inlined into a `format_args!("info: {}")`.

It also means that `dbg!(x)` will have its file, line, and expression name inlined:

```rust
eprintln!("[{}:{}] {} = {:#?}", file!(), line!(), stringify!(x), x); // before
eprintln!("[example.rs:1] x = {:#?}", x); // after
```

Which can be nice in some cases, but also means a lot more unique static strings than before if dbg!() is used a lot.

6 files changed, 232 insertions, 24 deletions

diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
index 6503bf2bab7..5d164bc4b3c 100644
--- a/compiler/rustc_ast/src/ast.rs
+++ b/compiler/rustc_ast/src/ast.rs
@@ -1184,6 +1184,15 @@ impl Expr {
         expr
     }
 
+    pub fn peel_parens_and_refs(&self) -> &Expr {
+        let mut expr = self;
+        while let ExprKind::Paren(inner) | ExprKind::AddrOf(BorrowKind::Ref, _, inner) = &expr.kind
+        {
+            expr = inner;
+        }
+        expr
+    }
+
     /// Attempts to reparse as `Ty` (for diagnostic purposes).
     pub fn to_ty(&self) -> Option<P<Ty>> {
         let kind = match &self.kind {
diff --git a/compiler/rustc_ast/src/format.rs b/compiler/rustc_ast/src/format.rs
index d021bea5eca..356b9bb6371 100644
--- a/compiler/rustc_ast/src/format.rs
+++ b/compiler/rustc_ast/src/format.rs
@@ -131,8 +131,8 @@ impl FormatArguments {
         &self.arguments[..]
     }
 
-    pub fn all_args_mut(&mut self) -> &mut [FormatArgument] {
-        &mut self.arguments[..]
+    pub fn all_args_mut(&mut self) -> &mut Vec<FormatArgument> {
+        &mut self.arguments
     }
 }
 
diff --git a/compiler/rustc_ast_lowering/src/format.rs b/compiler/rustc_ast_lowering/src/format.rs
index 4095e225a80..72352b138cb 100644
--- a/compiler/rustc_ast_lowering/src/format.rs
+++ b/compiler/rustc_ast_lowering/src/format.rs
@@ -7,13 +7,172 @@ use rustc_hir as hir;
 use rustc_span::{
     sym,
     symbol::{kw, Ident},
-    Span,
+    Span, Symbol,
 };
+use std::borrow::Cow;
 
 impl<'hir> LoweringContext<'_, 'hir> {
     pub(crate) fn lower_format_args(&mut self, sp: Span, fmt: &FormatArgs) -> hir::ExprKind<'hir> {
-        expand_format_args(self, sp, fmt)
+        // Never call the const constructor of `fmt::Arguments` if the
+        // format_args!() had any arguments _before_ flattening/inlining.
+        let allow_const = fmt.arguments.all_args().is_empty();
+        let mut fmt = Cow::Borrowed(fmt);
+        if self.tcx.sess.opts.unstable_opts.flatten_format_args {
+            fmt = flatten_format_args(fmt);
+            fmt = inline_literals(fmt);
+        }
+        expand_format_args(self, sp, &fmt, allow_const)
+    }
+}
+
+/// Flattens nested `format_args!()` into one.
+///
+/// Turns
+///
+/// `format_args!("a {} {} {}.", 1, format_args!("b{}!", 2), 3)`
+///
+/// into
+///
+/// `format_args!("a {} b{}! {}.", 1, 2, 3)`.
+fn flatten_format_args(mut fmt: Cow<'_, FormatArgs>) -> Cow<'_, FormatArgs> {
+    let mut i = 0;
+    while i < fmt.template.len() {
+        if let FormatArgsPiece::Placeholder(placeholder) = &fmt.template[i]
+            && let FormatTrait::Display | FormatTrait::Debug = &placeholder.format_trait
+            && let Ok(arg_index) = placeholder.argument.index
+            && let arg = fmt.arguments.all_args()[arg_index].expr.peel_parens_and_refs()
+            && let ExprKind::FormatArgs(_) = &arg.kind
+            // Check that this argument is not used by any other placeholders.
+            && fmt.template.iter().enumerate().all(|(j, p)|
+                i == j ||
+                !matches!(p, FormatArgsPiece::Placeholder(placeholder)
+                    if placeholder.argument.index == Ok(arg_index))
+            )
+        {
+            // Now we need to mutate the outer FormatArgs.
+            // If this is the first time, this clones the outer FormatArgs.
+            let fmt = fmt.to_mut();
+
+            // Take the inner FormatArgs out of the outer arguments, and
+            // replace it by the inner arguments. (We can't just put those at
+            // the end, because we need to preserve the order of evaluation.)
+
+            let args = fmt.arguments.all_args_mut();
+            let remaining_args = args.split_off(arg_index + 1);
+            let old_arg_offset = args.len();
+            let mut fmt2 = &mut args.pop().unwrap().expr; // The inner FormatArgs.
+            let fmt2 = loop { // Unwrap the Expr to get to the FormatArgs.
+                match &mut fmt2.kind {
+                    ExprKind::Paren(inner) | ExprKind::AddrOf(BorrowKind::Ref, _, inner) => fmt2 = inner,
+                    ExprKind::FormatArgs(fmt2) => break fmt2,
+                    _ => unreachable!(),
+                }
+            };
+
+            args.append(fmt2.arguments.all_args_mut());
+            let new_arg_offset = args.len();
+            args.extend(remaining_args);
+
+            // Correct the indexes that refer to the arguments after the newly inserted arguments.
+            for_all_argument_indexes(&mut fmt.template, |index| {
+                if *index >= old_arg_offset {
+                    *index -= old_arg_offset;
+                    *index += new_arg_offset;
+                }
+            });
+
+            // Now merge the placeholders:
+
+            let rest = fmt.template.split_off(i + 1);
+            fmt.template.pop(); // remove the placeholder for the nested fmt args.
+            // Insert the pieces from the nested format args, but correct any
+            // placeholders to point to the correct argument index.
+            for_all_argument_indexes(&mut fmt2.template, |index| *index += arg_index);
+            fmt.template.append(&mut fmt2.template);
+            fmt.template.extend(rest);
+
+            // Don't increment `i` here, so we recurse into the newly added pieces.
+        } else {
+            i += 1;
+        }
     }
+    fmt
+}
+
+/// Inline literals into the format string.
+///
+/// Turns
+///
+/// `format_args!("Hello, {}! {} {}", "World", 123, x)`
+///
+/// into
+///
+/// `format_args!("Hello, World! 123 {}", x)`.
+fn inline_literals(mut fmt: Cow<'_, FormatArgs>) -> Cow<'_, FormatArgs> {
+    let mut was_inlined = vec![false; fmt.arguments.all_args().len()];
+    let mut inlined_anything = false;
+
+    for i in 0..fmt.template.len() {
+        let FormatArgsPiece::Placeholder(placeholder) = &fmt.template[i] else { continue };
+        let Ok(arg_index) = placeholder.argument.index else { continue };
+
+        let mut literal = None;
+
+        if let FormatTrait::Display = placeholder.format_trait
+            && placeholder.format_options == Default::default()
+            && let arg = fmt.arguments.all_args()[arg_index].expr.peel_parens_and_refs()
+            && let ExprKind::Lit(lit) = arg.kind
+        {
+            if let token::LitKind::Str | token::LitKind::StrRaw(_) = lit.kind
+                && let Ok(LitKind::Str(s, _)) = LitKind::from_token_lit(lit)
+            {
+                literal = Some(s);
+            } else if let token::LitKind::Integer = lit.kind
+                && let Ok(LitKind::Int(n, _)) = LitKind::from_token_lit(lit)
+            {
+                literal = Some(Symbol::intern(&n.to_string()));
+            }
+        }
+
+        if let Some(literal) = literal {
+            // Now we need to mutate the outer FormatArgs.
+            // If this is the first time, this clones the outer FormatArgs.
+            let fmt = fmt.to_mut();
+            // Replace the placeholder with the literal.
+            fmt.template[i] = FormatArgsPiece::Literal(literal);
+            was_inlined[arg_index] = true;
+            inlined_anything = true;
+        }
+    }
+
+    // Remove the arguments that were inlined.
+    if inlined_anything {
+        let fmt = fmt.to_mut();
+
+        let mut remove = was_inlined;
+
+        // Don't remove anything that's still used.
+        for_all_argument_indexes(&mut fmt.template, |index| remove[*index] = false);
+
+        // Drop all the arguments that are marked for removal.
+        let mut remove_it = remove.iter();
+        fmt.arguments.all_args_mut().retain(|_| remove_it.next() != Some(&true));
+
+        // Calculate the mapping of old to new indexes for the remaining arguments.
+        let index_map: Vec<usize> = remove
+            .into_iter()
+            .scan(0, |i, remove| {
+                let mapped = *i;
+                *i += !remove as usize;
+                Some(mapped)
+            })
+            .collect();
+
+        // Correct the indexes that refer to arguments that have shifted position.
+        for_all_argument_indexes(&mut fmt.template, |index| *index = index_map[*index]);
+    }
+
+    fmt
 }
 
 #[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
@@ -189,11 +348,26 @@ fn expand_format_args<'hir>(
     ctx: &mut LoweringContext<'_, 'hir>,
     macsp: Span,
     fmt: &FormatArgs,
+    allow_const: bool,
 ) -> hir::ExprKind<'hir> {
+    let mut incomplete_lit = String::new();
     let lit_pieces =
         ctx.arena.alloc_from_iter(fmt.template.iter().enumerate().filter_map(|(i, piece)| {
             match piece {
-                &FormatArgsPiece::Literal(s) => Some(ctx.expr_str(fmt.span, s)),
+                &FormatArgsPiece::Literal(s) => {
+                    // Coalesce adjacent literal pieces.
+                    if let Some(FormatArgsPiece::Literal(_)) = fmt.template.get(i + 1) {
+                        incomplete_lit.push_str(s.as_str());
+                        None
+                    } else if !incomplete_lit.is_empty() {
+                        incomplete_lit.push_str(s.as_str());
+                        let s = Symbol::intern(&incomplete_lit);
+                        incomplete_lit.clear();
+                        Some(ctx.expr_str(fmt.span, s))
+                    } else {
+                        Some(ctx.expr_str(fmt.span, s))
+                    }
+                }
                 &FormatArgsPiece::Placeholder(_) => {
                     // Inject empty string before placeholders when not already preceded by a literal piece.
                     if i == 0 || matches!(fmt.template[i - 1], FormatArgsPiece::Placeholder(_)) {
@@ -244,6 +418,18 @@ 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(lit_pieces)
+        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([lit_pieces]);
+        return hir::ExprKind::Call(new, new_args);
+    }
+
     // If the args array contains exactly all the original arguments once,
     // in order, we can use a simple array instead of a `match` construction.
     // However, if there's a yield point in any argument except the first one,
@@ -290,25 +476,14 @@ fn expand_format_args<'hir>(
         let args_ident = Ident::new(sym::args, macsp);
         let (args_pat, args_hir_id) = ctx.pat_ident(macsp, args_ident);
         let args = ctx.arena.alloc_from_iter(argmap.iter().map(|&(arg_index, ty)| {
-            if let Some(arg) = arguments.get(arg_index) {
-                let sp = arg.expr.span.with_ctxt(macsp.ctxt());
-                let args_ident_expr = ctx.expr_ident(macsp, args_ident, args_hir_id);
-                let arg = ctx.arena.alloc(ctx.expr(
-                    sp,
-                    hir::ExprKind::Field(
-                        args_ident_expr,
-                        Ident::new(sym::integer(arg_index), macsp),
-                    ),
-                ));
-                make_argument(ctx, sp, arg, ty)
-            } else {
-                ctx.expr(
-                    macsp,
-                    hir::ExprKind::Err(
-                        ctx.tcx.sess.delay_span_bug(macsp, format!("no arg at {arg_index}")),
-                    ),
-                )
-            }
+            let arg = &arguments[arg_index];
+            let sp = arg.expr.span.with_ctxt(macsp.ctxt());
+            let args_ident_expr = ctx.expr_ident(macsp, args_ident, args_hir_id);
+            let arg = ctx.arena.alloc(ctx.expr(
+                sp,
+                hir::ExprKind::Field(args_ident_expr, Ident::new(sym::integer(arg_index), macsp)),
+            ));
+            make_argument(ctx, sp, arg, ty)
         }));
         let elements: Vec<_> = arguments
             .iter()
@@ -409,3 +584,22 @@ fn may_contain_yield_point(e: &ast::Expr) -> bool {
     visitor.visit_expr(e);
     visitor.0
 }
+
+fn for_all_argument_indexes(template: &mut [FormatArgsPiece], mut f: impl FnMut(&mut usize)) {
+    for piece in template {
+        let FormatArgsPiece::Placeholder(placeholder) = piece else { continue };
+        if let Ok(index) = &mut placeholder.argument.index {
+            f(index);
+        }
+        if let Some(FormatCount::Argument(FormatArgPosition { index: Ok(index), .. })) =
+            &mut placeholder.format_options.width
+        {
+            f(index);
+        }
+        if let Some(FormatCount::Argument(FormatArgPosition { index: Ok(index), .. })) =
+            &mut placeholder.format_options.precision
+        {
+            f(index);
+        }
+    }
+}
diff --git a/compiler/rustc_interface/src/tests.rs b/compiler/rustc_interface/src/tests.rs
index 18d84a7023a..014810dba9c 100644
--- a/compiler/rustc_interface/src/tests.rs
+++ b/compiler/rustc_interface/src/tests.rs
@@ -744,6 +744,7 @@ fn test_unstable_options_tracking_hash() {
     tracked!(emit_thin_lto, false);
     tracked!(export_executable_symbols, true);
     tracked!(fewer_names, Some(true));
+    tracked!(flatten_format_args, true);
     tracked!(force_unstable_if_unmarked, true);
     tracked!(fuel, Some(("abc".to_string(), 99)));
     tracked!(function_sections, Some(false));
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
index b466a3fcdee..0548379dc2f 100644
--- a/compiler/rustc_session/src/options.rs
+++ b/compiler/rustc_session/src/options.rs
@@ -1422,6 +1422,9 @@ options! {
     fewer_names: Option<bool> = (None, parse_opt_bool, [TRACKED],
         "reduce memory use by retaining fewer names within compilation artifacts (LLVM-IR) \
         (default: no)"),
+    flatten_format_args: bool = (false, parse_bool, [TRACKED],
+        "flatten nested format_args!() and literals into a simplified format_args!() call \
+        (default: no)"),
     force_unstable_if_unmarked: bool = (false, parse_bool, [TRACKED],
         "force all crates to be `rustc_private` unstable (default: no)"),
     fuel: Option<(String, u64)> = (None, parse_optimization_fuel, [TRACKED],
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
index 0154c719ef6..abe5af8f9e0 100644
--- a/compiler/rustc_span/src/symbol.rs
+++ b/compiler/rustc_span/src/symbol.rs
@@ -984,6 +984,7 @@ symbols! {
         never_type_fallback,
         new,
         new_binary,
+        new_const,
         new_debug,
         new_display,
         new_lower_exp,