Use span of placeholders in format_args!() expansion.

Author: m-ou-se

compiler/rustc_ast_lowering/src/format.rs

@@ -2,7 +2,7 @@ use super::LoweringContext;
 use rustc_ast as ast;
 use rustc_ast::visit::{self, Visitor};
 use rustc_ast::*;
-use rustc_data_structures::fx::FxIndexSet;
+use rustc_data_structures::fx::FxIndexMap;
 use rustc_hir as hir;
 use rustc_span::{
     sym,
@@ -238,7 +238,7 @@ fn make_count<'hir>(
     ctx: &mut LoweringContext<'_, 'hir>,
     sp: Span,
     count: &Option<FormatCount>,
-    argmap: &mut FxIndexSet<(usize, ArgumentType)>,
+    argmap: &mut FxIndexMap<(usize, ArgumentType), Option<Span>>,
 ) -> hir::Expr<'hir> {
     match count {
         Some(FormatCount::Literal(n)) => {
@@ -252,7 +252,7 @@ fn make_count<'hir>(
         }
         Some(FormatCount::Argument(arg)) => {
             if let Ok(arg_index) = arg.index {
-                let (i, _) = argmap.insert_full((arg_index, ArgumentType::Usize));
+                let (i, _) = argmap.insert_full((arg_index, ArgumentType::Usize), arg.span);
                 let count_param = ctx.arena.alloc(ctx.expr_lang_item_type_relative(
                     sp,
                     hir::LangItem::FormatCount,
@@ -291,12 +291,14 @@ fn make_format_spec<'hir>(
     ctx: &mut LoweringContext<'_, 'hir>,
     sp: Span,
     placeholder: &FormatPlaceholder,
-    argmap: &mut FxIndexSet<(usize, ArgumentType)>,
+    argmap: &mut FxIndexMap<(usize, ArgumentType), Option<Span>>,
 ) -> hir::Expr<'hir> {
     let position = match placeholder.argument.index {
         Ok(arg_index) => {
-            let (i, _) =
-                argmap.insert_full((arg_index, ArgumentType::Format(placeholder.format_trait)));
+            let (i, _) = argmap.insert_full(
+                (arg_index, ArgumentType::Format(placeholder.format_trait)),
+                placeholder.span,
+            );
             ctx.expr_usize(sp, i)
         }
         Err(_) => ctx.expr(
@@ -386,15 +388,18 @@ fn expand_format_args<'hir>(
 
     // Create a list of all _unique_ (argument, format trait) combinations.
     // E.g. "{0} {0:x} {0} {1}" -> [(0, Display), (0, LowerHex), (1, Display)]
-    let mut argmap = FxIndexSet::default();
+    let mut argmap = FxIndexMap::default();
     for piece in &fmt.template {
         let FormatArgsPiece::Placeholder(placeholder) = piece else { continue };
         if placeholder.format_options != Default::default() {
             // Can't use basic form if there's any formatting options.
             use_format_options = true;
         }
         if let Ok(index) = placeholder.argument.index {
-            if !argmap.insert((index, ArgumentType::Format(placeholder.format_trait))) {
+            if argmap
+                .insert((index, ArgumentType::Format(placeholder.format_trait)), placeholder.span)
+                .is_some()
+            {
                 // Duplicate (argument, format trait) combination,
                 // which we'll only put once in the args array.
                 use_format_options = true;
@@ -438,7 +443,7 @@ fn expand_format_args<'hir>(
     // This is an optimization, speeding up compilation about 1-2% in some cases.
     // See https://github.com/rust-lang/rust/pull/106770#issuecomment-1380790609
     let use_simple_array = argmap.len() == arguments.len()
-        && argmap.iter().enumerate().all(|(i, &(j, _))| i == j)
+        && argmap.iter().enumerate().all(|(i, (&(j, _), _))| i == j)
         && arguments.iter().skip(1).all(|arg| !may_contain_yield_point(&arg.expr));
 
     let args = if use_simple_array {
@@ -452,14 +457,19 @@ fn expand_format_args<'hir>(
         let elements: Vec<_> = arguments
             .iter()
             .zip(argmap)
-            .map(|(arg, (_, ty))| {
-                let sp = arg.expr.span.with_ctxt(macsp.ctxt());
+            .map(|(arg, ((_, ty), placeholder_span))| {
+                let placeholder_span =
+                    placeholder_span.unwrap_or(arg.expr.span).with_ctxt(macsp.ctxt());
+                let arg_span = match arg.kind {
+                    FormatArgumentKind::Captured(_) => placeholder_span,
+                    _ => arg.expr.span.with_ctxt(macsp.ctxt()),
+                };
                 let arg = ctx.lower_expr(&arg.expr);
                 let ref_arg = ctx.arena.alloc(ctx.expr(
-                    sp,
+                    arg_span,
                     hir::ExprKind::AddrOf(hir::BorrowKind::Ref, hir::Mutability::Not, arg),
                 ));
-                make_argument(ctx, sp, ref_arg, ty)
+                make_argument(ctx, placeholder_span, ref_arg, ty)
             })
             .collect();
         ctx.expr_array_ref(macsp, ctx.arena.alloc_from_iter(elements))
@@ -475,16 +485,26 @@ 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)| {
-            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 args = ctx.arena.alloc_from_iter(argmap.iter().map(
+            |(&(arg_index, ty), &placeholder_span)| {
+                let arg = &arguments[arg_index];
+                let placeholder_span =
+                    placeholder_span.unwrap_or(arg.expr.span).with_ctxt(macsp.ctxt());
+                let arg_span = match arg.kind {
+                    FormatArgumentKind::Captured(_) => placeholder_span,
+                    _ => 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(
+                    arg_span,
+                    hir::ExprKind::Field(
+                        args_ident_expr,
+                        Ident::new(sym::integer(arg_index), macsp),
+                    ),
+                ));
+                make_argument(ctx, placeholder_span, arg, ty)
+            },
+        ));
         let elements: Vec<_> = arguments
             .iter()
             .map(|arg| {