From 123ea25542ba00e92bf6d19084cad6e7a24453f0 Mon Sep 17 00:00:00 2001
From: Guillaume Gomez <guillaume1.gomez@gmail.com>
Date: Fri, 7 Mar 2025 15:20:14 +0100
Subject: [PATCH] Correctly handle line comments in attributes and generate
 extern crates outside of wrapping function

---
 src/librustdoc/doctest/make.rs                | 50 +++++++++++++------
 src/librustdoc/doctest/tests.rs               | 30 +++++++++--
 tests/run-make/rustdoc-error-lines/rmake.rs   |  6 +--
 .../rustdoc-ui/doctest/display-output.stdout  |  6 +--
 tests/rustdoc-ui/doctest/extern-crate.rs      | 23 +++++++++
 tests/rustdoc-ui/doctest/extern-crate.stdout  |  6 +++
 tests/rustdoc/playground.rs                   |  2 +-
 7 files changed, 98 insertions(+), 25 deletions(-)
 create mode 100644 tests/rustdoc-ui/doctest/extern-crate.rs
 create mode 100644 tests/rustdoc-ui/doctest/extern-crate.stdout

diff --git a/src/librustdoc/doctest/make.rs b/src/librustdoc/doctest/make.rs
index 810f53636ce..0ed2d37f74c 100644
--- a/src/librustdoc/doctest/make.rs
+++ b/src/librustdoc/doctest/make.rs
@@ -176,9 +176,24 @@ impl DocTestBuilder {
 
         // Now push any outer attributes from the example, assuming they
         // are intended to be crate attributes.
-        prog.push_str(&self.crate_attrs);
-        prog.push_str(&self.maybe_crate_attrs);
-        prog.push_str(&self.crates);
+        if !self.crate_attrs.is_empty() {
+            prog.push_str(&self.crate_attrs);
+            if !self.crate_attrs.ends_with('\n') {
+                prog.push('\n');
+            }
+        }
+        if !self.maybe_crate_attrs.is_empty() {
+            prog.push_str(&self.maybe_crate_attrs);
+            if !self.maybe_crate_attrs.ends_with('\n') {
+                prog.push('\n');
+            }
+        }
+        if !self.crates.is_empty() {
+            prog.push_str(&self.crates);
+            if !self.crates.ends_with('\n') {
+                prog.push('\n');
+            }
+        }
 
         // Don't inject `extern crate std` because it's already injected by the
         // compiler.
@@ -276,7 +291,6 @@ const DOCTEST_CODE_WRAPPER: &str = "fn f(){";
 fn parse_source(source: &str, crate_name: &Option<&str>) -> Result<ParseSourceInfo, ()> {
     use rustc_errors::DiagCtxt;
     use rustc_errors::emitter::{Emitter, HumanEmitter};
-    // use rustc_parse::parser::ForceCollect;
     use rustc_span::source_map::FilePathMapping;
 
     let mut info =
@@ -338,7 +352,8 @@ fn parse_source(source: &str, crate_name: &Option<&str>) -> Result<ParseSourceIn
         info: &mut ParseSourceInfo,
         crate_name: &Option<&str>,
         is_top_level: bool,
-    ) {
+    ) -> bool {
+        let mut is_crate = false;
         if !info.has_global_allocator
             && item.attrs.iter().any(|attr| attr.name_or_empty() == sym::global_allocator)
         {
@@ -354,12 +369,13 @@ fn parse_source(source: &str, crate_name: &Option<&str>) -> Result<ParseSourceIn
                 if let Some(ref body) = fn_item.body {
                     for stmt in &body.stmts {
                         if let StmtKind::Item(ref item) = stmt.kind {
-                            check_item(item, info, crate_name, false)
+                            is_crate |= check_item(item, info, crate_name, false);
                         }
                     }
                 }
             }
             ast::ItemKind::ExternCrate(original) => {
+                is_crate = true;
                 if !info.found_extern_crate
                     && let Some(crate_name) = crate_name
                 {
@@ -374,6 +390,7 @@ fn parse_source(source: &str, crate_name: &Option<&str>) -> Result<ParseSourceIn
             }
             _ => {}
         }
+        is_crate
     }
 
     let mut prev_span_hi = None;
@@ -412,8 +429,11 @@ fn parse_source(source: &str, crate_name: &Option<&str>) -> Result<ParseSourceIn
                 }
             }
             for stmt in &body.stmts {
+                let mut is_crate = false;
                 match stmt.kind {
-                    StmtKind::Item(ref item) => check_item(&item, &mut info, crate_name, true),
+                    StmtKind::Item(ref item) => {
+                        is_crate = check_item(&item, &mut info, crate_name, true);
+                    }
                     StmtKind::Expr(ref expr) if matches!(expr.kind, ast::ExprKind::Err(_)) => {
                         cancel_error_count(&psess);
                         return Err(());
@@ -450,15 +470,15 @@ fn parse_source(source: &str, crate_name: &Option<&str>) -> Result<ParseSourceIn
                 if info.everything_else.is_empty()
                     && (!info.maybe_crate_attrs.is_empty() || !info.crate_attrs.is_empty())
                 {
-                    // We add potential backlines/comments into attributes if there are some.
-                    push_to_s(
-                        &mut info.maybe_crate_attrs,
-                        source,
-                        span.shrink_to_lo(),
-                        &mut prev_span_hi,
-                    );
+                    // We add potential backlines/comments if there are some in items generated
+                    // before the wrapping function.
+                    push_to_s(&mut info.crates, source, span.shrink_to_lo(), &mut prev_span_hi);
+                }
+                if !is_crate {
+                    push_to_s(&mut info.everything_else, source, span, &mut prev_span_hi);
+                } else {
+                    push_to_s(&mut info.crates, source, span, &mut prev_span_hi);
                 }
-                push_to_s(&mut info.everything_else, source, span, &mut prev_span_hi);
             }
             Ok(info)
         }
diff --git a/src/librustdoc/doctest/tests.rs b/src/librustdoc/doctest/tests.rs
index 2a63d005082..59cc33558ed 100644
--- a/src/librustdoc/doctest/tests.rs
+++ b/src/librustdoc/doctest/tests.rs
@@ -127,8 +127,8 @@ fn make_test_manual_extern_crate() {
 use asdf::qwop;
 assert_eq!(2+2, 4);";
     let expected = "#![allow(unused)]
-fn main() {
 extern crate asdf;
+fn main() {
 use asdf::qwop;
 assert_eq!(2+2, 4);
 }"
@@ -144,8 +144,8 @@ fn make_test_manual_extern_crate_with_macro_use() {
 use asdf::qwop;
 assert_eq!(2+2, 4);";
     let expected = "#![allow(unused)]
-fn main() {
 #[macro_use] extern crate asdf;
+fn main() {
 use asdf::qwop;
 assert_eq!(2+2, 4);
 }"
@@ -197,6 +197,7 @@ fn make_test_crate_attrs() {
 assert_eq!(2+2, 4);";
     let expected = "#![allow(unused)]
 #![feature(sick_rad)]
+
 fn main() {
 assert_eq!(2+2, 4);
 }"
@@ -277,10 +278,10 @@ fn make_test_issues_33731() {
 assert_eq!(asdf::foo, 4);";
 
     let expected = "#![allow(unused)]
+extern crate hella_qwop;
 #[allow(unused_extern_crates)]
 extern crate r#asdf;
 fn main() {
-extern crate hella_qwop;
 assert_eq!(asdf::foo, 4);
 }"
     .to_string();
@@ -401,3 +402,26 @@ fn check_split_args() {
     compare("a\n\t \rb", &["a", "b"]);
     compare("a\n\t1 \rb", &["a", "1", "b"]);
 }
+
+#[test]
+fn comment_in_attrs() {
+    // if input already has a fn main, it should insert a space before it
+    let opts = default_global_opts("");
+    let input = "\
+#![feature(rustdoc_internals)]
+#![allow(internal_features)]
+#![doc(rust_logo)]
+//! This crate has the Rust(tm) branding on it.";
+    let expected = "\
+#![allow(unused)]
+#![feature(rustdoc_internals)]
+#![allow(internal_features)]
+#![doc(rust_logo)]
+//! This crate has the Rust(tm) branding on it.
+fn main() {
+
+}"
+    .to_string();
+    let (output, len) = make_test(input, None, false, &opts, None);
+    assert_eq!((output, len), (expected, 2));
+}
diff --git a/tests/run-make/rustdoc-error-lines/rmake.rs b/tests/run-make/rustdoc-error-lines/rmake.rs
index ea5ec2faed9..0d8c500ed1e 100644
--- a/tests/run-make/rustdoc-error-lines/rmake.rs
+++ b/tests/run-make/rustdoc-error-lines/rmake.rs
@@ -8,11 +8,11 @@ fn main() {
 
     let should_contain = &[
         "input.rs - foo (line 5)",
-        "input.rs:7:15",
+        "input.rs:8:15",
         "input.rs - bar (line 13)",
-        "input.rs:15:15",
+        "input.rs:16:15",
         "input.rs - bar (line 22)",
-        "input.rs:24:15",
+        "input.rs:25:15",
     ];
     for text in should_contain {
         assert!(output.contains(text), "output doesn't contains {:?}", text);
diff --git a/tests/rustdoc-ui/doctest/display-output.stdout b/tests/rustdoc-ui/doctest/display-output.stdout
index ad25d1ce541..45e107b2c70 100644
--- a/tests/rustdoc-ui/doctest/display-output.stdout
+++ b/tests/rustdoc-ui/doctest/display-output.stdout
@@ -6,7 +6,7 @@ successes:
 
 ---- $DIR/display-output.rs - foo (line 9) stdout ----
 warning: unused variable: `x`
-  --> $DIR/display-output.rs:11:5
+  --> $DIR/display-output.rs:12:5
    |
 LL | let x = 12;
    |     ^ help: if this is intentional, prefix it with an underscore: `_x`
@@ -19,13 +19,13 @@ LL | #![warn(unused)]
    = note: `#[warn(unused_variables)]` implied by `#[warn(unused)]`
 
 warning: unused variable: `x`
-  --> $DIR/display-output.rs:13:8
+  --> $DIR/display-output.rs:14:8
    |
 LL | fn foo(x: &dyn std::fmt::Display) {}
    |        ^ help: if this is intentional, prefix it with an underscore: `_x`
 
 warning: function `foo` is never used
-  --> $DIR/display-output.rs:13:4
+  --> $DIR/display-output.rs:14:4
    |
 LL | fn foo(x: &dyn std::fmt::Display) {}
    |    ^^^
diff --git a/tests/rustdoc-ui/doctest/extern-crate.rs b/tests/rustdoc-ui/doctest/extern-crate.rs
new file mode 100644
index 00000000000..0415d33bb72
--- /dev/null
+++ b/tests/rustdoc-ui/doctest/extern-crate.rs
@@ -0,0 +1,23 @@
+//@ check-pass
+//@ 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"
+
+// This test ensures that crate imports are placed outside of the `main` function
+// so they work all the time (even in 2015 edition).
+
+/// ```rust
+/// #![feature(test)]
+///
+/// extern crate test;
+/// use test::Bencher;
+///
+/// #[bench]
+/// fn bench_xor_1000_ints(b: &mut Bencher) {
+///     b.iter(|| {
+///         (0..1000).fold(0, |old, new| old ^ new);
+///     });
+/// }
+/// ```
+///
+pub fn foo() {}
diff --git a/tests/rustdoc-ui/doctest/extern-crate.stdout b/tests/rustdoc-ui/doctest/extern-crate.stdout
new file mode 100644
index 00000000000..b103343afdd
--- /dev/null
+++ b/tests/rustdoc-ui/doctest/extern-crate.stdout
@@ -0,0 +1,6 @@
+
+running 1 test
+test $DIR/extern-crate.rs - foo (line 9) ... ok
+
+test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in $TIME
+
diff --git a/tests/rustdoc/playground.rs b/tests/rustdoc/playground.rs
index db2d1669df6..65dad2a5195 100644
--- a/tests/rustdoc/playground.rs
+++ b/tests/rustdoc/playground.rs
@@ -24,4 +24,4 @@
 
 //@ matches foo/index.html '//a[@class="test-arrow"][@href="https://www.example.com/?code=%23!%5Ballow(unused)%5D%0Afn+main()+%7B%0A++++println!(%22Hello,+world!%22);%0A%7D&edition=2015"]' ""
 //@ matches foo/index.html '//a[@class="test-arrow"][@href="https://www.example.com/?code=%23!%5Ballow(unused)%5D%0Afn+main()+%7B%0A++++println!(%22Hello,+world!%22);%0A%7D&edition=2015"]' ""
-//@ matches foo/index.html '//a[@class="test-arrow"][@href="https://www.example.com/?code=%23!%5Ballow(unused)%5D%0A%23!%5Bfeature(something)%5D%0A%0Afn+main()+%7B%0A++++println!(%22Hello,+world!%22);%0A%7D&version=nightly&edition=2015"]' ""
+//@ matches foo/index.html '//a[@class="test-arrow"][@href="https://www.example.com/?code=%23!%5Ballow(unused)%5D%0A%23!%5Bfeature(something)%5D%0A%0A%0Afn+main()+%7B%0A++++println!(%22Hello,+world!%22);%0A%7D&version=nightly&edition=2015"]' ""