From 81afdbc161300549bb5bbfe0bec5c1a60e11243d Mon Sep 17 00:00:00 2001
From: Nicholas Nethercote <n.nethercote@gmail.com>
Date: Mon, 23 Sep 2024 13:04:36 +1000
Subject: [PATCH] Fix a problem with metavars and inner attributes.

---
 compiler/rustc_ast/src/tokenstream.rs     | 61 +++++++++++++++--------
 tests/ui/attributes/inner-attr-metavar.rs | 14 ++++++
 2 files changed, 53 insertions(+), 22 deletions(-)
 create mode 100644 tests/ui/attributes/inner-attr-metavar.rs

diff --git a/compiler/rustc_ast/src/tokenstream.rs b/compiler/rustc_ast/src/tokenstream.rs
index 161f8e67551f..fb331e74aeb8 100644
--- a/compiler/rustc_ast/src/tokenstream.rs
+++ b/compiler/rustc_ast/src/tokenstream.rs
@@ -233,35 +233,52 @@ fn attrs_and_tokens_to_token_trees(
 
     // Insert inner attribute tokens.
     if !inner_attrs.is_empty() {
-        let mut found = false;
-        // Check the last two trees (to account for a trailing semi)
-        for tree in res.iter_mut().rev().take(2) {
-            if let TokenTree::Delimited(span, spacing, delim, delim_tokens) = tree {
-                // Inner attributes are only supported on extern blocks, functions,
-                // impls, and modules. All of these have their inner attributes
-                // placed at the beginning of the rightmost outermost braced group:
-                // e.g. fn foo() { #![my_attr] }
-                //
-                // Therefore, we can insert them back into the right location
-                // without needing to do any extra position tracking.
-                //
-                // Note: Outline modules are an exception - they can
-                // have attributes like `#![my_attr]` at the start of a file.
-                // Support for custom attributes in this position is not
-                // properly implemented - we always synthesize fake tokens,
-                // so we never reach this code.
+        let found = insert_inner_attrs(inner_attrs, res);
+        assert!(found, "Failed to find trailing delimited group in: {res:?}");
+    }
+
+    // Inner attributes are only supported on blocks, functions, impls, and
+    // modules. All of these have their inner attributes placed at the
+    // beginning of the rightmost outermost braced group:
+    // e.g. `fn foo() { #![my_attr] }`. (Note: the braces may be within
+    // invisible delimiters.)
+    //
+    // Therefore, we can insert them back into the right location without
+    // needing to do any extra position tracking.
+    //
+    // Note: Outline modules are an exception - they can have attributes like
+    // `#![my_attr]` at the start of a file. Support for custom attributes in
+    // this position is not properly implemented - we always synthesize fake
+    // tokens, so we never reach this code.
+    fn insert_inner_attrs(inner_attrs: &[Attribute], tts: &mut Vec<TokenTree>) -> bool {
+        for tree in tts.iter_mut().rev() {
+            if let TokenTree::Delimited(span, spacing, Delimiter::Brace, stream) = tree {
+                // Found it: the rightmost, outermost braced group.
                 let mut tts = vec![];
                 for inner_attr in inner_attrs {
                     tts.extend(inner_attr.token_trees());
                 }
-                tts.extend(delim_tokens.0.iter().cloned());
+                tts.extend(stream.0.iter().cloned());
                 let stream = TokenStream::new(tts);
-                *tree = TokenTree::Delimited(*span, *spacing, *delim, stream);
-                found = true;
-                break;
+                *tree = TokenTree::Delimited(*span, *spacing, Delimiter::Brace, stream);
+                return true;
+            } else if let TokenTree::Delimited(span, spacing, Delimiter::Invisible(src), stream) =
+                tree
+            {
+                // Recurse inside invisible delimiters.
+                let mut vec: Vec<_> = stream.iter().cloned().collect();
+                if insert_inner_attrs(inner_attrs, &mut vec) {
+                    *tree = TokenTree::Delimited(
+                        *span,
+                        *spacing,
+                        Delimiter::Invisible(*src),
+                        TokenStream::new(vec),
+                    );
+                    return true;
+                }
             }
         }
-        assert!(found, "Failed to find trailing delimited group in: {res:?}");
+        false
     }
 }
 
diff --git a/tests/ui/attributes/inner-attr-metavar.rs b/tests/ui/attributes/inner-attr-metavar.rs
new file mode 100644
index 000000000000..345121c252e8
--- /dev/null
+++ b/tests/ui/attributes/inner-attr-metavar.rs
@@ -0,0 +1,14 @@
+//@ check-pass
+//
+// During `Nonterminal` removal (#124141) there was at one point a problem with
+// calling from_ast on expressions with inner attributes within metavars -- the
+// inner attributes were being inserted in the wrong place in `from_ast`. This
+// test covers that case.
+
+macro_rules! m3 { ($e:expr) => {} }
+macro_rules! m2 { ($e:expr) => { m3!($e); } }
+macro_rules! m1 { ($e:expr) => { m2!($e); } }
+
+m1!({ #![allow(unused)] 0 });
+
+fn main() {}