Rollup merge of #140420 - fmease:rustdoc-fix-doctest-heur, r=GuillaumeGomez

rustdoc: Fix doctest heuristic for main fn wrapping

Fixes #140412 which regressed in #140220 that I reviewed. As mentioned in https://github.com/rust-lang/rust/pull/140220#issuecomment-2837061779, at the time I didn't have the time to re-review its latest changes and should've therefore invalided my previous "r=me" and blocked the PR on another review given the fragile nature of the doctest impl. This didn't happen which is my fault.

Contains some other small changes. Diff best reviewed modulo whitespace.
r? ``@GuillaumeGomez``

12 files changed, 182 insertions, 88 deletions

diff --git a/src/librustdoc/doctest/make.rs b/src/librustdoc/doctest/make.rs
index 4194abc8d57..d4fbfb12582 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,41 @@ 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;
+                                    break;
                                 }
-                            {
-                                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/auxiliary/items.rs b/tests/rustdoc-ui/doctest/auxiliary/items.rs
new file mode 100644
index 00000000000..40d4eb261e5
--- /dev/null
+++ b/tests/rustdoc-ui/doctest/auxiliary/items.rs
@@ -0,0 +1 @@
+fn item() {}
diff --git a/tests/rustdoc-ui/doctest/auxiliary/macro-after-main.rs b/tests/rustdoc-ui/doctest/auxiliary/macro-after-main.rs
deleted file mode 100644
index ed7584b7425..00000000000
--- a/tests/rustdoc-ui/doctest/auxiliary/macro-after-main.rs
+++ /dev/null
@@ -1 +0,0 @@
-use std::string::String;
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 0a42343f1c2..00000000000
--- 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 72ffe2b5a27..00000000000
--- 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
-
diff --git a/tests/rustdoc-ui/doctest/main-alongside-macro-calls.fail.stdout b/tests/rustdoc-ui/doctest/main-alongside-macro-calls.fail.stdout
new file mode 100644
index 00000000000..65989a8ef47
--- /dev/null
+++ b/tests/rustdoc-ui/doctest/main-alongside-macro-calls.fail.stdout
@@ -0,0 +1,60 @@
+
+running 4 tests
+test $DIR/main-alongside-macro-calls.rs - (line 19) ... ok
+test $DIR/main-alongside-macro-calls.rs - (line 24) ... ok
+test $DIR/main-alongside-macro-calls.rs - (line 28) ... FAILED
+test $DIR/main-alongside-macro-calls.rs - (line 33) ... FAILED
+
+failures:
+
+---- $DIR/main-alongside-macro-calls.rs - (line 28) stdout ----
+error: macros that expand to items must be delimited with braces or followed by a semicolon
+  --> $DIR/main-alongside-macro-calls.rs:30:1
+   |
+LL | println!();
+   | ^^^^^^^^^^
+   |
+   = note: this error originates in the macro `println` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error: macro expansion ignores `{` and any tokens following
+  --> $SRC_DIR/std/src/macros.rs:LL:COL
+   |
+  ::: $DIR/main-alongside-macro-calls.rs:30:1
+   |
+LL | println!();
+   | ---------- caused by the macro expansion here
+   |
+   = note: the usage of `print!` is likely invalid in item context
+
+error: aborting due to 2 previous errors
+
+Couldn't compile the test.
+---- $DIR/main-alongside-macro-calls.rs - (line 33) stdout ----
+error: macros that expand to items must be delimited with braces or followed by a semicolon
+  --> $DIR/main-alongside-macro-calls.rs:34:1
+   |
+LL | println!();
+   | ^^^^^^^^^^
+   |
+   = note: this error originates in the macro `println` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error: macro expansion ignores `{` and any tokens following
+  --> $SRC_DIR/std/src/macros.rs:LL:COL
+   |
+  ::: $DIR/main-alongside-macro-calls.rs:34:1
+   |
+LL | println!();
+   | ---------- caused by the macro expansion here
+   |
+   = note: the usage of `print!` is likely invalid in item context
+
+error: aborting due to 2 previous errors
+
+Couldn't compile the test.
+
+failures:
+    $DIR/main-alongside-macro-calls.rs - (line 28)
+    $DIR/main-alongside-macro-calls.rs - (line 33)
+
+test result: FAILED. 2 passed; 2 failed; 0 ignored; 0 measured; 0 filtered out; finished in $TIME
+
diff --git a/tests/rustdoc-ui/doctest/main-alongside-macro-calls.pass.stdout b/tests/rustdoc-ui/doctest/main-alongside-macro-calls.pass.stdout
new file mode 100644
index 00000000000..93a4bbd8736
--- /dev/null
+++ b/tests/rustdoc-ui/doctest/main-alongside-macro-calls.pass.stdout
@@ -0,0 +1,9 @@
+
+running 4 tests
+test $DIR/main-alongside-macro-calls.rs - (line 19) ... ok
+test $DIR/main-alongside-macro-calls.rs - (line 24) ... ok
+test $DIR/main-alongside-macro-calls.rs - (line 28) - compile fail ... ok
+test $DIR/main-alongside-macro-calls.rs - (line 33) - compile fail ... ok
+
+test result: ok. 4 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in $TIME
+
diff --git a/tests/rustdoc-ui/doctest/main-alongside-macro-calls.rs b/tests/rustdoc-ui/doctest/main-alongside-macro-calls.rs
new file mode 100644
index 00000000000..b455d8b0cc3
--- /dev/null
+++ b/tests/rustdoc-ui/doctest/main-alongside-macro-calls.rs
@@ -0,0 +1,44 @@
+// This test ensures that if there is are any macro calls alongside a `main` function,
+// it will indeed consider the `main` function as the program entry point and *won't*
+// generate its own `main` function to wrap everything even though macro calls are
+// valid in statement contexts, too, and could just as well expand to statements or
+// expressions (we don't perform any macro expansion to find `main`, see also
+// <https://github.com/rust-lang/rust/issues/57415>).
+//
+// See <./main-alongside-stmts.rs> for comparison.
+//
+//@ compile-flags:--test --test-args --test-threads=1
+//@ normalize-stdout: "tests/rustdoc-ui/doctest" -> "$$DIR"
+//@ normalize-stdout: "finished in \d+\.\d+s" -> "finished in $$TIME"
+//@ revisions: pass fail
+//@[pass] check-pass
+//@[fail] failure-status: 101
+
+// Regression test for <https://github.com/rust-lang/rust/pull/140220#issuecomment-2831872920>:
+
+//! ```
+//! fn main() {}
+//! include!("./auxiliary/items.rs");
+//! ```
+//!
+//! ```
+//! include!("./auxiliary/items.rs");
+//! fn main() {}
+//! ```
+
+// Regression test for <https://github.com/rust-lang/rust/issues/140412>:
+// We test the "same" thing twice: Once via `compile_fail` to more closely mirror the reported
+// regression and once without it to make sure that it leads to the expected rustc errors,
+// namely `println!(…)` not being valid in item contexts.
+
+#![cfg_attr(pass, doc = " ```compile_fail")]
+#![cfg_attr(fail, doc = " ```")]
+//! fn main() {}
+//! println!();
+//! ```
+//!
+#![cfg_attr(pass, doc = " ```compile_fail")]
+#![cfg_attr(fail, doc = " ```")]
+//! println!();
+//! fn main() {}
+//! ```
diff --git a/tests/rustdoc-ui/doctest/main-alongside-stmts.rs b/tests/rustdoc-ui/doctest/main-alongside-stmts.rs
new file mode 100644
index 00000000000..5965f928cdd
--- /dev/null
+++ b/tests/rustdoc-ui/doctest/main-alongside-stmts.rs
@@ -0,0 +1,33 @@
+// This test ensures that if there is are any statements alongside a `main` function,
+// it will not consider the `main` function as the program entry point but instead
+// will generate its own `main` function to wrap everything as it needs to reside in a
+// module where only *items* are permitted syntactically.
+//
+// See <./main-alongside-macro-calls.rs> for comparison.
+//
+// This is a regression test for:
+// * <https://github.com/rust-lang/rust/issues/140162>
+// * <https://github.com/rust-lang/rust/issues/139651>
+//
+//@ compile-flags:--test --test-args --test-threads=1
+//@ normalize-stdout: "tests/rustdoc-ui/doctest" -> "$$DIR"
+//@ normalize-stdout: "finished in \d+\.\d+s" -> "finished in $$TIME"
+//@ check-pass
+
+//! ```
+//! # if cfg!(miri) { return; }
+//! use std::ops::Deref;
+//!
+//! fn main() {
+//!     assert!(false);
+//! }
+//! ```
+//!
+//! ```
+//! let x = 2;
+//! assert_eq!(x, 2);
+//!
+//! fn main() {
+//!     assert!(false);
+//! }
+//! ```
diff --git a/tests/rustdoc-ui/doctest/main-alongside-stmts.stdout b/tests/rustdoc-ui/doctest/main-alongside-stmts.stdout
new file mode 100644
index 00000000000..9b9a3fe8a68
--- /dev/null
+++ b/tests/rustdoc-ui/doctest/main-alongside-stmts.stdout
@@ -0,0 +1,7 @@
+
+running 2 tests
+test $DIR/main-alongside-stmts.rs - (line 17) ... ok
+test $DIR/main-alongside-stmts.rs - (line 26) ... ok
+
+test result: ok. 2 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in $TIME
+
diff --git a/tests/rustdoc-ui/doctest/test-main-alongside-exprs.rs b/tests/rustdoc-ui/doctest/test-main-alongside-exprs.rs
deleted file mode 100644
index ee2299c0fd8..00000000000
--- a/tests/rustdoc-ui/doctest/test-main-alongside-exprs.rs
+++ /dev/null
@@ -1,22 +0,0 @@
-// This test ensures that if there is an expression alongside a `main`
-// function, it will not consider the entire code to be part of the `main`
-// function and will generate its own function to wrap everything.
-//
-// This is a regression test for:
-// * <https://github.com/rust-lang/rust/issues/140162>
-// * <https://github.com/rust-lang/rust/issues/139651>
-//@ compile-flags:--test
-//@ normalize-stdout: "tests/rustdoc-ui/doctest" -> "$$DIR"
-//@ normalize-stdout: "finished in \d+\.\d+s" -> "finished in $$TIME"
-//@ check-pass
-
-#![crate_name = "foo"]
-
-//! ```
-//! # if cfg!(miri) { return; }
-//! use std::ops::Deref;
-//!
-//! fn main() {
-//!     println!("Hi!");
-//! }
-//! ```
diff --git a/tests/rustdoc-ui/doctest/test-main-alongside-exprs.stdout b/tests/rustdoc-ui/doctest/test-main-alongside-exprs.stdout
deleted file mode 100644
index 90d7c3546bf..00000000000
--- a/tests/rustdoc-ui/doctest/test-main-alongside-exprs.stdout
+++ /dev/null
@@ -1,6 +0,0 @@
-
-running 1 test
-test $DIR/test-main-alongside-exprs.rs - (line 15) ... ok
-
-test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in $TIME
-