Improve heuristics whether `format_args` string is a source literal

Previously, it only checked whether there was _a_ literal at the span of
the first argument, not whether the literal actually matched up. This
caused issues when a proc macro was generating a different literal with
the same span.

This requires an annoying special case for literals ending in `\n`
because otherwise `println` wouldn't give detailed diagnostics anymore
which would be bad.

4 files changed, 78 insertions, 2 deletions

diff --git a/compiler/rustc_parse_format/src/lib.rs b/compiler/rustc_parse_format/src/lib.rs
index ab0463045fa..84243e53faf 100644
--- a/compiler/rustc_parse_format/src/lib.rs
+++ b/compiler/rustc_parse_format/src/lib.rs
@@ -20,6 +20,7 @@ pub use Flag::*;
 pub use Piece::*;
 pub use Position::*;
 
+use rustc_lexer::unescape;
 use std::iter;
 use std::str;
 use std::string;
@@ -306,7 +307,7 @@ impl<'a> Parser<'a> {
         append_newline: bool,
         mode: ParseMode,
     ) -> Parser<'a> {
-        let (width_map, is_literal) = find_width_map_from_snippet(snippet, style);
+        let (width_map, is_literal) = find_width_map_from_snippet(s, snippet, style);
         Parser {
             mode,
             input: s,
@@ -844,6 +845,7 @@ impl<'a> Parser<'a> {
 /// written code (code snippet) and the `InternedString` that gets processed in the `Parser`
 /// in order to properly synthesise the intra-string `Span`s for error diagnostics.
 fn find_width_map_from_snippet(
+    input: &str,
     snippet: Option<string::String>,
     str_style: Option<usize>,
 ) -> (Vec<InnerWidthMapping>, bool) {
@@ -856,8 +858,27 @@ fn find_width_map_from_snippet(
         return (vec![], true);
     }
 
+    // Strip quotes.
     let snippet = &snippet[1..snippet.len() - 1];
 
+    // Macros like `println` add a newline at the end. That technically doens't make them "literals" anymore, but it's fine
+    // since we will never need to point our spans there, so we lie about it here by ignoring it.
+    // Since there might actually be newlines in the source code, we need to normalize away all trailing newlines.
+    // If we only trimmed it off the input, `format!("\n")` would cause a mismatch as here we they actually match up.
+    // Alternatively, we could just count the trailing newlines and only trim one from the input if they don't match up.
+    let input_no_nl = input.trim_end_matches('\n');
+    let Ok(unescaped) = unescape_string(snippet) else {
+        return (vec![], false);
+    };
+
+    let unescaped_no_nl = unescaped.trim_end_matches('\n');
+
+    if unescaped_no_nl != input_no_nl {
+        // The source string that we're pointing at isn't our input, so spans pointing at it will be incorrect.
+        // This can for example happen with proc macros that respan generated literals.
+        return (vec![], false);
+    }
+
     let mut s = snippet.char_indices();
     let mut width_mappings = vec![];
     while let Some((pos, c)) = s.next() {
@@ -936,9 +957,23 @@ fn find_width_map_from_snippet(
             _ => {}
         }
     }
+
     (width_mappings, true)
 }
 
+fn unescape_string(string: &str) -> Result<string::String, unescape::EscapeError> {
+    let mut buf = string::String::new();
+    let mut error = Ok(());
+    unescape::unescape_literal(string, unescape::Mode::Str, &mut |_, unescaped_char| {
+        match unescaped_char {
+            Ok(c) => buf.push(c),
+            Err(err) => error = Err(err),
+        }
+    });
+
+    error.map(|_| buf)
+}
+
 // Assert a reasonable size for `Piece`
 #[cfg(all(target_arch = "x86_64", target_pointer_width = "64"))]
 rustc_data_structures::static_assert_size!(Piece<'_>, 16);
diff --git a/src/test/ui/fmt/auxiliary/format-string-proc-macro.rs b/src/test/ui/fmt/auxiliary/format-string-proc-macro.rs
index e44a84776bc..539c8fb27b3 100644
--- a/src/test/ui/fmt/auxiliary/format-string-proc-macro.rs
+++ b/src/test/ui/fmt/auxiliary/format-string-proc-macro.rs
@@ -5,7 +5,8 @@
 
 extern crate proc_macro;
 
-use proc_macro::{Literal, Span, TokenStream, TokenTree};
+use proc_macro::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream, TokenTree};
+use std::iter::FromIterator;
 
 #[proc_macro]
 pub fn foo_with_input_span(input: TokenStream) -> TokenStream {
@@ -26,3 +27,14 @@ pub fn err_with_input_span(input: TokenStream) -> TokenStream {
 
     TokenStream::from(TokenTree::Literal(lit))
 }
+
+#[proc_macro]
+pub fn respan_to_invalid_format_literal(input: TokenStream) -> TokenStream {
+    let mut s = Literal::string("{");
+    s.set_span(input.into_iter().next().unwrap().span());
+    TokenStream::from_iter([
+        TokenTree::from(Ident::new("format", Span::call_site())),
+        TokenTree::from(Punct::new('!', Spacing::Alone)),
+        TokenTree::from(Group::new(Delimiter::Parenthesis, TokenTree::from(s).into())),
+    ])
+}
diff --git a/src/test/ui/fmt/respanned-literal-issue-106191.rs b/src/test/ui/fmt/respanned-literal-issue-106191.rs
new file mode 100644
index 00000000000..44642a10fc0
--- /dev/null
+++ b/src/test/ui/fmt/respanned-literal-issue-106191.rs
@@ -0,0 +1,10 @@
+// aux-build:format-string-proc-macro.rs
+
+extern crate format_string_proc_macro;
+
+fn main() {
+    format_string_proc_macro::respan_to_invalid_format_literal!("¡");
+    //~^ ERROR invalid format string: expected `'}'` but string was terminated
+    format_args!(r#concat!("¡        {"));
+    //~^ ERROR invalid format string: expected `'}'` but string was terminated
+}
diff --git a/src/test/ui/fmt/respanned-literal-issue-106191.stderr b/src/test/ui/fmt/respanned-literal-issue-106191.stderr
new file mode 100644
index 00000000000..73a3af65a38
--- /dev/null
+++ b/src/test/ui/fmt/respanned-literal-issue-106191.stderr
@@ -0,0 +1,19 @@
+error: invalid format string: expected `'}'` but string was terminated
+  --> $DIR/respanned-literal-issue-106191.rs:6:65
+   |
+LL |     format_string_proc_macro::respan_to_invalid_format_literal!("¡");
+   |                                                                 ^^^ expected `'}'` in format string
+   |
+   = note: if you intended to print `{`, you can escape it using `{{`
+
+error: invalid format string: expected `'}'` but string was terminated
+  --> $DIR/respanned-literal-issue-106191.rs:8:18
+   |
+LL |     format_args!(r#concat!("¡        {"));
+   |                  ^^^^^^^^^^^^^^^^^^^^^^^ expected `'}'` in format string
+   |
+   = note: if you intended to print `{`, you can escape it using `{{`
+   = note: this error originates in the macro `concat` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error: aborting due to 2 previous errors
+