Allow skipping extra paren insertion during AST pretty-printing

Fixes #74616
Makes progress towards #43081
Unblocks PR #76130

When pretty-printing an AST node, we may insert additional parenthesis
to ensure that precedence is properly preserved in code we output.
However, the proc macro implementation relies on comparing a
pretty-printed AST node to the captured `TokenStream`. Inserting extra
parenthesis changes the structure of the reparsed `TokenStream`, making
the comparison fail.

This PR refactors the AST pretty-printing code to allow skipping the
insertion of additional parenthesis. Several freestanding methods are
moved to trait methods on `PrintState`, which keep track of an internal
`insert_extra_parens` flag. This flag is normally `true`, but we expose
a public method which allows pretty-printing a nonterminal with
`insert_extra_parens = false`.

To avoid changing the public interface of `rustc_ast_pretty`, the
freestanding `_to_string` methods are changed to delegate to a
newly-crated `State`. The main pretty-printing code is moved to a new
`state` module to ensure that it does not accidentally call any of these
public helper functions (instead, the internal functions with the same
name should be used).

6 files changed, 219 insertions, 12 deletions

diff --git a/compiler/rustc_ast_pretty/src/pprust/mod.rs b/compiler/rustc_ast_pretty/src/pprust/mod.rs
index b88699f6ee1..b34ea41ab55 100644
--- a/compiler/rustc_ast_pretty/src/pprust/mod.rs
+++ b/compiler/rustc_ast_pretty/src/pprust/mod.rs
@@ -8,6 +8,11 @@ use rustc_ast as ast;
 use rustc_ast::token::{Nonterminal, Token, TokenKind};
 use rustc_ast::tokenstream::{TokenStream, TokenTree};
 
+pub fn nonterminal_to_string_no_extra_parens(nt: &Nonterminal) -> String {
+    let state = State::without_insert_extra_parens();
+    state.nonterminal_to_string(nt)
+}
+
 pub fn nonterminal_to_string(nt: &Nonterminal) -> String {
     State::new().nonterminal_to_string(nt)
 }
diff --git a/compiler/rustc_ast_pretty/src/pprust/state.rs b/compiler/rustc_ast_pretty/src/pprust/state.rs
index 14feb498903..9aa066370bb 100644
--- a/compiler/rustc_ast_pretty/src/pprust/state.rs
+++ b/compiler/rustc_ast_pretty/src/pprust/state.rs
@@ -20,9 +20,6 @@ use rustc_span::{BytePos, FileName, Span};
 
 use std::borrow::Cow;
 
-#[cfg(test)]
-mod tests;
-
 pub enum MacHeader<'a> {
     Path(&'a ast::Path),
     Keyword(&'static str),
@@ -91,6 +88,13 @@ pub struct State<'a> {
     comments: Option<Comments<'a>>,
     ann: &'a (dyn PpAnn + 'a),
     is_expanded: bool,
+    // If `true`, additional parenthesis (separate from `ExprKind::Paren`)
+    // are inserted to ensure that proper precedence is preserved
+    // in the pretty-printed output.
+    //
+    // This is usually `true`, except when performing the pretty-print/reparse
+    // check in `nt_to_tokenstream`
+    insert_extra_parens: bool,
 }
 
 crate const INDENT_UNIT: usize = 4;
@@ -112,6 +116,7 @@ pub fn print_crate<'a>(
         comments: Some(Comments::new(sm, filename, input)),
         ann,
         is_expanded,
+        insert_extra_parens: true,
     };
 
     if is_expanded && has_injected_crate {
@@ -225,7 +230,7 @@ pub fn literal_to_string(lit: token::Lit) -> String {
 }
 
 fn visibility_qualified(vis: &ast::Visibility, s: &str) -> String {
-    format!("{}{}", State::new().to_string(|s| s.print_visibility(vis)), s) 
+    format!("{}{}", State::new().to_string(|s| s.print_visibility(vis)), s)
 }
 
 impl std::ops::Deref for State<'_> {
@@ -242,6 +247,7 @@ impl std::ops::DerefMut for State<'_> {
 }
 
 pub trait PrintState<'a>: std::ops::Deref<Target = pp::Printer> + std::ops::DerefMut {
+    fn insert_extra_parens(&self) -> bool;
     fn comments(&mut self) -> &mut Option<Comments<'a>>;
     fn print_ident(&mut self, ident: Ident);
     fn print_generic_args(&mut self, args: &ast::GenericArgs, colons_before_params: bool);
@@ -827,12 +833,16 @@ pub trait PrintState<'a>: std::ops::Deref<Target = pp::Printer> + std::ops::Dere
 
     fn to_string(&self, f: impl FnOnce(&mut State<'_>)) -> String {
         let mut printer = State::new();
+        printer.insert_extra_parens = self.insert_extra_parens();
         f(&mut printer);
         printer.s.eof()
     }
 }
 
 impl<'a> PrintState<'a> for State<'a> {
+    fn insert_extra_parens(&self) -> bool {
+        self.insert_extra_parens
+    }
     fn comments(&mut self) -> &mut Option<Comments<'a>> {
         &mut self.comments
     }
@@ -874,9 +884,14 @@ impl<'a> State<'a> {
             comments: None,
             ann: &NoAnn,
             is_expanded: false,
+            insert_extra_parens: true,
         }
     }
 
+    pub(super) fn without_insert_extra_parens() -> State<'a> {
+        State { insert_extra_parens: false, ..State::new() }
+    }
+
     // Synthesizes a comment that was not textually present in the original source
     // file.
     pub fn synth_comment(&mut self, text: String) {
@@ -1679,7 +1694,8 @@ impl<'a> State<'a> {
     }
 
     /// Prints `expr` or `(expr)` when `needs_par` holds.
-    fn print_expr_cond_paren(&mut self, expr: &ast::Expr, needs_par: bool) {
+    fn print_expr_cond_paren(&mut self, expr: &ast::Expr, mut needs_par: bool) {
+        needs_par &= self.insert_extra_parens;
         if needs_par {
             self.popen();
         }
diff --git a/compiler/rustc_hir_pretty/src/lib.rs b/compiler/rustc_hir_pretty/src/lib.rs
index f6e4b1fb418..72011f04d9a 100644
--- a/compiler/rustc_hir_pretty/src/lib.rs
+++ b/compiler/rustc_hir_pretty/src/lib.rs
@@ -141,6 +141,9 @@ impl std::ops::DerefMut for State<'_> {
 }
 
 impl<'a> PrintState<'a> for State<'a> {
+    fn insert_extra_parens(&self) -> bool {
+        true
+    }
     fn comments(&mut self) -> &mut Option<Comments<'a>> {
         &mut self.comments
     }
diff --git a/compiler/rustc_parse/src/lib.rs b/compiler/rustc_parse/src/lib.rs
index b68d36c9a8e..51038b7d3aa 100644
--- a/compiler/rustc_parse/src/lib.rs
+++ b/compiler/rustc_parse/src/lib.rs
@@ -297,7 +297,11 @@ pub fn nt_to_tokenstream(nt: &Nonterminal, sess: &ParseSess, span: Span) -> Toke
     };
 
     // FIXME(#43081): Avoid this pretty-print + reparse hack
-    let source = pprust::nonterminal_to_string(nt);
+    // Pretty-print the AST struct without inserting any parenthesis
+    // beyond those explicitly written by the user (e.g. `ExpnKind::Paren`).
+    // The resulting stream may have incorrect precedence, but it's only
+    // ever used for a comparison against the capture tokenstream.
+    let source = pprust::nonterminal_to_string_no_extra_parens(nt);
     let filename = FileName::macro_expansion_source_code(&source);
     let reparsed_tokens = parse_stream_from_source_str(filename, source, sess, Some(span));
 
@@ -325,9 +329,28 @@ pub fn nt_to_tokenstream(nt: &Nonterminal, sess: &ParseSess, span: Span) -> Toke
     // modifications, including adding/removing typically non-semantic
     // tokens such as extra braces and commas, don't happen.
     if let Some(tokens) = tokens {
+        // If the streams match, then the AST hasn't been modified. Return the captured
+        // `TokenStream`.
         if tokenstream_probably_equal_for_proc_macro(&tokens, &reparsed_tokens, sess) {
             return tokens;
         }
+
+        // The check failed. This time, we pretty-print the AST struct with parenthesis
+        // inserted to preserve precedence. This may cause `None`-delimiters in the captured
+        // token stream to match up with inserted parenthesis in the reparsed stream.
+        let source_with_parens = pprust::nonterminal_to_string(nt);
+        let filename_with_parens = FileName::macro_expansion_source_code(&source_with_parens);
+        let tokens_with_parens = parse_stream_from_source_str(
+            filename_with_parens,
+            source_with_parens,
+            sess,
+            Some(span),
+        );
+
+        if tokenstream_probably_equal_for_proc_macro(&tokens, &tokens_with_parens, sess) {
+            return tokens;
+        }
+
         info!(
             "cached tokens found, but they're not \"probably equal\", \
                 going with stringified version"
@@ -489,12 +512,12 @@ pub fn tokentree_probably_equal_for_proc_macro(
         (TokenTree::Token(token), TokenTree::Token(reparsed_token)) => {
             token_probably_equal_for_proc_macro(token, reparsed_token)
         }
-        (
-            TokenTree::Delimited(_, delim, tokens),
-            TokenTree::Delimited(_, reparsed_delim, reparsed_tokens),
-        ) => {
-            delim == reparsed_delim
-                && tokenstream_probably_equal_for_proc_macro(tokens, reparsed_tokens, sess)
+        (TokenTree::Delimited(_, delim, tts), TokenTree::Delimited(_, delim2, tts2)) => {
+            // `NoDelim` delimiters can appear in the captured tokenstream, but not
+            // in the reparsed tokenstream. Allow them to match with anything, so
+            // that we check if the two streams are structurally equivalent.
+            (delim == delim2 || *delim == DelimToken::NoDelim || *delim2 == DelimToken::NoDelim)
+                && tokenstream_probably_equal_for_proc_macro(&tts, &tts2, sess)
         }
         _ => false,
     }
diff --git a/src/test/ui/proc-macro/issue-75734-pp-paren.rs b/src/test/ui/proc-macro/issue-75734-pp-paren.rs
new file mode 100644
index 00000000000..faa93787d13
--- /dev/null
+++ b/src/test/ui/proc-macro/issue-75734-pp-paren.rs
@@ -0,0 +1,26 @@
+// Regression test for issue #75734
+// Ensures that we don't lose tokens when pretty-printing would
+// normally insert extra parentheses.
+
+// check-pass
+// aux-build:test-macros.rs
+// compile-flags: -Z span-debug
+
+#![no_std] // Don't load unnecessary hygiene information from std
+extern crate std;
+
+#[macro_use]
+extern crate test_macros;
+
+macro_rules! mul_2 {
+    ($val:expr) => {
+        print_bang!($val * 2);
+    };
+}
+
+
+#[print_attr]
+fn main() {
+    &|_: u8| {};
+    mul_2!(1 + 1);
+}
diff --git a/src/test/ui/proc-macro/issue-75734-pp-paren.stdout b/src/test/ui/proc-macro/issue-75734-pp-paren.stdout
new file mode 100644
index 00000000000..b33b85f1705
--- /dev/null
+++ b/src/test/ui/proc-macro/issue-75734-pp-paren.stdout
@@ -0,0 +1,134 @@
+PRINT-ATTR INPUT (DISPLAY): fn main() { & | _ : u8 | { } ; mul_2 ! (1 + 1) ; }
+PRINT-ATTR INPUT (DEBUG): TokenStream [
+    Ident {
+        ident: "fn",
+        span: $DIR/issue-75734-pp-paren.rs:23:1: 23:3 (#0),
+    },
+    Ident {
+        ident: "main",
+        span: $DIR/issue-75734-pp-paren.rs:23:4: 23:8 (#0),
+    },
+    Group {
+        delimiter: Parenthesis,
+        stream: TokenStream [],
+        span: $DIR/issue-75734-pp-paren.rs:23:8: 23:10 (#0),
+    },
+    Group {
+        delimiter: Brace,
+        stream: TokenStream [
+            Punct {
+                ch: '&',
+                spacing: Joint,
+                span: $DIR/issue-75734-pp-paren.rs:24:5: 24:6 (#0),
+            },
+            Punct {
+                ch: '|',
+                spacing: Alone,
+                span: $DIR/issue-75734-pp-paren.rs:24:6: 24:7 (#0),
+            },
+            Ident {
+                ident: "_",
+                span: $DIR/issue-75734-pp-paren.rs:24:7: 24:8 (#0),
+            },
+            Punct {
+                ch: ':',
+                spacing: Alone,
+                span: $DIR/issue-75734-pp-paren.rs:24:8: 24:9 (#0),
+            },
+            Ident {
+                ident: "u8",
+                span: $DIR/issue-75734-pp-paren.rs:24:10: 24:12 (#0),
+            },
+            Punct {
+                ch: '|',
+                spacing: Alone,
+                span: $DIR/issue-75734-pp-paren.rs:24:12: 24:13 (#0),
+            },
+            Group {
+                delimiter: Brace,
+                stream: TokenStream [],
+                span: $DIR/issue-75734-pp-paren.rs:24:14: 24:16 (#0),
+            },
+            Punct {
+                ch: ';',
+                spacing: Alone,
+                span: $DIR/issue-75734-pp-paren.rs:24:16: 24:17 (#0),
+            },
+            Ident {
+                ident: "mul_2",
+                span: $DIR/issue-75734-pp-paren.rs:25:5: 25:10 (#0),
+            },
+            Punct {
+                ch: '!',
+                spacing: Alone,
+                span: $DIR/issue-75734-pp-paren.rs:25:10: 25:11 (#0),
+            },
+            Group {
+                delimiter: Parenthesis,
+                stream: TokenStream [
+                    Literal {
+                        kind: Integer,
+                        symbol: "1",
+                        suffix: None,
+                        span: $DIR/issue-75734-pp-paren.rs:25:12: 25:13 (#0),
+                    },
+                    Punct {
+                        ch: '+',
+                        spacing: Alone,
+                        span: $DIR/issue-75734-pp-paren.rs:25:14: 25:15 (#0),
+                    },
+                    Literal {
+                        kind: Integer,
+                        symbol: "1",
+                        suffix: None,
+                        span: $DIR/issue-75734-pp-paren.rs:25:16: 25:17 (#0),
+                    },
+                ],
+                span: $DIR/issue-75734-pp-paren.rs:25:11: 25:18 (#0),
+            },
+            Punct {
+                ch: ';',
+                spacing: Alone,
+                span: $DIR/issue-75734-pp-paren.rs:25:18: 25:19 (#0),
+            },
+        ],
+        span: $DIR/issue-75734-pp-paren.rs:23:11: 26:2 (#0),
+    },
+]
+PRINT-BANG INPUT (DISPLAY): 1 + 1 * 2
+PRINT-BANG INPUT (DEBUG): TokenStream [
+    Group {
+        delimiter: None,
+        stream: TokenStream [
+            Literal {
+                kind: Integer,
+                symbol: "1",
+                suffix: None,
+                span: $DIR/issue-75734-pp-paren.rs:25:12: 25:13 (#0),
+            },
+            Punct {
+                ch: '+',
+                spacing: Alone,
+                span: $DIR/issue-75734-pp-paren.rs:25:14: 25:15 (#0),
+            },
+            Literal {
+                kind: Integer,
+                symbol: "1",
+                suffix: None,
+                span: $DIR/issue-75734-pp-paren.rs:25:16: 25:17 (#0),
+            },
+        ],
+        span: $DIR/issue-75734-pp-paren.rs:17:21: 17:25 (#7),
+    },
+    Punct {
+        ch: '*',
+        spacing: Alone,
+        span: $DIR/issue-75734-pp-paren.rs:17:26: 17:27 (#7),
+    },
+    Literal {
+        kind: Integer,
+        symbol: "2",
+        suffix: None,
+        span: $DIR/issue-75734-pp-paren.rs:17:28: 17:29 (#7),
+    },
+]