1//! Logic for transforming the raw code given by the user into something actually2//! runnable, e.g. by adding a `main` function if it doesn't already exist.34use std::fmt::{self, Write as _};5use std::io;6use std::sync::Arc;78use rustc_ast::token::{Delimiter, TokenKind};9use rustc_ast::tokenstream::TokenTree;10use rustc_ast::{self as ast, AttrStyle, HasAttrs, StmtKind};11use rustc_errors::emitter::get_stderr_color_choice;12use rustc_errors::{AutoStream, ColorChoice, ColorConfig, DiagCtxtHandle};13use rustc_parse::lexer::StripTokens;14use rustc_parse::new_parser_from_source_str;15use rustc_session::parse::ParseSess;16use rustc_span::edition::{DEFAULT_EDITION, Edition};17use rustc_span::source_map::SourceMap;18use rustc_span::symbol::sym;19use rustc_span::{DUMMY_SP, FileName, Span, kw};20use tracing::debug;2122use super::GlobalTestOptions;23use crate::config::MergeDoctests;24use crate::display::Joined as _;25use crate::html::markdown::LangString;2627#[derive(Default)]28struct ParseSourceInfo {29    has_main_fn: bool,30    already_has_extern_crate: bool,31    supports_color: bool,32    has_global_allocator: bool,33    has_macro_def: bool,34    everything_else: String,35    crates: String,36    crate_attrs: String,37    maybe_crate_attrs: String,38}3940/// Builder type for `DocTestBuilder`.41pub(crate) struct BuildDocTestBuilder<'a> {42    source: &'a str,43    crate_name: Option<&'a str>,44    edition: Edition,45    can_merge_doctests: MergeDoctests,46    // If `test_id` is `None`, it means we're generating code for a code example "run" link.47    test_id: Option<String>,48    lang_str: Option<&'a LangString>,49    span: Span,50    global_crate_attrs: Vec<String>,51}5253impl<'a> BuildDocTestBuilder<'a> {54    pub(crate) fn new(source: &'a str) -> Self {55        Self {56            source,57            crate_name: None,58            edition: DEFAULT_EDITION,59            can_merge_doctests: MergeDoctests::Never,60            test_id: None,61            lang_str: None,62            span: DUMMY_SP,63            global_crate_attrs: Vec::new(),64        }65    }6667    #[inline]68    pub(crate) fn crate_name(mut self, crate_name: &'a str) -> Self {69        self.crate_name = Some(crate_name);70        self71    }7273    #[inline]74    pub(crate) fn can_merge_doctests(mut self, can_merge_doctests: MergeDoctests) -> Self {75        self.can_merge_doctests = can_merge_doctests;76        self77    }7879    #[inline]80    pub(crate) fn test_id(mut self, test_id: String) -> Self {81        self.test_id = Some(test_id);82        self83    }8485    #[inline]86    pub(crate) fn lang_str(mut self, lang_str: &'a LangString) -> Self {87        self.lang_str = Some(lang_str);88        self89    }9091    #[inline]92    pub(crate) fn span(mut self, span: Span) -> Self {93        self.span = span;94        self95    }9697    #[inline]98    pub(crate) fn edition(mut self, edition: Edition) -> Self {99        self.edition = edition;100        self101    }102103    #[inline]104    pub(crate) fn global_crate_attrs(mut self, global_crate_attrs: Vec<String>) -> Self {105        self.global_crate_attrs = global_crate_attrs;106        self107    }108109    pub(crate) fn build(self, dcx: Option<DiagCtxtHandle<'_>>) -> DocTestBuilder {110        let BuildDocTestBuilder {111            source,112            crate_name,113            edition,114            can_merge_doctests,115            // If `test_id` is `None`, it means we're generating code for a code example "run" link.116            test_id,117            lang_str,118            span,119            global_crate_attrs,120        } = self;121122        let result = rustc_driver::catch_fatal_errors(|| {123            rustc_span::create_session_if_not_set_then(edition, |_| {124                parse_source(source, &crate_name, dcx, span)125            })126        });127128        let Ok(Ok(ParseSourceInfo {129            has_main_fn,130            already_has_extern_crate,131            supports_color,132            has_global_allocator,133            has_macro_def,134            everything_else,135            crates,136            crate_attrs,137            maybe_crate_attrs,138        })) = result139        else {140            // If the AST returned an error, we don't want this doctest to be merged with the141            // others.142            return DocTestBuilder::invalid(143                Vec::new(),144                String::new(),145                String::new(),146                String::new(),147                source.to_string(),148                test_id,149            );150        };151152        debug!("crate_attrs:\n{crate_attrs}{maybe_crate_attrs}");153        debug!("crates:\n{crates}");154        debug!("after:\n{everything_else}");155        debug!("merge-doctests: {can_merge_doctests:?}");156157        // Up until now, we've been dealing with settings for the whole crate.158        // Now, infer settings for this particular test.159        //160        // Avoid tests with incompatible attributes.161        let opt_out = lang_str.is_some_and(|lang_str| {162            lang_str.compile_fail || lang_str.test_harness || lang_str.standalone_crate163        });164        let can_be_merged = if can_merge_doctests == MergeDoctests::Auto {165            // We try to look at the contents of the test to detect whether it should be merged.166            // This is not a complete list of possible failures, but it catches many cases.167            let will_probably_fail = has_global_allocator168                || !crate_attrs.is_empty()169                // If this is a merged doctest and a defined macro uses `$crate`, then the path will170                // not work, so better not put it into merged doctests.171                || (has_macro_def && everything_else.contains("$crate"));172            !opt_out && !will_probably_fail173        } else {174            can_merge_doctests != MergeDoctests::Never && !opt_out175        };176        DocTestBuilder {177            supports_color,178            has_main_fn,179            global_crate_attrs,180            crate_attrs,181            maybe_crate_attrs,182            crates,183            everything_else,184            already_has_extern_crate,185            test_id,186            invalid_ast: false,187            can_be_merged,188        }189    }190}191192/// This struct contains information about the doctest itself which is then used to generate193/// doctest source code appropriately.194pub(crate) struct DocTestBuilder {195    pub(crate) supports_color: bool,196    pub(crate) already_has_extern_crate: bool,197    pub(crate) has_main_fn: bool,198    pub(crate) global_crate_attrs: Vec<String>,199    pub(crate) crate_attrs: String,200    /// If this is a merged doctest, it will be put into `everything_else`, otherwise it will201    /// put into `crate_attrs`.202    pub(crate) maybe_crate_attrs: String,203    pub(crate) crates: String,204    pub(crate) everything_else: String,205    pub(crate) test_id: Option<String>,206    pub(crate) invalid_ast: bool,207    pub(crate) can_be_merged: bool,208}209210/// Contains needed information for doctest to be correctly generated with expected "wrapping".211pub(crate) struct WrapperInfo {212    pub(crate) before: String,213    pub(crate) after: String,214    pub(crate) returns_result: bool,215    insert_indent_space: bool,216}217218impl WrapperInfo {219    fn len(&self) -> usize {220        self.before.len() + self.after.len()221    }222}223224/// Contains a doctest information. Can be converted into code with the `to_string()` method.225pub(crate) enum DocTestWrapResult {226    Valid {227        crate_level_code: String,228        /// This field can be `None` if one of the following conditions is true:229        ///230        /// * The doctest's codeblock has the `test_harness` attribute.231        /// * The doctest has a `main` function.232        /// * The doctest has the `![no_std]` attribute.233        wrapper: Option<WrapperInfo>,234        /// Contains the doctest processed code without the wrappers (which are stored in the235        /// `wrapper` field).236        code: String,237    },238    /// Contains the original source code.239    SyntaxError(String),240}241242impl std::string::ToString for DocTestWrapResult {243    fn to_string(&self) -> String {244        match self {245            Self::SyntaxError(s) => s.clone(),246            Self::Valid { crate_level_code, wrapper, code } => {247                let mut prog_len = code.len() + crate_level_code.len();248                if let Some(wrapper) = wrapper {249                    prog_len += wrapper.len();250                    if wrapper.insert_indent_space {251                        prog_len += code.lines().count() * 4;252                    }253                }254                let mut prog = String::with_capacity(prog_len);255256                prog.push_str(crate_level_code);257                if let Some(wrapper) = wrapper {258                    prog.push_str(&wrapper.before);259260                    // add extra 4 spaces for each line to offset the code block261                    if wrapper.insert_indent_space {262                        write!(263                            prog,264                            "{}",265                            fmt::from_fn(|f| code266                                .lines()267                                .map(|line| fmt::from_fn(move |f| write!(f, "    {line}")))268                                .joined("\n", f))269                        )270                        .unwrap();271                    } else {272                        prog.push_str(code);273                    }274                    prog.push_str(&wrapper.after);275                } else {276                    prog.push_str(code);277                }278                prog279            }280        }281    }282}283284impl DocTestBuilder {285    fn invalid(286        global_crate_attrs: Vec<String>,287        crate_attrs: String,288        maybe_crate_attrs: String,289        crates: String,290        everything_else: String,291        test_id: Option<String>,292    ) -> Self {293        Self {294            supports_color: false,295            has_main_fn: false,296            global_crate_attrs,297            crate_attrs,298            maybe_crate_attrs,299            crates,300            everything_else,301            already_has_extern_crate: false,302            test_id,303            invalid_ast: true,304            can_be_merged: false,305        }306    }307308    /// Transforms a test into code that can be compiled into a Rust binary, and returns the number of309    /// lines before the test code begins.310    pub(crate) fn generate_unique_doctest(311        &self,312        test_code: &str,313        dont_insert_main: bool,314        opts: &GlobalTestOptions,315        crate_name: Option<&str>,316    ) -> (DocTestWrapResult, usize) {317        if self.invalid_ast {318            // If the AST failed to compile, no need to go generate a complete doctest, the error319            // will be better this way.320            debug!("invalid AST:\n{test_code}");321            return (DocTestWrapResult::SyntaxError(test_code.to_string()), 0);322        }323        let mut line_offset = 0;324        let mut crate_level_code = String::new();325        let processed_code = self.everything_else.trim();326        if self.global_crate_attrs.is_empty() {327            // If there aren't any attributes supplied by #![doc(test(attr(...)))], then allow some328            // lints that are commonly triggered in doctests. The crate-level test attributes are329            // commonly used to make tests fail in case they trigger warnings, so having this there in330            // that case may cause some tests to pass when they shouldn't have.331            crate_level_code.push_str("#![allow(unused)]\n");332            line_offset += 1;333        }334335        // Next, any attributes that came from #![doc(test(attr(...)))].336        for attr in &self.global_crate_attrs {337            crate_level_code.push_str(&format!("#![{attr}]\n"));338            line_offset += 1;339        }340341        // Now push any outer attributes from the example, assuming they342        // are intended to be crate attributes.343        if !self.crate_attrs.is_empty() {344            crate_level_code.push_str(&self.crate_attrs);345            if !self.crate_attrs.ends_with('\n') {346                crate_level_code.push('\n');347            }348        }349        if !self.maybe_crate_attrs.is_empty() {350            crate_level_code.push_str(&self.maybe_crate_attrs);351            if !self.maybe_crate_attrs.ends_with('\n') {352                crate_level_code.push('\n');353            }354        }355        if !self.crates.is_empty() {356            crate_level_code.push_str(&self.crates);357            if !self.crates.ends_with('\n') {358                crate_level_code.push('\n');359            }360        }361362        // Don't inject `extern crate std` because it's already injected by the363        // compiler.364        if !self.already_has_extern_crate &&365            !opts.no_crate_inject &&366            let Some(crate_name) = crate_name &&367            crate_name != "std" &&368            // Don't inject `extern crate` if the crate is never used.369            // NOTE: this is terribly inaccurate because it doesn't actually370            // parse the source, but only has false positives, not false371            // negatives.372            test_code.contains(crate_name)373        {374            // rustdoc implicitly inserts an `extern crate` item for the own crate375            // which may be unused, so we need to allow the lint.376            crate_level_code.push_str("#[allow(unused_extern_crates)]\n");377378            crate_level_code.push_str(&format!("extern crate r#{crate_name};\n"));379            line_offset += 1;380        }381382        // FIXME: This code cannot yet handle no_std test cases yet383        let wrapper = if dont_insert_main384            || self.has_main_fn385            || crate_level_code.contains("![no_std]")386        {387            None388        } else {389            let returns_result = processed_code.ends_with("(())");390            // Give each doctest main function a unique name.391            // This is for example needed for the tooling around `-C instrument-coverage`.392            let inner_fn_name = if let Some(ref test_id) = self.test_id {393                format!("_doctest_main_{test_id}")394            } else {395                "_inner".into()396            };397            let inner_attr = if self.test_id.is_some() { "#[allow(non_snake_case)] " } else { "" };398            let (main_pre, main_post) = if returns_result {399                (400                    format!(401                        "fn main() {{ {inner_attr}fn {inner_fn_name}() -> core::result::Result<(), impl core::fmt::Debug> {{\n",402                    ),403                    format!("\n}} {inner_fn_name}().unwrap() }}"),404                )405            } else if self.test_id.is_some() {406                (407                    format!("fn main() {{ {inner_attr}fn {inner_fn_name}() {{\n",),408                    format!("\n}} {inner_fn_name}() }}"),409                )410            } else {411                ("fn main() {\n".into(), "\n}".into())412            };413            // Note on newlines: We insert a line/newline *before*, and *after*414            // the doctest and adjust the `line_offset` accordingly.415            // In the case of `-C instrument-coverage`, this means that the generated416            // inner `main` function spans from the doctest opening codeblock to the417            // closing one. For example418            // /// ``` <- start of the inner main419            // /// <- code under doctest420            // /// ``` <- end of the inner main421            line_offset += 1;422423            Some(WrapperInfo {424                before: main_pre,425                after: main_post,426                returns_result,427                insert_indent_space: opts.insert_indent_space,428            })429        };430431        (432            DocTestWrapResult::Valid {433                code: processed_code.to_string(),434                wrapper,435                crate_level_code,436            },437            line_offset,438        )439    }440}441442fn reset_error_count(psess: &ParseSess) {443    // Reset errors so that they won't be reported as compiler bugs when dropping the444    // dcx. Any errors in the tests will be reported when the test file is compiled,445    // Note that we still need to cancel the errors above otherwise `Diag` will panic on446    // drop.447    psess.dcx().reset_err_count();448}449450const DOCTEST_CODE_WRAPPER: &str = "fn f(){";451452fn parse_source(453    source: &str,454    crate_name: &Option<&str>,455    parent_dcx: Option<DiagCtxtHandle<'_>>,456    span: Span,457) -> Result<ParseSourceInfo, ()> {458    use rustc_errors::DiagCtxt;459    use rustc_errors::annotate_snippet_emitter_writer::AnnotateSnippetEmitter;460    use rustc_span::source_map::FilePathMapping;461462    let mut info =463        ParseSourceInfo { already_has_extern_crate: crate_name.is_none(), ..Default::default() };464465    let wrapped_source = format!("{DOCTEST_CODE_WRAPPER}{source}\n}}");466467    let filename = FileName::anon_source_code(&wrapped_source);468469    let sm = Arc::new(SourceMap::new(FilePathMapping::empty()));470    let supports_color = match get_stderr_color_choice(ColorConfig::Auto, &std::io::stderr()) {471        ColorChoice::Auto => unreachable!(),472        ColorChoice::AlwaysAnsi | ColorChoice::Always => true,473        ColorChoice::Never => false,474    };475    info.supports_color = supports_color;476    // Any errors in parsing should also appear when the doctest is compiled for real, so just477    // send all the errors that the parser emits directly into a `Sink` instead of stderr.478    let emitter = AnnotateSnippetEmitter::new(AutoStream::never(Box::new(io::sink())));479480    // FIXME(misdreavus): pass `-Z treat-err-as-bug` to the doctest parser481    let dcx = DiagCtxt::new(Box::new(emitter)).disable_warnings();482    let psess = ParseSess::with_dcx(dcx, sm);483484    // Don't strip any tokens; it wouldn't matter anyway because the source is wrapped in a function.485    let mut parser =486        match new_parser_from_source_str(&psess, filename, wrapped_source, StripTokens::Nothing) {487            Ok(p) => p,488            Err(errs) => {489                errs.into_iter().for_each(|err| err.cancel());490                reset_error_count(&psess);491                return Err(());492            }493        };494495    fn push_to_s(s: &mut String, source: &str, span: rustc_span::Span, prev_span_hi: &mut usize) {496        let extra_len = DOCTEST_CODE_WRAPPER.len();497        // We need to shift by the length of `DOCTEST_CODE_WRAPPER` because we498        // added it at the beginning of the source we provided to the parser.499        let mut hi = span.hi().0 as usize - extra_len;500        if hi > source.len() {501            hi = source.len();502        }503        s.push_str(&source[*prev_span_hi..hi]);504        *prev_span_hi = hi;505    }506507    fn check_item(item: &ast::Item, info: &mut ParseSourceInfo, crate_name: &Option<&str>) -> bool {508        let mut is_extern_crate = false;509        if !info.has_global_allocator510            && item.attrs.iter().any(|attr| attr.has_name(sym::global_allocator))511        {512            info.has_global_allocator = true;513        }514        match item.kind {515            ast::ItemKind::Fn(ref fn_item) if !info.has_main_fn => {516                if fn_item.ident.name == sym::main {517                    info.has_main_fn = true;518                }519            }520            ast::ItemKind::ExternCrate(original, ident) => {521                is_extern_crate = true;522                if !info.already_has_extern_crate523                    && let Some(crate_name) = crate_name524                {525                    info.already_has_extern_crate = match original {526                        Some(name) => name.as_str() == *crate_name,527                        None => ident.as_str() == *crate_name,528                    };529                }530            }531            ast::ItemKind::MacroDef(..) => {532                info.has_macro_def = true;533            }534            _ => {}535        }536        is_extern_crate537    }538539    let mut prev_span_hi = 0;540    let not_crate_attrs = &[sym::forbid, sym::allow, sym::warn, sym::deny, sym::expect];541    let parsed = parser.parse_item(542        rustc_parse::parser::ForceCollect::No,543        rustc_parse::parser::AllowConstBlockItems::No,544    );545546    let result = match parsed {547        Ok(Some(ref item))548            if let ast::ItemKind::Fn(ref fn_item) = item.kind549                && let Some(ref body) = fn_item.body =>550        {551            for attr in &item.attrs {552                if attr.style == AttrStyle::Outer || attr.has_any_name(not_crate_attrs) {553                    // There is one exception to these attributes:554                    // `#![allow(internal_features)]`. If this attribute is used, we need to555                    // consider it only as a crate-level attribute.556                    if attr.has_name(sym::allow)557                        && let Some(list) = attr.meta_item_list()558                        && list.iter().any(|sub_attr| {559                            sub_attr.has_name(sym::internal_features)560                                || sub_attr.has_name(sym::incomplete_features)561                        })562                    {563                        push_to_s(&mut info.crate_attrs, source, attr.span, &mut prev_span_hi);564                    } else {565                        push_to_s(566                            &mut info.maybe_crate_attrs,567                            source,568                            attr.span,569                            &mut prev_span_hi,570                        );571                    }572                } else {573                    push_to_s(&mut info.crate_attrs, source, attr.span, &mut prev_span_hi);574                }575            }576            let mut has_non_items = false;577            for stmt in &body.stmts {578                let mut is_extern_crate = false;579                match stmt.kind {580                    StmtKind::Item(ref item) => {581                        is_extern_crate = check_item(item, &mut info, crate_name);582                    }583                    // We assume that the macro calls will expand to item(s) even though they could584                    // expand to statements and expressions.585                    StmtKind::MacCall(ref mac_call) => {586                        if !info.has_main_fn {587                            // For backward compatibility, we look for the token sequence `fn main(…)`588                            // in the macro input (!) to crudely detect main functions "masked by a589                            // wrapper macro". For the record, this is a horrible heuristic!590                            // See <https://github.com/rust-lang/rust/issues/56898>.591                            let mut iter = mac_call.mac.args.tokens.iter();592                            while let Some(token) = iter.next() {593                                if let TokenTree::Token(token, _) = token594                                    && let TokenKind::Ident(kw::Fn, _) = token.kind595                                    && let Some(TokenTree::Token(ident, _)) = iter.peek()596                                    && let TokenKind::Ident(sym::main, _) = ident.kind597                                    && let Some(TokenTree::Delimited(.., Delimiter::Parenthesis, _)) = {598                                        iter.next();599                                        iter.peek()600                                    }601                                {602                                    info.has_main_fn = true;603                                    break;604                                }605                            }606                        }607                    }608                    StmtKind::Expr(ref expr) => {609                        if matches!(expr.kind, ast::ExprKind::Err(_)) {610                            reset_error_count(&psess);611                            return Err(());612                        }613                        has_non_items = true;614                    }615                    StmtKind::Let(_) | StmtKind::Semi(_) | StmtKind::Empty => has_non_items = true,616                }617618                // Weirdly enough, the `Stmt` span doesn't include its attributes, so we need to619                // tweak the span to include the attributes as well.620                let mut span = stmt.span;621                if let Some(attr) =622                    stmt.kind.attrs().iter().find(|attr| attr.style == AttrStyle::Outer)623                {624                    span = span.with_lo(attr.span.lo());625                }626                if info.everything_else.is_empty()627                    && (!info.maybe_crate_attrs.is_empty() || !info.crate_attrs.is_empty())628                {629                    // To keep the doctest code "as close as possible" to the original, we insert630                    // all the code located between this new span and the previous span which631                    // might contain code comments and backlines.632                    push_to_s(&mut info.crates, source, span.shrink_to_lo(), &mut prev_span_hi);633                }634                if !is_extern_crate {635                    push_to_s(&mut info.everything_else, source, span, &mut prev_span_hi);636                } else {637                    push_to_s(&mut info.crates, source, span, &mut prev_span_hi);638                }639            }640            if has_non_items {641                if info.has_main_fn642                    && let Some(dcx) = parent_dcx643                    && !span.is_dummy()644                {645                    dcx.span_warn(646                        span,647                        "the `main` function of this doctest won't be run as it contains \648                         expressions at the top level, meaning that the whole doctest code will be \649                         wrapped in a function",650                    );651                }652                info.has_main_fn = false;653            }654            Ok(info)655        }656        Err(e) => {657            e.cancel();658            Err(())659        }660        _ => Err(()),661    };662663    reset_error_count(&psess);664    result665}