From 714f6f2b065eb105a3df12e2398a4723b18a6473 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Le=C3=B3n=20Orell=20Valerian=20Liehr?= <me@fmease.dev>
Date: Tue, 29 Apr 2025 02:11:41 +0200
Subject: [PATCH] rustdoc: Fix doctest heuristic for main fn wrapping

---
 src/librustdoc/doctest/make.rs                | 64 ++++++++-----------
 tests/rustdoc-ui/doctest/macro-after-main.rs  | 16 -----
 .../doctest/macro-after-main.stdout           |  6 --
 3 files changed, 27 insertions(+), 59 deletions(-)
 delete mode 100644 tests/rustdoc-ui/doctest/macro-after-main.rs
 delete mode 100644 tests/rustdoc-ui/doctest/macro-after-main.stdout

diff --git a/src/librustdoc/doctest/make.rs b/src/librustdoc/doctest/make.rs
index 4194abc8d5742..2b8ed6f0e426e 100644
--- a/src/librustdoc/doctest/make.rs
+++ b/src/librustdoc/doctest/make.rs
@@ -301,8 +301,6 @@ fn parse_source(source: &str, crate_name: &Option<&str>) -> Result<ParseSourceIn
 
     let filename = FileName::anon_source_code(&wrapped_source);
 
-    // Any errors in parsing should also appear when the doctest is compiled for real, so just
-    // send all the errors that librustc_ast emits directly into a `Sink` instead of stderr.
     let sm = Arc::new(SourceMap::new(FilePathMapping::empty()));
     let fallback_bundle = rustc_errors::fallback_fluent_bundle(
         rustc_driver::DEFAULT_LOCALE_RESOURCES.to_vec(),
@@ -311,7 +309,8 @@ fn parse_source(source: &str, crate_name: &Option<&str>) -> Result<ParseSourceIn
     info.supports_color =
         HumanEmitter::new(stderr_destination(ColorConfig::Auto), fallback_bundle.clone())
             .supports_color();
-
+    // Any errors in parsing should also appear when the doctest is compiled for real, so just
+    // send all the errors that the parser emits directly into a `Sink` instead of stderr.
     let emitter = HumanEmitter::new(Box::new(io::sink()), fallback_bundle);
 
     // FIXME(misdreavus): pass `-Z treat-err-as-bug` to the doctest parser
@@ -339,9 +338,6 @@ fn parse_source(source: &str, crate_name: &Option<&str>) -> Result<ParseSourceIn
         *prev_span_hi = hi;
     }
 
-    // Recurse through functions body. It is necessary because the doctest source code is
-    // wrapped in a function to limit the number of AST errors. If we don't recurse into
-    // functions, we would thing all top-level items (so basically nothing).
     fn check_item(item: &ast::Item, info: &mut ParseSourceInfo, crate_name: &Option<&str>) -> bool {
         let mut is_extern_crate = false;
         if !info.has_global_allocator
@@ -351,8 +347,6 @@ fn parse_source(source: &str, crate_name: &Option<&str>) -> Result<ParseSourceIn
         }
         match item.kind {
             ast::ItemKind::Fn(ref fn_item) if !info.has_main_fn => {
-                // We only push if it's the top item because otherwise, we would duplicate
-                // its content since the top-level item was already added.
                 if fn_item.ident.name == sym::main {
                     info.has_main_fn = true;
                 }
@@ -412,44 +406,40 @@ fn parse_source(source: &str, crate_name: &Option<&str>) -> Result<ParseSourceIn
                 let mut is_extern_crate = false;
                 match stmt.kind {
                     StmtKind::Item(ref item) => {
-                        is_extern_crate = check_item(&item, &mut info, crate_name);
-                    }
-                    StmtKind::Expr(ref expr) => {
-                        if matches!(expr.kind, ast::ExprKind::Err(_)) {
-                            reset_error_count(&psess);
-                            return Err(());
-                        }
-                        has_non_items = true;
+                        is_extern_crate = check_item(item, &mut info, crate_name);
                     }
                     // We assume that the macro calls will expand to item(s) even though they could
-                    // expand to statements and expressions. And the simple fact that we're trying
-                    // to retrieve a `main` function inside it is a terrible idea.
+                    // expand to statements and expressions.
                     StmtKind::MacCall(ref mac_call) => {
-                        if info.has_main_fn {
-                            continue;
-                        }
-                        let mut iter = mac_call.mac.args.tokens.iter();
-
-                        while let Some(token) = iter.next() {
-                            if let TokenTree::Token(token, _) = token
-                                && let TokenKind::Ident(name, _) = token.kind
-                                && name == kw::Fn
-                                && let Some(TokenTree::Token(fn_token, _)) = iter.peek()
-                                && let TokenKind::Ident(fn_name, _) = fn_token.kind
-                                && fn_name == sym::main
-                                && let Some(TokenTree::Delimited(_, _, Delimiter::Parenthesis, _)) = {
-                                    iter.next();
-                                    iter.peek()
+                        if !info.has_main_fn {
+                            // For backward compatibility, we look for the token sequence `fn main(…)`
+                            // in the macro input (!) to crudely detect main functions "masked by a
+                            // wrapper macro". For the record, this is a horrible heuristic!
+                            // See <https://github.com/rust-lang/rust/issues/56898>.
+                            let mut iter = mac_call.mac.args.tokens.iter();
+                            while let Some(token) = iter.next() {
+                                if let TokenTree::Token(token, _) = token
+                                    && let TokenKind::Ident(kw::Fn, _) = token.kind
+                                    && let Some(TokenTree::Token(ident, _)) = iter.peek()
+                                    && let TokenKind::Ident(sym::main, _) = ident.kind
+                                    && let Some(TokenTree::Delimited(.., Delimiter::Parenthesis, _)) = {
+                                        iter.next();
+                                        iter.peek()
+                                    }
+                                {
+                                    info.has_main_fn = true;
                                 }
-                            {
-                                info.has_main_fn = true;
-                                break;
                             }
                         }
                     }
-                    _ => {
+                    StmtKind::Expr(ref expr) => {
+                        if matches!(expr.kind, ast::ExprKind::Err(_)) {
+                            reset_error_count(&psess);
+                            return Err(());
+                        }
                         has_non_items = true;
                     }
+                    StmtKind::Let(_) | StmtKind::Semi(_) | StmtKind::Empty => has_non_items = true,
                 }
 
                 // Weirdly enough, the `Stmt` span doesn't include its attributes, so we need to
diff --git a/tests/rustdoc-ui/doctest/macro-after-main.rs b/tests/rustdoc-ui/doctest/macro-after-main.rs
deleted file mode 100644
index 0a42343f1c272..0000000000000
--- a/tests/rustdoc-ui/doctest/macro-after-main.rs
+++ /dev/null
@@ -1,16 +0,0 @@
-// This test checks a corner case where the macro calls used to be skipped,
-// making them considered as statement, and therefore some cases where
-// `include!` macro was then put into a function body, making the doctest
-// compilation fail.
-
-//@ compile-flags:--test
-//@ normalize-stdout: "tests/rustdoc-ui/doctest" -> "$$DIR"
-//@ normalize-stdout: "finished in \d+\.\d+s" -> "finished in $$TIME"
-//@ check-pass
-
-//! ```
-//! include!("./auxiliary/macro-after-main.rs");
-//!
-//! fn main() {}
-//! eprintln!();
-//! ```
diff --git a/tests/rustdoc-ui/doctest/macro-after-main.stdout b/tests/rustdoc-ui/doctest/macro-after-main.stdout
deleted file mode 100644
index 72ffe2b5a27c5..0000000000000
--- a/tests/rustdoc-ui/doctest/macro-after-main.stdout
+++ /dev/null
@@ -1,6 +0,0 @@
-
-running 1 test
-test $DIR/macro-after-main.rs - (line 11) ... ok
-
-test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in $TIME
-