Parse full doctest source; extract helper for parsing code

It doesn't really make sense to skip part of the source when we're
parsing it, so parse the whole doctest. This simplifies things too.

1 files changed, 104 insertions, 90 deletions

diff --git a/src/librustdoc/doctest/make.rs b/src/librustdoc/doctest/make.rs
index 1eafad9cfc6..90efcbdd8b1 100644
--- a/src/librustdoc/doctest/make.rs
+++ b/src/librustdoc/doctest/make.rs
@@ -6,7 +6,7 @@ use std::io;
 use rustc_ast as ast;
 use rustc_data_structures::sync::Lrc;
 use rustc_errors::emitter::stderr_destination;
-use rustc_errors::ColorConfig;
+use rustc_errors::{ColorConfig, FatalError};
 use rustc_parse::new_parser_from_source_str;
 use rustc_parse::parser::attr::InnerAttrPolicy;
 use rustc_session::parse::ParseSess;
@@ -55,6 +55,95 @@ pub(crate) fn make_test(
 
     // Uses librustc_ast to parse the doctest and find if there's a main fn and the extern
     // crate already is included.
+    let Ok((already_has_main, already_has_extern_crate)) =
+        check_for_main_and_extern_crate(crate_name, s.to_owned(), edition, &mut supports_color)
+    else {
+        // If the parser panicked due to a fatal error, pass the test code through unchanged.
+        // The error will be reported during compilation.
+        return (s.to_owned(), 0, false);
+    };
+
+    // Don't inject `extern crate std` because it's already injected by the
+    // compiler.
+    if !already_has_extern_crate && !opts.no_crate_inject && crate_name != Some("std") {
+        if let Some(crate_name) = crate_name {
+            // Don't inject `extern crate` if the crate is never used.
+            // NOTE: this is terribly inaccurate because it doesn't actually
+            // parse the source, but only has false positives, not false
+            // negatives.
+            if s.contains(crate_name) {
+                // rustdoc implicitly inserts an `extern crate` item for the own crate
+                // which may be unused, so we need to allow the lint.
+                prog.push_str("#[allow(unused_extern_crates)]\n");
+
+                prog.push_str(&format!("extern crate r#{crate_name};\n"));
+                line_offset += 1;
+            }
+        }
+    }
+
+    // FIXME: This code cannot yet handle no_std test cases yet
+    if dont_insert_main || already_has_main || prog.contains("![no_std]") {
+        prog.push_str(everything_else);
+    } else {
+        let returns_result = everything_else.trim_end().ends_with("(())");
+        // Give each doctest main function a unique name.
+        // This is for example needed for the tooling around `-C instrument-coverage`.
+        let inner_fn_name = if let Some(test_id) = test_id {
+            format!("_doctest_main_{test_id}")
+        } else {
+            "_inner".into()
+        };
+        let inner_attr = if test_id.is_some() { "#[allow(non_snake_case)] " } else { "" };
+        let (main_pre, main_post) = if returns_result {
+            (
+                format!(
+                    "fn main() {{ {inner_attr}fn {inner_fn_name}() -> Result<(), impl core::fmt::Debug> {{\n",
+                ),
+                format!("\n}} {inner_fn_name}().unwrap() }}"),
+            )
+        } else if test_id.is_some() {
+            (
+                format!("fn main() {{ {inner_attr}fn {inner_fn_name}() {{\n",),
+                format!("\n}} {inner_fn_name}() }}"),
+            )
+        } else {
+            ("fn main() {\n".into(), "\n}".into())
+        };
+        // Note on newlines: We insert a line/newline *before*, and *after*
+        // the doctest and adjust the `line_offset` accordingly.
+        // In the case of `-C instrument-coverage`, this means that the generated
+        // inner `main` function spans from the doctest opening codeblock to the
+        // closing one. For example
+        // /// ``` <- start of the inner main
+        // /// <- code under doctest
+        // /// ``` <- end of the inner main
+        line_offset += 1;
+
+        // add extra 4 spaces for each line to offset the code block
+        let content = if opts.insert_indent_space {
+            everything_else
+                .lines()
+                .map(|line| format!("    {}", line))
+                .collect::<Vec<String>>()
+                .join("\n")
+        } else {
+            everything_else.to_string()
+        };
+        prog.extend([&main_pre, content.as_str(), &main_post].iter().cloned());
+    }
+
+    debug!("final doctest:\n{prog}");
+
+    (prog, line_offset, supports_color)
+}
+
+fn check_for_main_and_extern_crate(
+    crate_name: Option<&str>,
+    source: String,
+    edition: Edition,
+    supports_color: &mut bool,
+) -> Result<(bool, bool), FatalError> {
     let result = rustc_driver::catch_fatal_errors(|| {
         rustc_span::create_session_if_not_set_then(edition, |_| {
             use rustc_errors::emitter::{Emitter, HumanEmitter};
@@ -62,8 +151,7 @@ pub(crate) fn make_test(
             use rustc_parse::parser::ForceCollect;
             use rustc_span::source_map::FilePathMapping;
 
-            let filename = FileName::anon_source_code(s);
-            let source = crates + everything_else;
+            let filename = FileName::anon_source_code(&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.
@@ -72,7 +160,7 @@ pub(crate) fn make_test(
                 rustc_driver::DEFAULT_LOCALE_RESOURCES.to_vec(),
                 false,
             );
-            supports_color =
+            *supports_color =
                 HumanEmitter::new(stderr_destination(ColorConfig::Auto), fallback_bundle.clone())
                     .supports_color();
 
@@ -86,13 +174,14 @@ pub(crate) fn make_test(
             let mut found_extern_crate = crate_name.is_none();
             let mut found_macro = false;
 
-            let mut parser = match new_parser_from_source_str(&psess, filename, source) {
-                Ok(p) => p,
-                Err(errs) => {
-                    errs.into_iter().for_each(|err| err.cancel());
-                    return (found_main, found_extern_crate, found_macro);
-                }
-            };
+            let mut parser =
+                match new_parser_from_source_str(&psess, filename, source.clone()) {
+                    Ok(p) => p,
+                    Err(errs) => {
+                        errs.into_iter().for_each(|err| err.cancel());
+                        return (found_main, found_extern_crate, found_macro);
+                    }
+                };
 
             loop {
                 match parser.parse_item(ForceCollect::No) {
@@ -146,18 +235,15 @@ pub(crate) fn make_test(
             (found_main, found_extern_crate, found_macro)
         })
     });
-    let Ok((already_has_main, already_has_extern_crate, found_macro)) = result else {
-        // If the parser panicked due to a fatal error, pass the test code through unchanged.
-        // The error will be reported during compilation.
-        return (s.to_owned(), 0, false);
-    };
+    let (already_has_main, already_has_extern_crate, found_macro) = result?;
 
     // If a doctest's `fn main` is being masked by a wrapper macro, the parsing loop above won't
     // see it. In that case, run the old text-based scan to see if they at least have a main
     // function written inside a macro invocation. See
     // https://github.com/rust-lang/rust/issues/56898
     let already_has_main = if found_macro && !already_has_main {
-        s.lines()
+        source
+            .lines()
             .map(|line| {
                 let comment = line.find("//");
                 if let Some(comment_begins) = comment { &line[0..comment_begins] } else { line }
@@ -167,79 +253,7 @@ pub(crate) fn make_test(
         already_has_main
     };
 
-    // Don't inject `extern crate std` because it's already injected by the
-    // compiler.
-    if !already_has_extern_crate && !opts.no_crate_inject && crate_name != Some("std") {
-        if let Some(crate_name) = crate_name {
-            // Don't inject `extern crate` if the crate is never used.
-            // NOTE: this is terribly inaccurate because it doesn't actually
-            // parse the source, but only has false positives, not false
-            // negatives.
-            if s.contains(crate_name) {
-                // rustdoc implicitly inserts an `extern crate` item for the own crate
-                // which may be unused, so we need to allow the lint.
-                prog.push_str("#[allow(unused_extern_crates)]\n");
-
-                prog.push_str(&format!("extern crate r#{crate_name};\n"));
-                line_offset += 1;
-            }
-        }
-    }
-
-    // FIXME: This code cannot yet handle no_std test cases yet
-    if dont_insert_main || already_has_main || prog.contains("![no_std]") {
-        prog.push_str(everything_else);
-    } else {
-        let returns_result = everything_else.trim_end().ends_with("(())");
-        // Give each doctest main function a unique name.
-        // This is for example needed for the tooling around `-C instrument-coverage`.
-        let inner_fn_name = if let Some(test_id) = test_id {
-            format!("_doctest_main_{test_id}")
-        } else {
-            "_inner".into()
-        };
-        let inner_attr = if test_id.is_some() { "#[allow(non_snake_case)] " } else { "" };
-        let (main_pre, main_post) = if returns_result {
-            (
-                format!(
-                    "fn main() {{ {inner_attr}fn {inner_fn_name}() -> Result<(), impl core::fmt::Debug> {{\n",
-                ),
-                format!("\n}} {inner_fn_name}().unwrap() }}"),
-            )
-        } else if test_id.is_some() {
-            (
-                format!("fn main() {{ {inner_attr}fn {inner_fn_name}() {{\n",),
-                format!("\n}} {inner_fn_name}() }}"),
-            )
-        } else {
-            ("fn main() {\n".into(), "\n}".into())
-        };
-        // Note on newlines: We insert a line/newline *before*, and *after*
-        // the doctest and adjust the `line_offset` accordingly.
-        // In the case of `-C instrument-coverage`, this means that the generated
-        // inner `main` function spans from the doctest opening codeblock to the
-        // closing one. For example
-        // /// ``` <- start of the inner main
-        // /// <- code under doctest
-        // /// ``` <- end of the inner main
-        line_offset += 1;
-
-        // add extra 4 spaces for each line to offset the code block
-        let content = if opts.insert_indent_space {
-            everything_else
-                .lines()
-                .map(|line| format!("    {}", line))
-                .collect::<Vec<String>>()
-                .join("\n")
-        } else {
-            everything_else.to_string()
-        };
-        prog.extend([&main_pre, content.as_str(), &main_post].iter().cloned());
-    }
-
-    debug!("final doctest:\n{prog}");
-
-    (prog, line_offset, supports_color)
+    Ok((already_has_main, already_has_extern_crate))
 }
 
 fn check_if_attr_is_complete(source: &str, edition: Edition) -> bool {