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, InnerSpan, Span, kw};20use tracing::debug;2122use super::{CodeLineMapping, 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    /// Inner attributes (`#![...]`) from the source that have to be put at the crate level.37    crate_attrs: String,38    /// Inner attributes (`#![...]`) from the source that can be put into a module and therefore do39    /// not inhibit merging: even in the merged test, the attributes can be isolated to the test.40    module_attrs: String,41}4243/// Builder type for `DocTestBuilder`.44pub(crate) struct BuildDocTestBuilder<'a> {45    source: &'a str,46    crate_name: Option<&'a str>,47    edition: Edition,48    can_merge_doctests: MergeDoctests,49    // If `test_id` is `None`, it means we're generating code for a code example "run" link.50    test_id: Option<String>,51    lang_str: Option<&'a LangString>,52    span: Span,53    code_mappings: &'a [CodeLineMapping],54    global_crate_attrs: Vec<String>,55}5657impl<'a> BuildDocTestBuilder<'a> {58    pub(crate) fn new(source: &'a str) -> Self {59        Self {60            source,61            crate_name: None,62            edition: DEFAULT_EDITION,63            can_merge_doctests: MergeDoctests::Never,64            test_id: None,65            lang_str: None,66            span: DUMMY_SP,67            code_mappings: &[],68            global_crate_attrs: Vec::new(),69        }70    }7172    #[inline]73    pub(crate) fn crate_name(mut self, crate_name: &'a str) -> Self {74        self.crate_name = Some(crate_name);75        self76    }7778    #[inline]79    pub(crate) fn can_merge_doctests(mut self, can_merge_doctests: MergeDoctests) -> Self {80        self.can_merge_doctests = can_merge_doctests;81        self82    }8384    #[inline]85    pub(crate) fn test_id(mut self, test_id: String) -> Self {86        self.test_id = Some(test_id);87        self88    }8990    #[inline]91    pub(crate) fn lang_str(mut self, lang_str: &'a LangString) -> Self {92        self.lang_str = Some(lang_str);93        self94    }9596    #[inline]97    pub(crate) fn span(mut self, span: Span) -> Self {98        self.span = span;99        self100    }101102    #[inline]103    pub(crate) fn code_mappings(mut self, code_mappings: &'a [CodeLineMapping]) -> Self {104        self.code_mappings = code_mappings;105        self106    }107108    #[inline]109    pub(crate) fn edition(mut self, edition: Edition) -> Self {110        self.edition = edition;111        self112    }113114    #[inline]115    pub(crate) fn global_crate_attrs(mut self, global_crate_attrs: Vec<String>) -> Self {116        self.global_crate_attrs = global_crate_attrs;117        self118    }119120    pub(crate) fn build(self, dcx: Option<DiagCtxtHandle<'_>>) -> DocTestBuilder {121        let BuildDocTestBuilder {122            source,123            crate_name,124            edition,125            can_merge_doctests,126            // If `test_id` is `None`, it means we're generating code for a code example "run" link.127            test_id,128            lang_str,129            span,130            code_mappings,131            global_crate_attrs,132        } = self;133134        let result = rustc_driver::catch_fatal_errors(|| {135            rustc_span::create_session_if_not_set_then(edition, |_| {136                parse_source(source, &crate_name, dcx, span, code_mappings)137            })138        });139140        let Ok(Ok(ParseSourceInfo {141            has_main_fn,142            already_has_extern_crate,143            supports_color,144            has_global_allocator,145            has_macro_def,146            everything_else,147            crates,148            crate_attrs,149            module_attrs,150        })) = result151        else {152            // If the AST returned an error, we don't want this doctest to be merged with the153            // others.154            return DocTestBuilder::invalid(155                Vec::new(),156                String::new(),157                String::new(),158                String::new(),159                source.to_string(),160                test_id,161            );162        };163164        debug!("crate_attrs:\n{crate_attrs}{module_attrs}");165        debug!("crates:\n{crates}");166        debug!("after:\n{everything_else}");167        debug!("merge-doctests: {can_merge_doctests:?}");168169        // Up until now, we've been dealing with settings for the whole crate.170        // Now, infer settings for this particular test.171        //172        // Avoid tests with incompatible attributes.173        let opt_out = lang_str.is_some_and(|lang_str| {174            lang_str.compile_fail || lang_str.test_harness || lang_str.standalone_crate175        });176        let can_be_merged = if can_merge_doctests == MergeDoctests::Auto {177            // We try to look at the contents of the test to detect whether it should be merged.178            // This is not a complete list of possible failures, but it catches many cases.179            let will_probably_fail = has_global_allocator180                || !crate_attrs.is_empty()181                // If this is a merged doctest and a defined macro uses `$crate`, then the path will182                // not work, so better not put it into merged doctests.183                || (has_macro_def && everything_else.contains("$crate"));184            !opt_out && !will_probably_fail185        } else {186            can_merge_doctests != MergeDoctests::Never && !opt_out187        };188        DocTestBuilder {189            supports_color,190            has_main_fn,191            global_crate_attrs,192            crate_attrs,193            module_attrs,194            crates,195            everything_else,196            already_has_extern_crate,197            test_id,198            invalid_ast: false,199            can_be_merged,200        }201    }202}203204/// This struct contains information about the doctest itself which is then used to generate205/// doctest source code appropriately.206pub(crate) struct DocTestBuilder {207    pub(crate) supports_color: bool,208    pub(crate) already_has_extern_crate: bool,209    pub(crate) has_main_fn: bool,210    pub(crate) global_crate_attrs: Vec<String>,211    pub(crate) crate_attrs: String,212    /// If this is a merged doctest, it will be put into `everything_else`, otherwise it will213    /// put into `crate_attrs`.214    pub(crate) module_attrs: String,215    pub(crate) crates: String,216    pub(crate) everything_else: String,217    pub(crate) test_id: Option<String>,218    pub(crate) invalid_ast: bool,219    pub(crate) can_be_merged: bool,220}221222/// Contains needed information for doctest to be correctly generated with expected "wrapping".223pub(crate) struct WrapperInfo {224    pub(crate) before: String,225    pub(crate) after: String,226    pub(crate) returns_result: bool,227    insert_indent_space: bool,228}229230impl WrapperInfo {231    fn len(&self) -> usize {232        self.before.len() + self.after.len()233    }234}235236/// Contains a doctest information. Can be converted into code with the `to_string()` method.237pub(crate) enum DocTestWrapResult {238    Valid {239        crate_level_code: String,240        /// This field can be `None` if one of the following conditions is true:241        ///242        /// * The doctest's codeblock has the `test_harness` attribute.243        /// * The doctest has a `main` function.244        /// * The doctest has the `![no_std]` attribute.245        wrapper: Option<WrapperInfo>,246        /// Contains the doctest processed code without the wrappers (which are stored in the247        /// `wrapper` field).248        code: String,249    },250    /// Contains the original source code.251    SyntaxError(String),252}253254impl std::string::ToString for DocTestWrapResult {255    fn to_string(&self) -> String {256        match self {257            Self::SyntaxError(s) => s.clone(),258            Self::Valid { crate_level_code, wrapper, code } => {259                let mut prog_len = code.len() + crate_level_code.len();260                if let Some(wrapper) = wrapper {261                    prog_len += wrapper.len();262                    if wrapper.insert_indent_space {263                        prog_len += code.lines().count() * 4;264                    }265                }266                let mut prog = String::with_capacity(prog_len);267268                prog.push_str(crate_level_code);269                if let Some(wrapper) = wrapper {270                    prog.push_str(&wrapper.before);271272                    // add extra 4 spaces for each line to offset the code block273                    if wrapper.insert_indent_space {274                        write!(275                            prog,276                            "{}",277                            fmt::from_fn(|f| code278                                .lines()279                                .map(|line| fmt::from_fn(move |f| write!(f, "    {line}")))280                                .joined("\n", f))281                        )282                        .unwrap();283                    } else {284                        prog.push_str(code);285                    }286                    prog.push_str(&wrapper.after);287                } else {288                    prog.push_str(code);289                }290                prog291            }292        }293    }294}295296impl DocTestBuilder {297    fn invalid(298        global_crate_attrs: Vec<String>,299        crate_attrs: String,300        module_attrs: String,301        crates: String,302        everything_else: String,303        test_id: Option<String>,304    ) -> Self {305        Self {306            supports_color: false,307            has_main_fn: false,308            global_crate_attrs,309            crate_attrs,310            module_attrs,311            crates,312            everything_else,313            already_has_extern_crate: false,314            test_id,315            invalid_ast: true,316            can_be_merged: false,317        }318    }319320    /// Transforms a test into code that can be compiled into a Rust binary, and returns the number of321    /// lines before the test code begins.322    pub(crate) fn generate_unique_doctest(323        &self,324        test_code: &str,325        dont_insert_main: bool,326        opts: &GlobalTestOptions,327        crate_name: Option<&str>,328    ) -> (DocTestWrapResult, usize) {329        if self.invalid_ast {330            // If the AST failed to compile, no need to go generate a complete doctest, the error331            // will be better this way.332            debug!("invalid AST:\n{test_code}");333            return (DocTestWrapResult::SyntaxError(test_code.to_string()), 0);334        }335        let mut line_offset = 0;336        let mut crate_level_code = String::new();337        let processed_code = self.everything_else.trim();338        if self.global_crate_attrs.is_empty() {339            // If there aren't any attributes supplied by #![doc(test(attr(...)))], then allow some340            // lints that are commonly triggered in doctests. The crate-level test attributes are341            // commonly used to make tests fail in case they trigger warnings, so having this there in342            // that case may cause some tests to pass when they shouldn't have.343            crate_level_code.push_str("#![allow(unused)]\n");344            line_offset += 1;345        }346347        // Next, any attributes that came from #![doc(test(attr(...)))].348        for attr in &self.global_crate_attrs {349            crate_level_code.push_str(&format!("#![{attr}]\n"));350            line_offset += 1;351        }352353        // Now push any outer attributes from the example (both crate and module attributes).354        if !self.crate_attrs.is_empty() {355            crate_level_code.push_str(&self.crate_attrs);356            if !self.crate_attrs.ends_with('\n') {357                crate_level_code.push('\n');358            }359        }360        if !self.module_attrs.is_empty() {361            crate_level_code.push_str(&self.module_attrs);362            if !self.module_attrs.ends_with('\n') {363                crate_level_code.push('\n');364            }365        }366        if !self.crates.is_empty() {367            crate_level_code.push_str(&self.crates);368            if !self.crates.ends_with('\n') {369                crate_level_code.push('\n');370            }371        }372373        // Don't inject `extern crate std` because it's already injected by the374        // compiler.375        if !self.already_has_extern_crate &&376            !opts.no_crate_inject &&377            let Some(crate_name) = crate_name &&378            crate_name != "std" &&379            // Don't inject `extern crate` if the crate is never used.380            // NOTE: this is terribly inaccurate because it doesn't actually381            // parse the source, but only has false positives, not false382            // negatives.383            test_code.contains(crate_name)384        {385            // rustdoc implicitly inserts an `extern crate` item for the own crate386            // which may be unused, so we need to allow the lint.387            crate_level_code.push_str("#[allow(unused_extern_crates)]\n");388389            crate_level_code.push_str(&format!("extern crate r#{crate_name};\n"));390            line_offset += 1;391        }392393        // FIXME: This code cannot yet handle no_std test cases yet394        let wrapper = if dont_insert_main395            || self.has_main_fn396            || crate_level_code.contains("![no_std]")397        {398            None399        } else {400            let returns_result = processed_code.ends_with("(())");401            // Give each doctest main function a unique name.402            // This is for example needed for the tooling around `-C instrument-coverage`.403            let inner_fn_name = if let Some(ref test_id) = self.test_id {404                format!("_doctest_main_{test_id}")405            } else {406                "_inner".into()407            };408            let inner_attr = if self.test_id.is_some() { "#[allow(non_snake_case)] " } else { "" };409            let (main_pre, main_post) = if returns_result {410                (411                    format!(412                        "fn main() {{ {inner_attr}fn {inner_fn_name}() -> core::result::Result<(), impl core::fmt::Debug> {{\n",413                    ),414                    format!("\n}} {inner_fn_name}().unwrap() }}"),415                )416            } else if self.test_id.is_some() {417                (418                    format!("fn main() {{ {inner_attr}fn {inner_fn_name}() {{\n",),419                    format!("\n}} {inner_fn_name}() }}"),420                )421            } else {422                ("fn main() {\n".into(), "\n}".into())423            };424            // Note on newlines: We insert a line/newline *before*, and *after*425            // the doctest and adjust the `line_offset` accordingly.426            // In the case of `-C instrument-coverage`, this means that the generated427            // inner `main` function spans from the doctest opening codeblock to the428            // closing one. For example429            // /// ``` <- start of the inner main430            // /// <- code under doctest431            // /// ``` <- end of the inner main432            line_offset += 1;433434            Some(WrapperInfo {435                before: main_pre,436                after: main_post,437                returns_result,438                insert_indent_space: opts.insert_indent_space,439            })440        };441442        (443            DocTestWrapResult::Valid {444                code: processed_code.to_string(),445                wrapper,446                crate_level_code,447            },448            line_offset,449        )450    }451}452453fn reset_error_count(psess: &ParseSess) {454    // Reset errors so that they won't be reported as compiler bugs when dropping the455    // dcx. Any errors in the tests will be reported when the test file is compiled,456    // Note that we still need to cancel the errors above otherwise `Diag` will panic on457    // drop.458    psess.dcx().reset_err_count();459}460461const DOCTEST_CODE_WRAPPER: &str = "fn f(){";462463fn parse_source(464    source: &str,465    crate_name: &Option<&str>,466    parent_dcx: Option<DiagCtxtHandle<'_>>,467    span: Span,468    code_mappings: &[CodeLineMapping],469) -> Result<ParseSourceInfo, ()> {470    use rustc_errors::DiagCtxt;471    use rustc_errors::annotate_snippet_emitter_writer::AnnotateSnippetEmitter;472    use rustc_span::source_map::FilePathMapping;473474    let mut info =475        ParseSourceInfo { already_has_extern_crate: crate_name.is_none(), ..Default::default() };476477    let wrapped_source = format!("{DOCTEST_CODE_WRAPPER}{source}\n}}");478479    let filename = FileName::anon_source_code(&wrapped_source);480481    let sm = Arc::new(SourceMap::new(FilePathMapping::empty()));482    let supports_color = match get_stderr_color_choice(ColorConfig::Auto, &std::io::stderr()) {483        ColorChoice::Auto => unreachable!(),484        ColorChoice::AlwaysAnsi | ColorChoice::Always => true,485        ColorChoice::Never => false,486    };487    info.supports_color = supports_color;488    // Any errors in parsing should also appear when the doctest is compiled for real, so just489    // send all the errors that the parser emits directly into a `Sink` instead of stderr.490    let emitter = AnnotateSnippetEmitter::new(AutoStream::never(Box::new(io::sink())));491492    // FIXME(misdreavus): pass `-Z treat-err-as-bug` to the doctest parser493    let dcx = DiagCtxt::new(Box::new(emitter)).disable_warnings();494    let psess = ParseSess::with_dcx(dcx, sm);495496    // Don't strip any tokens; it wouldn't matter anyway because the source is wrapped in a function.497    let mut parser =498        match new_parser_from_source_str(&psess, filename, wrapped_source, StripTokens::Nothing) {499            Ok(p) => p,500            Err(errs) => {501                errs.into_iter().for_each(|err| err.cancel());502                reset_error_count(&psess);503                return Err(());504            }505        };506507    fn push_to_s(s: &mut String, source: &str, span: rustc_span::Span, prev_span_hi: &mut usize) {508        let extra_len = DOCTEST_CODE_WRAPPER.len();509        // We need to shift by the length of `DOCTEST_CODE_WRAPPER` because we510        // added it at the beginning of the source we provided to the parser.511        let mut hi = span.hi().0 as usize - extra_len;512        if hi > source.len() {513            hi = source.len();514        }515        s.push_str(&source[*prev_span_hi..hi]);516        *prev_span_hi = hi;517    }518519    fn span_in_doctest_source(span: Span, code_mappings: &[CodeLineMapping]) -> Option<Span> {520        let extra_len = DOCTEST_CODE_WRAPPER.len();521        let lo = (span.lo().0 as usize).checked_sub(extra_len)?;522        let hi = (span.hi().0 as usize).checked_sub(extra_len)?;523        if hi < lo {524            return None;525        }526        code_mappings.iter().find_map(|mapping| {527            if mapping.generated.start <= lo && hi <= mapping.generated.end {528                let start = lo - mapping.generated.start;529                let end = hi - mapping.generated.start;530                Some(mapping.original.from_inner(InnerSpan::new(start, end)))531            } else {532                None533            }534        })535    }536537    fn check_item(item: &ast::Item, info: &mut ParseSourceInfo, crate_name: &Option<&str>) -> bool {538        let mut is_extern_crate = false;539        if !info.has_global_allocator540            && item.attrs.iter().any(|attr| attr.has_name(sym::global_allocator))541        {542            info.has_global_allocator = true;543        }544        match item.kind {545            ast::ItemKind::Fn(ref fn_item) if !info.has_main_fn => {546                if fn_item.ident.name == sym::main {547                    info.has_main_fn = true;548                }549            }550            ast::ItemKind::ExternCrate(original, ident) => {551                is_extern_crate = true;552                if !info.already_has_extern_crate553                    && let Some(crate_name) = crate_name554                {555                    info.already_has_extern_crate = match original {556                        Some(name) => name.as_str() == *crate_name,557                        None => ident.as_str() == *crate_name,558                    };559                }560            }561            ast::ItemKind::MacroDef(..) => {562                info.has_macro_def = true;563            }564            _ => {}565        }566        is_extern_crate567    }568569    let mut prev_span_hi = 0;570    let not_crate_attrs = &[sym::forbid, sym::allow, sym::warn, sym::deny, sym::expect];571    let parsed = parser.parse_item(572        rustc_parse::parser::ForceCollect::No,573        rustc_parse::parser::AllowConstBlockItems::No,574    );575576    let result = match parsed {577        Ok(Some(ref item))578            if let ast::ItemKind::Fn(ref fn_item) = item.kind579                && let Some(ref body) = fn_item.body =>580        {581            for attr in &item.attrs {582                if attr.style == AttrStyle::Outer || attr.has_any_name(not_crate_attrs) {583                    // There is one exception to these attributes:584                    // `#![allow(internal_features)]`. If this attribute is used, we need to585                    // consider it only as a crate-level attribute.586                    if attr.has_name(sym::allow)587                        && let Some(list) = attr.meta_item_list()588                        && list.iter().any(|sub_attr| {589                            sub_attr.has_name(sym::internal_features)590                                || sub_attr.has_name(sym::incomplete_features)591                        })592                    {593                        push_to_s(&mut info.crate_attrs, source, attr.span, &mut prev_span_hi);594                    } else {595                        push_to_s(&mut info.module_attrs, source, attr.span, &mut prev_span_hi);596                    }597                } else {598                    push_to_s(&mut info.crate_attrs, source, attr.span, &mut prev_span_hi);599                }600            }601            let mut has_non_items = false;602            let mut first_non_item_span = None;603            for stmt in &body.stmts {604                let mut is_extern_crate = false;605                match stmt.kind {606                    StmtKind::Item(ref item) => {607                        is_extern_crate = check_item(item, &mut info, crate_name);608                    }609                    // We assume that the macro calls will expand to item(s) even though they could610                    // expand to statements and expressions.611                    StmtKind::MacCall(ref mac_call) => {612                        if !info.has_main_fn {613                            // For backward compatibility, we look for the token sequence `fn main(…)`614                            // in the macro input (!) to crudely detect main functions "masked by a615                            // wrapper macro". For the record, this is a horrible heuristic!616                            // See <https://github.com/rust-lang/rust/issues/56898>.617                            let mut iter = mac_call.mac.args.tokens.iter();618                            while let Some(token) = iter.next() {619                                if let TokenTree::Token(token, _) = token620                                    && let TokenKind::Ident(kw::Fn, _) = token.kind621                                    && let Some(TokenTree::Token(ident, _)) = iter.peek()622                                    && let TokenKind::Ident(sym::main, _) = ident.kind623                                    && let Some(TokenTree::Delimited(.., Delimiter::Parenthesis, _)) = {624                                        iter.next();625                                        iter.peek()626                                    }627                                {628                                    info.has_main_fn = true;629                                    break;630                                }631                            }632                        }633                    }634                    StmtKind::Expr(ref expr) => {635                        if matches!(expr.kind, ast::ExprKind::Err(_)) {636                            reset_error_count(&psess);637                            return Err(());638                        }639                        has_non_items = true;640                        first_non_item_span.get_or_insert(stmt.span);641                    }642                    StmtKind::Let(_) | StmtKind::Semi(_) | StmtKind::Empty => {643                        has_non_items = true;644                        first_non_item_span.get_or_insert(stmt.span);645                    }646                }647648                // Weirdly enough, the `Stmt` span doesn't include its attributes, so we need to649                // tweak the span to include the attributes as well.650                let mut span = stmt.span;651                if let Some(attr) =652                    stmt.kind.attrs().iter().find(|attr| attr.style == AttrStyle::Outer)653                {654                    span = span.with_lo(attr.span.lo());655                }656                if info.everything_else.is_empty()657                    && (!info.module_attrs.is_empty() || !info.crate_attrs.is_empty())658                {659                    // To keep the doctest code "as close as possible" to the original, we insert660                    // all the code located between this new span and the previous span which661                    // might contain code comments and backlines.662                    push_to_s(&mut info.crates, source, span.shrink_to_lo(), &mut prev_span_hi);663                }664                if !is_extern_crate {665                    push_to_s(&mut info.everything_else, source, span, &mut prev_span_hi);666                } else {667                    push_to_s(&mut info.crates, source, span, &mut prev_span_hi);668                }669            }670            if has_non_items {671                let warning_span = first_non_item_span672                    .and_then(|span| span_in_doctest_source(span, code_mappings))673                    .unwrap_or(span);674                if info.has_main_fn675                    && let Some(dcx) = parent_dcx676                    && !warning_span.is_dummy()677                {678                    dcx.span_warn(679                        warning_span,680                        "the `main` function of this doctest won't be run as it contains \681                         expressions at the top level, meaning that the whole doctest code will be \682                         wrapped in a function",683                    );684                }685                info.has_main_fn = false;686            }687            Ok(info)688        }689        Err(e) => {690            e.cancel();691            Err(())692        }693        _ => Err(()),694    };695696    reset_error_count(&psess);697    result698}