Rollup merge of #77153 - Aaron1011:fix/better-recursive-expand, r=petrochenkov

Fix recursive nonterminal expansion during pretty-print/reparse check

Makes progress towards #43081

In PR #73084, we started recursively expanded nonterminals during the
pretty-print/reparse check, allowing them to be properly compared
against the reparsed tokenstream.

Unfortunately, the recursive logic in that PR only handles the case
where a nonterminal appears inside a `TokenTree::Delimited`. If a
nonterminal appears directly in the expanded tokens of another
nonterminal, the inner nonterminal will not be expanded.

This PR fixes the recursive expansion of nonterminals, ensuring that
they are expanded wherever they occur.

3 files changed, 120 insertions, 23 deletions

diff --git a/compiler/rustc_parse/src/lib.rs b/compiler/rustc_parse/src/lib.rs
index d59dd4016a9..b68d36c9a8e 100644
--- a/compiler/rustc_parse/src/lib.rs
+++ b/compiler/rustc_parse/src/lib.rs
@@ -8,7 +8,7 @@
 
 use rustc_ast as ast;
 use rustc_ast::token::{self, Nonterminal, Token, TokenKind};
-use rustc_ast::tokenstream::{self, Spacing, TokenStream, TokenTree};
+use rustc_ast::tokenstream::{self, TokenStream, TokenTree};
 use rustc_ast_pretty::pprust;
 use rustc_data_structures::sync::Lrc;
 use rustc_errors::{Diagnostic, FatalError, Level, PResult};
@@ -435,31 +435,42 @@ pub fn tokenstream_probably_equal_for_proc_macro(
         token_trees.into_iter()
     }
 
-    let expand_nt = |tree: TokenTree| {
-        if let TokenTree::Token(Token { kind: TokenKind::Interpolated(nt), span }) = &tree {
-            // When checking tokenstreams for 'probable equality', we are comparing
-            // a captured (from parsing) `TokenStream` to a reparsed tokenstream.
-            // The reparsed Tokenstream will never have `None`-delimited groups,
-            // since they are only ever inserted as a result of macro expansion.
-            // Therefore, inserting a `None`-delimtied group here (when we
-            // convert a nested `Nonterminal` to a tokenstream) would cause
-            // a mismatch with the reparsed tokenstream.
-            //
-            // Note that we currently do not handle the case where the
-            // reparsed stream has a `Parenthesis`-delimited group
-            // inserted. This will cause a spurious mismatch:
-            // issue #75734 tracks resolving this.
-            nt_to_tokenstream(nt, sess, *span).into_trees()
-        } else {
-            TokenStream::new(vec![(tree, Spacing::Alone)]).into_trees()
-        }
-    };
+    fn expand_token(tree: TokenTree, sess: &ParseSess) -> impl Iterator<Item = TokenTree> {
+        // When checking tokenstreams for 'probable equality', we are comparing
+        // a captured (from parsing) `TokenStream` to a reparsed tokenstream.
+        // The reparsed Tokenstream will never have `None`-delimited groups,
+        // since they are only ever inserted as a result of macro expansion.
+        // Therefore, inserting a `None`-delimtied group here (when we
+        // convert a nested `Nonterminal` to a tokenstream) would cause
+        // a mismatch with the reparsed tokenstream.
+        //
+        // Note that we currently do not handle the case where the
+        // reparsed stream has a `Parenthesis`-delimited group
+        // inserted. This will cause a spurious mismatch:
+        // issue #75734 tracks resolving this.
+
+        let expanded: SmallVec<[_; 1]> =
+            if let TokenTree::Token(Token { kind: TokenKind::Interpolated(nt), span }) = &tree {
+                nt_to_tokenstream(nt, sess, *span)
+                    .into_trees()
+                    .flat_map(|t| expand_token(t, sess))
+                    .collect()
+            } else {
+                // Filter before and after breaking tokens,
+                // since we may want to ignore both glued and unglued tokens.
+                std::iter::once(tree)
+                    .filter(semantic_tree)
+                    .flat_map(break_tokens)
+                    .filter(semantic_tree)
+                    .collect()
+            };
+        expanded.into_iter()
+    }
 
     // Break tokens after we expand any nonterminals, so that we break tokens
     // that are produced as a result of nonterminal expansion.
-    let tokens = tokens.trees().filter(semantic_tree).flat_map(expand_nt).flat_map(break_tokens);
-    let reparsed_tokens =
-        reparsed_tokens.trees().filter(semantic_tree).flat_map(expand_nt).flat_map(break_tokens);
+    let tokens = tokens.trees().flat_map(|t| expand_token(t, sess));
+    let reparsed_tokens = reparsed_tokens.trees().flat_map(|t| expand_token(t, sess));
 
     tokens.eq_by(reparsed_tokens, |t, rt| tokentree_probably_equal_for_proc_macro(&t, &rt, sess))
 }
diff --git a/src/test/ui/proc-macro/nested-nonterminal-tokens.rs b/src/test/ui/proc-macro/nested-nonterminal-tokens.rs
new file mode 100644
index 00000000000..2f5af10a40a
--- /dev/null
+++ b/src/test/ui/proc-macro/nested-nonterminal-tokens.rs
@@ -0,0 +1,26 @@
+// check-pass
+// edition:2018
+// compile-flags: -Z span-debug
+// aux-build:test-macros.rs
+
+// Tests that we properly pass tokens to proc-macro when nested
+// nonterminals are involved.
+
+#![no_std] // Don't load unnecessary hygiene information from std
+extern crate std;
+
+#[macro_use]
+extern crate test_macros;
+
+
+macro_rules! wrap {
+    (first, $e:expr) => { wrap!(second, $e + 1) };
+    (second, $e:expr) => { wrap!(third, $e + 2) };
+    (third, $e:expr) => {
+        print_bang!($e + 3);
+    };
+}
+
+fn main() {
+    let _ = wrap!(first, 0);
+}
diff --git a/src/test/ui/proc-macro/nested-nonterminal-tokens.stdout b/src/test/ui/proc-macro/nested-nonterminal-tokens.stdout
new file mode 100644
index 00000000000..a3d24dd26fe
--- /dev/null
+++ b/src/test/ui/proc-macro/nested-nonterminal-tokens.stdout
@@ -0,0 +1,60 @@
+PRINT-BANG INPUT (DISPLAY): 0 + 1 + 2 + 3
+PRINT-BANG INPUT (DEBUG): TokenStream [
+    Group {
+        delimiter: None,
+        stream: TokenStream [
+            Group {
+                delimiter: None,
+                stream: TokenStream [
+                    Group {
+                        delimiter: None,
+                        stream: TokenStream [
+                            Literal {
+                                kind: Integer,
+                                symbol: "0",
+                                suffix: None,
+                                span: $DIR/nested-nonterminal-tokens.rs:25:26: 25:27 (#0),
+                            },
+                        ],
+                        span: $DIR/nested-nonterminal-tokens.rs:17:41: 17:43 (#4),
+                    },
+                    Punct {
+                        ch: '+',
+                        spacing: Alone,
+                        span: $DIR/nested-nonterminal-tokens.rs:17:44: 17:45 (#4),
+                    },
+                    Literal {
+                        kind: Integer,
+                        symbol: "1",
+                        suffix: None,
+                        span: $DIR/nested-nonterminal-tokens.rs:17:46: 17:47 (#4),
+                    },
+                ],
+                span: $DIR/nested-nonterminal-tokens.rs:18:41: 18:43 (#5),
+            },
+            Punct {
+                ch: '+',
+                spacing: Alone,
+                span: $DIR/nested-nonterminal-tokens.rs:18:44: 18:45 (#5),
+            },
+            Literal {
+                kind: Integer,
+                symbol: "2",
+                suffix: None,
+                span: $DIR/nested-nonterminal-tokens.rs:18:46: 18:47 (#5),
+            },
+        ],
+        span: $DIR/nested-nonterminal-tokens.rs:20:21: 20:23 (#6),
+    },
+    Punct {
+        ch: '+',
+        spacing: Alone,
+        span: $DIR/nested-nonterminal-tokens.rs:20:24: 20:25 (#6),
+    },
+    Literal {
+        kind: Integer,
+        symbol: "3",
+        suffix: None,
+        span: $DIR/nested-nonterminal-tokens.rs:20:26: 20:27 (#6),
+    },
+]