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    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    code_mappings: &'a [CodeLineMapping],51    global_crate_attrs: Vec<String>,52}5354impl<'a> BuildDocTestBuilder<'a> {55    pub(crate) fn new(source: &'a str) -> Self {56        Self {57            source,58            crate_name: None,59            edition: DEFAULT_EDITION,60            can_merge_doctests: MergeDoctests::Never,61            test_id: None,62            lang_str: None,63            span: DUMMY_SP,64            code_mappings: &[],65            global_crate_attrs: Vec::new(),66        }67    }6869    #[inline]70    pub(crate) fn crate_name(mut self, crate_name: &'a str) -> Self {71        self.crate_name = Some(crate_name);72        self73    }7475    #[inline]76    pub(crate) fn can_merge_doctests(mut self, can_merge_doctests: MergeDoctests) -> Self {77        self.can_merge_doctests = can_merge_doctests;78        self79    }8081    #[inline]82    pub(crate) fn test_id(mut self, test_id: String) -> Self {83        self.test_id = Some(test_id);84        self85    }8687    #[inline]88    pub(crate) fn lang_str(mut self, lang_str: &'a LangString) -> Self {89        self.lang_str = Some(lang_str);90        self91    }9293    #[inline]94    pub(crate) fn span(mut self, span: Span) -> Self {95        self.span = span;96        self97    }9899    #[inline]100    pub(crate) fn code_mappings(mut self, code_mappings: &'a [CodeLineMapping]) -> Self {101        self.code_mappings = code_mappings;102        self103    }104105    #[inline]106    pub(crate) fn edition(mut self, edition: Edition) -> Self {107        self.edition = edition;108        self109    }110111    #[inline]112    pub(crate) fn global_crate_attrs(mut self, global_crate_attrs: Vec<String>) -> Self {113        self.global_crate_attrs = global_crate_attrs;114        self115    }116117    pub(crate) fn build(self, dcx: Option<DiagCtxtHandle<'_>>) -> DocTestBuilder {118        let BuildDocTestBuilder {119            source,120            crate_name,121            edition,122            can_merge_doctests,123            // If `test_id` is `None`, it means we're generating code for a code example "run" link.124            test_id,125            lang_str,126            span,127            code_mappings,128            global_crate_attrs,129        } = self;130131        let result = rustc_driver::catch_fatal_errors(|| {132            rustc_span::create_session_if_not_set_then(edition, |_| {133                parse_source(source, &crate_name, dcx, span, code_mappings)134            })135        });136137        let Ok(Ok(ParseSourceInfo {138            has_main_fn,139            already_has_extern_crate,140            supports_color,141            has_global_allocator,142            has_macro_def,143            everything_else,144            crates,145            crate_attrs,146            maybe_crate_attrs,147        })) = result148        else {149            // If the AST returned an error, we don't want this doctest to be merged with the150            // others.151            return DocTestBuilder::invalid(152                Vec::new(),153                String::new(),154                String::new(),155                String::new(),156                source.to_string(),157                test_id,158            );159        };160161        debug!("crate_attrs:\n{crate_attrs}{maybe_crate_attrs}");162        debug!("crates:\n{crates}");163        debug!("after:\n{everything_else}");164        debug!("merge-doctests: {can_merge_doctests:?}");165166        // Up until now, we've been dealing with settings for the whole crate.167        // Now, infer settings for this particular test.168        //169        // Avoid tests with incompatible attributes.170        let opt_out = lang_str.is_some_and(|lang_str| {171            lang_str.compile_fail || lang_str.test_harness || lang_str.standalone_crate172        });173        let can_be_merged = if can_merge_doctests == MergeDoctests::Auto {174            // We try to look at the contents of the test to detect whether it should be merged.175            // This is not a complete list of possible failures, but it catches many cases.176            let will_probably_fail = has_global_allocator177                || !crate_attrs.is_empty()178                // If this is a merged doctest and a defined macro uses `$crate`, then the path will179                // not work, so better not put it into merged doctests.180                || (has_macro_def && everything_else.contains("$crate"));181            !opt_out && !will_probably_fail182        } else {183            can_merge_doctests != MergeDoctests::Never && !opt_out184        };185        DocTestBuilder {186            supports_color,187            has_main_fn,188            global_crate_attrs,189            crate_attrs,190            maybe_crate_attrs,191            crates,192            everything_else,193            already_has_extern_crate,194            test_id,195            invalid_ast: false,196            can_be_merged,197        }198    }199}200201/// This struct contains information about the doctest itself which is then used to generate202/// doctest source code appropriately.203pub(crate) struct DocTestBuilder {204    pub(crate) supports_color: bool,205    pub(crate) already_has_extern_crate: bool,206    pub(crate) has_main_fn: bool,207    pub(crate) global_crate_attrs: Vec<String>,208    pub(crate) crate_attrs: String,209    /// If this is a merged doctest, it will be put into `everything_else`, otherwise it will210    /// put into `crate_attrs`.211    pub(crate) maybe_crate_attrs: String,212    pub(crate) crates: String,213    pub(crate) everything_else: String,214    pub(crate) test_id: Option<String>,215    pub(crate) invalid_ast: bool,216    pub(crate) can_be_merged: bool,217}218219/// Contains needed information for doctest to be correctly generated with expected "wrapping".220pub(crate) struct WrapperInfo {221    pub(crate) before: String,222    pub(crate) after: String,223    pub(crate) returns_result: bool,224    insert_indent_space: bool,225}226227impl WrapperInfo {228    fn len(&self) -> usize {229        self.before.len() + self.after.len()230    }231}232233/// Contains a doctest information. Can be converted into code with the `to_string()` method.234pub(crate) enum DocTestWrapResult {235    Valid {236        crate_level_code: String,237        /// This field can be `None` if one of the following conditions is true:238        ///239        /// * The doctest's codeblock has the `test_harness` attribute.240        /// * The doctest has a `main` function.241        /// * The doctest has the `![no_std]` attribute.242        wrapper: Option<WrapperInfo>,243        /// Contains the doctest processed code without the wrappers (which are stored in the244        /// `wrapper` field).245        code: String,246    },247    /// Contains the original source code.248    SyntaxError(String),249}250251impl std::string::ToString for DocTestWrapResult {252    fn to_string(&self) -> String {253        match self {254            Self::SyntaxError(s) => s.clone(),255            Self::Valid { crate_level_code, wrapper, code } => {256                let mut prog_len = code.len() + crate_level_code.len();257                if let Some(wrapper) = wrapper {258                    prog_len += wrapper.len();259                    if wrapper.insert_indent_space {260                        prog_len += code.lines().count() * 4;261                    }262                }263                let mut prog = String::with_capacity(prog_len);264265                prog.push_str(crate_level_code);266                if let Some(wrapper) = wrapper {267                    prog.push_str(&wrapper.before);268269                    // add extra 4 spaces for each line to offset the code block270                    if wrapper.insert_indent_space {271                        write!(272                            prog,273                            "{}",274                            fmt::from_fn(|f| code275                                .lines()276                                .map(|line| fmt::from_fn(move |f| write!(f, "    {line}")))277                                .joined("\n", f))278                        )279                        .unwrap();280                    } else {281                        prog.push_str(code);282                    }283                    prog.push_str(&wrapper.after);284                } else {285                    prog.push_str(code);286                }287                prog288            }289        }290    }291}292293impl DocTestBuilder {294    fn invalid(295        global_crate_attrs: Vec<String>,296        crate_attrs: String,297        maybe_crate_attrs: String,298        crates: String,299        everything_else: String,300        test_id: Option<String>,301    ) -> Self {302        Self {303            supports_color: false,304            has_main_fn: false,305            global_crate_attrs,306            crate_attrs,307            maybe_crate_attrs,308            crates,309            everything_else,310            already_has_extern_crate: false,311            test_id,312            invalid_ast: true,313            can_be_merged: false,314        }315    }316317    /// Transforms a test into code that can be compiled into a Rust binary, and returns the number of318    /// lines before the test code begins.319    pub(crate) fn generate_unique_doctest(320        &self,321        test_code: &str,322        dont_insert_main: bool,323        opts: &GlobalTestOptions,324        crate_name: Option<&str>,325    ) -> (DocTestWrapResult, usize) {326        if self.invalid_ast {327            // If the AST failed to compile, no need to go generate a complete doctest, the error328            // will be better this way.329            debug!("invalid AST:\n{test_code}");330            return (DocTestWrapResult::SyntaxError(test_code.to_string()), 0);331        }332        let mut line_offset = 0;333        let mut crate_level_code = String::new();334        let processed_code = self.everything_else.trim();335        if self.global_crate_attrs.is_empty() {336            // If there aren't any attributes supplied by #![doc(test(attr(...)))], then allow some337            // lints that are commonly triggered in doctests. The crate-level test attributes are338            // commonly used to make tests fail in case they trigger warnings, so having this there in339            // that case may cause some tests to pass when they shouldn't have.340            crate_level_code.push_str("#![allow(unused)]\n");341            line_offset += 1;342        }343344        // Next, any attributes that came from #![doc(test(attr(...)))].345        for attr in &self.global_crate_attrs {346            crate_level_code.push_str(&format!("#![{attr}]\n"));347            line_offset += 1;348        }349350        // Now push any outer attributes from the example, assuming they351        // are intended to be crate attributes.352        if !self.crate_attrs.is_empty() {353            crate_level_code.push_str(&self.crate_attrs);354            if !self.crate_attrs.ends_with('\n') {355                crate_level_code.push('\n');356            }357        }358        if !self.maybe_crate_attrs.is_empty() {359            crate_level_code.push_str(&self.maybe_crate_attrs);360            if !self.maybe_crate_attrs.ends_with('\n') {361                crate_level_code.push('\n');362            }363        }364        if !self.crates.is_empty() {365            crate_level_code.push_str(&self.crates);366            if !self.crates.ends_with('\n') {367                crate_level_code.push('\n');368            }369        }370371        // Don't inject `extern crate std` because it's already injected by the372        // compiler.373        if !self.already_has_extern_crate &&374            !opts.no_crate_inject &&375            let Some(crate_name) = crate_name &&376            crate_name != "std" &&377            // Don't inject `extern crate` if the crate is never used.378            // NOTE: this is terribly inaccurate because it doesn't actually379            // parse the source, but only has false positives, not false380            // negatives.381            test_code.contains(crate_name)382        {383            // rustdoc implicitly inserts an `extern crate` item for the own crate384            // which may be unused, so we need to allow the lint.385            crate_level_code.push_str("#[allow(unused_extern_crates)]\n");386387            crate_level_code.push_str(&format!("extern crate r#{crate_name};\n"));388            line_offset += 1;389        }390391        // FIXME: This code cannot yet handle no_std test cases yet392        let wrapper = if dont_insert_main393            || self.has_main_fn394            || crate_level_code.contains("![no_std]")395        {396            None397        } else {398            let returns_result = processed_code.ends_with("(())");399            // Give each doctest main function a unique name.400            // This is for example needed for the tooling around `-C instrument-coverage`.401            let inner_fn_name = if let Some(ref test_id) = self.test_id {402                format!("_doctest_main_{test_id}")403            } else {404                "_inner".into()405            };406            let inner_attr = if self.test_id.is_some() { "#[allow(non_snake_case)] " } else { "" };407            let (main_pre, main_post) = if returns_result {408                (409                    format!(410                        "fn main() {{ {inner_attr}fn {inner_fn_name}() -> core::result::Result<(), impl core::fmt::Debug> {{\n",411                    ),412                    format!("\n}} {inner_fn_name}().unwrap() }}"),413                )414            } else if self.test_id.is_some() {415                (416                    format!("fn main() {{ {inner_attr}fn {inner_fn_name}() {{\n",),417                    format!("\n}} {inner_fn_name}() }}"),418                )419            } else {420                ("fn main() {\n".into(), "\n}".into())421            };422            // Note on newlines: We insert a line/newline *before*, and *after*423            // the doctest and adjust the `line_offset` accordingly.424            // In the case of `-C instrument-coverage`, this means that the generated425            // inner `main` function spans from the doctest opening codeblock to the426            // closing one. For example427            // /// ``` <- start of the inner main428            // /// <- code under doctest429            // /// ``` <- end of the inner main430            line_offset += 1;431432            Some(WrapperInfo {433                before: main_pre,434                after: main_post,435                returns_result,436                insert_indent_space: opts.insert_indent_space,437            })438        };439440        (441            DocTestWrapResult::Valid {442                code: processed_code.to_string(),443                wrapper,444                crate_level_code,445            },446            line_offset,447        )448    }449}450451fn reset_error_count(psess: &ParseSess) {452    // Reset errors so that they won't be reported as compiler bugs when dropping the453    // dcx. Any errors in the tests will be reported when the test file is compiled,454    // Note that we still need to cancel the errors above otherwise `Diag` will panic on455    // drop.456    psess.dcx().reset_err_count();457}458459const DOCTEST_CODE_WRAPPER: &str = "fn f(){";460461fn parse_source(462    source: &str,463    crate_name: &Option<&str>,464    parent_dcx: Option<DiagCtxtHandle<'_>>,465    span: Span,466    code_mappings: &[CodeLineMapping],467) -> Result<ParseSourceInfo, ()> {468    use rustc_errors::DiagCtxt;469    use rustc_errors::annotate_snippet_emitter_writer::AnnotateSnippetEmitter;470    use rustc_span::source_map::FilePathMapping;471472    let mut info =473        ParseSourceInfo { already_has_extern_crate: crate_name.is_none(), ..Default::default() };474475    let wrapped_source = format!("{DOCTEST_CODE_WRAPPER}{source}\n}}");476477    let filename = FileName::anon_source_code(&wrapped_source);478479    let sm = Arc::new(SourceMap::new(FilePathMapping::empty()));480    let supports_color = match get_stderr_color_choice(ColorConfig::Auto, &std::io::stderr()) {481        ColorChoice::Auto => unreachable!(),482        ColorChoice::AlwaysAnsi | ColorChoice::Always => true,483        ColorChoice::Never => false,484    };485    info.supports_color = supports_color;486    // Any errors in parsing should also appear when the doctest is compiled for real, so just487    // send all the errors that the parser emits directly into a `Sink` instead of stderr.488    let emitter = AnnotateSnippetEmitter::new(AutoStream::never(Box::new(io::sink())));489490    // FIXME(misdreavus): pass `-Z treat-err-as-bug` to the doctest parser491    let dcx = DiagCtxt::new(Box::new(emitter)).disable_warnings();492    let psess = ParseSess::with_dcx(dcx, sm);493494    // Don't strip any tokens; it wouldn't matter anyway because the source is wrapped in a function.495    let mut parser =496        match new_parser_from_source_str(&psess, filename, wrapped_source, StripTokens::Nothing) {497            Ok(p) => p,498            Err(errs) => {499                errs.into_iter().for_each(|err| err.cancel());500                reset_error_count(&psess);501                return Err(());502            }503        };504505    fn push_to_s(s: &mut String, source: &str, span: rustc_span::Span, prev_span_hi: &mut usize) {506        let extra_len = DOCTEST_CODE_WRAPPER.len();507        // We need to shift by the length of `DOCTEST_CODE_WRAPPER` because we508        // added it at the beginning of the source we provided to the parser.509        let mut hi = span.hi().0 as usize - extra_len;510        if hi > source.len() {511            hi = source.len();512        }513        s.push_str(&source[*prev_span_hi..hi]);514        *prev_span_hi = hi;515    }516517    fn span_in_doctest_source(span: Span, code_mappings: &[CodeLineMapping]) -> Option<Span> {518        let extra_len = DOCTEST_CODE_WRAPPER.len();519        let lo = (span.lo().0 as usize).checked_sub(extra_len)?;520        let hi = (span.hi().0 as usize).checked_sub(extra_len)?;521        if hi < lo {522            return None;523        }524        code_mappings.iter().find_map(|mapping| {525            if mapping.generated.start <= lo && hi <= mapping.generated.end {526                let start = lo - mapping.generated.start;527                let end = hi - mapping.generated.start;528                Some(mapping.original.from_inner(InnerSpan::new(start, end)))529            } else {530                None531            }532        })533    }534535    fn check_item(item: &ast::Item, info: &mut ParseSourceInfo, crate_name: &Option<&str>) -> bool {536        let mut is_extern_crate = false;537        if !info.has_global_allocator538            && item.attrs.iter().any(|attr| attr.has_name(sym::global_allocator))539        {540            info.has_global_allocator = true;541        }542        match item.kind {543            ast::ItemKind::Fn(ref fn_item) if !info.has_main_fn => {544                if fn_item.ident.name == sym::main {545                    info.has_main_fn = true;546                }547            }548            ast::ItemKind::ExternCrate(original, ident) => {549                is_extern_crate = true;550                if !info.already_has_extern_crate551                    && let Some(crate_name) = crate_name552                {553                    info.already_has_extern_crate = match original {554                        Some(name) => name.as_str() == *crate_name,555                        None => ident.as_str() == *crate_name,556                    };557                }558            }559            ast::ItemKind::MacroDef(..) => {560                info.has_macro_def = true;561            }562            _ => {}563        }564        is_extern_crate565    }566567    let mut prev_span_hi = 0;568    let not_crate_attrs = &[sym::forbid, sym::allow, sym::warn, sym::deny, sym::expect];569    let parsed = parser.parse_item(570        rustc_parse::parser::ForceCollect::No,571        rustc_parse::parser::AllowConstBlockItems::No,572    );573574    let result = match parsed {575        Ok(Some(ref item))576            if let ast::ItemKind::Fn(ref fn_item) = item.kind577                && let Some(ref body) = fn_item.body =>578        {579            for attr in &item.attrs {580                if attr.style == AttrStyle::Outer || attr.has_any_name(not_crate_attrs) {581                    // There is one exception to these attributes:582                    // `#![allow(internal_features)]`. If this attribute is used, we need to583                    // consider it only as a crate-level attribute.584                    if attr.has_name(sym::allow)585                        && let Some(list) = attr.meta_item_list()586                        && list.iter().any(|sub_attr| {587                            sub_attr.has_name(sym::internal_features)588                                || sub_attr.has_name(sym::incomplete_features)589                        })590                    {591                        push_to_s(&mut info.crate_attrs, source, attr.span, &mut prev_span_hi);592                    } else {593                        push_to_s(594                            &mut info.maybe_crate_attrs,595                            source,596                            attr.span,597                            &mut prev_span_hi,598                        );599                    }600                } else {601                    push_to_s(&mut info.crate_attrs, source, attr.span, &mut prev_span_hi);602                }603            }604            let mut has_non_items = false;605            let mut first_non_item_span = None;606            for stmt in &body.stmts {607                let mut is_extern_crate = false;608                match stmt.kind {609                    StmtKind::Item(ref item) => {610                        is_extern_crate = check_item(item, &mut info, crate_name);611                    }612                    // We assume that the macro calls will expand to item(s) even though they could613                    // expand to statements and expressions.614                    StmtKind::MacCall(ref mac_call) => {615                        if !info.has_main_fn {616                            // For backward compatibility, we look for the token sequence `fn main(…)`617                            // in the macro input (!) to crudely detect main functions "masked by a618                            // wrapper macro". For the record, this is a horrible heuristic!619                            // See <https://github.com/rust-lang/rust/issues/56898>.620                            let mut iter = mac_call.mac.args.tokens.iter();621                            while let Some(token) = iter.next() {622                                if let TokenTree::Token(token, _) = token623                                    && let TokenKind::Ident(kw::Fn, _) = token.kind624                                    && let Some(TokenTree::Token(ident, _)) = iter.peek()625                                    && let TokenKind::Ident(sym::main, _) = ident.kind626                                    && let Some(TokenTree::Delimited(.., Delimiter::Parenthesis, _)) = {627                                        iter.next();628                                        iter.peek()629                                    }630                                {631                                    info.has_main_fn = true;632                                    break;633                                }634                            }635                        }636                    }637                    StmtKind::Expr(ref expr) => {638                        if matches!(expr.kind, ast::ExprKind::Err(_)) {639                            reset_error_count(&psess);640                            return Err(());641                        }642                        has_non_items = true;643                        first_non_item_span.get_or_insert(stmt.span);644                    }645                    StmtKind::Let(_) | StmtKind::Semi(_) | StmtKind::Empty => {646                        has_non_items = true;647                        first_non_item_span.get_or_insert(stmt.span);648                    }649                }650651                // Weirdly enough, the `Stmt` span doesn't include its attributes, so we need to652                // tweak the span to include the attributes as well.653                let mut span = stmt.span;654                if let Some(attr) =655                    stmt.kind.attrs().iter().find(|attr| attr.style == AttrStyle::Outer)656                {657                    span = span.with_lo(attr.span.lo());658                }659                if info.everything_else.is_empty()660                    && (!info.maybe_crate_attrs.is_empty() || !info.crate_attrs.is_empty())661                {662                    // To keep the doctest code "as close as possible" to the original, we insert663                    // all the code located between this new span and the previous span which664                    // might contain code comments and backlines.665                    push_to_s(&mut info.crates, source, span.shrink_to_lo(), &mut prev_span_hi);666                }667                if !is_extern_crate {668                    push_to_s(&mut info.everything_else, source, span, &mut prev_span_hi);669                } else {670                    push_to_s(&mut info.crates, source, span, &mut prev_span_hi);671                }672            }673            if has_non_items {674                let warning_span = first_non_item_span675                    .and_then(|span| span_in_doctest_source(span, code_mappings))676                    .unwrap_or(span);677                if info.has_main_fn678                    && let Some(dcx) = parent_dcx679                    && !warning_span.is_dummy()680                {681                    dcx.span_warn(682                        warning_span,683                        "the `main` function of this doctest won't be run as it contains \684                         expressions at the top level, meaning that the whole doctest code will be \685                         wrapped in a function",686                    );687                }688                info.has_main_fn = false;689            }690            Ok(info)691        }692        Err(e) => {693            e.cancel();694            Err(())695        }696        _ => Err(()),697    };698699    reset_error_count(&psess);700    result701}