1//! Logic for transforming the raw code given by the user into something actually
2//! runnable, e.g. by adding a `main` function if it doesn't already exist.
3
4use std::fmt::{self, Write as _};
5use std::io;
6use std::sync::Arc;
7
8use rustc_ast::token::{Delimiter, TokenKind};
9use rustc_ast::tokenstream::TokenTree;
10use rustc_ast::{self as ast, AttrStyle, HasAttrs, StmtKind};
11use rustc_errors::ColorConfig;
12use rustc_errors::emitter::stderr_destination;
13use rustc_parse::new_parser_from_source_str;
14use rustc_session::parse::ParseSess;
15use rustc_span::edition::Edition;
16use rustc_span::source_map::SourceMap;
17use rustc_span::symbol::sym;
18use rustc_span::{FileName, kw};
19use tracing::debug;
20
21use super::GlobalTestOptions;
22use crate::display::Joined as _;
23use crate::html::markdown::LangString;
24
25#[derive(Default)]
26struct ParseSourceInfo {
27 has_main_fn: bool,
28 already_has_extern_crate: bool,
29 supports_color: bool,
30 has_global_allocator: bool,
31 has_macro_def: bool,
32 everything_else: String,
33 crates: String,
34 crate_attrs: String,
35 maybe_crate_attrs: String,
36}
37
38/// This struct contains information about the doctest itself which is then used to generate
39/// doctest source code appropriately.
40pub(crate) struct DocTestBuilder {
41 pub(crate) supports_color: bool,
42 pub(crate) already_has_extern_crate: bool,
43 pub(crate) has_main_fn: bool,
44 pub(crate) crate_attrs: String,
45 /// If this is a merged doctest, it will be put into `everything_else`, otherwise it will
46 /// put into `crate_attrs`.
47 pub(crate) maybe_crate_attrs: String,
48 pub(crate) crates: String,
49 pub(crate) everything_else: String,
50 pub(crate) test_id: Option<String>,
51 pub(crate) invalid_ast: bool,
52 pub(crate) can_be_merged: bool,
53}
54
55impl DocTestBuilder {
56 pub(crate) fn new(
57 source: &str,
58 crate_name: Option<&str>,
59 edition: Edition,
60 can_merge_doctests: bool,
61 // If `test_id` is `None`, it means we're generating code for a code example "run" link.
62 test_id: Option<String>,
63 lang_str: Option<&LangString>,
64 ) -> Self {
65 let can_merge_doctests = can_merge_doctests
66 && lang_str.is_some_and(|lang_str| {
67 !lang_str.compile_fail && !lang_str.test_harness && !lang_str.standalone_crate
68 });
69
70 let result = rustc_driver::catch_fatal_errors(|| {
71 rustc_span::create_session_if_not_set_then(edition, |_| {
72 parse_source(source, &crate_name)
73 })
74 });
75
76 let Ok(Ok(ParseSourceInfo {
77 has_main_fn,
78 already_has_extern_crate,
79 supports_color,
80 has_global_allocator,
81 has_macro_def,
82 everything_else,
83 crates,
84 crate_attrs,
85 maybe_crate_attrs,
86 })) = result
87 else {
88 // If the AST returned an error, we don't want this doctest to be merged with the
89 // others.
90 return Self::invalid(
91 String::new(),
92 String::new(),
93 String::new(),
94 source.to_string(),
95 test_id,
96 );
97 };
98
99 debug!("crate_attrs:\n{crate_attrs}{maybe_crate_attrs}");
100 debug!("crates:\n{crates}");
101 debug!("after:\n{everything_else}");
102
103 // If it contains `#[feature]` or `#[no_std]`, we don't want it to be merged either.
104 let can_be_merged = can_merge_doctests
105 && !has_global_allocator
106 && crate_attrs.is_empty()
107 // If this is a merged doctest and a defined macro uses `$crate`, then the path will
108 // not work, so better not put it into merged doctests.
109 && !(has_macro_def && everything_else.contains("$crate"));
110 Self {
111 supports_color,
112 has_main_fn,
113 crate_attrs,
114 maybe_crate_attrs,
115 crates,
116 everything_else,
117 already_has_extern_crate,
118 test_id,
119 invalid_ast: false,
120 can_be_merged,
121 }
122 }
123
124 fn invalid(
125 crate_attrs: String,
126 maybe_crate_attrs: String,
127 crates: String,
128 everything_else: String,
129 test_id: Option<String>,
130 ) -> Self {
131 Self {
132 supports_color: false,
133 has_main_fn: false,
134 crate_attrs,
135 maybe_crate_attrs,
136 crates,
137 everything_else,
138 already_has_extern_crate: false,
139 test_id,
140 invalid_ast: true,
141 can_be_merged: false,
142 }
143 }
144
145 /// Transforms a test into code that can be compiled into a Rust binary, and returns the number of
146 /// lines before the test code begins.
147 pub(crate) fn generate_unique_doctest(
148 &self,
149 test_code: &str,
150 dont_insert_main: bool,
151 opts: &GlobalTestOptions,
152 crate_name: Option<&str>,
153 ) -> (String, usize) {
154 if self.invalid_ast {
155 // If the AST failed to compile, no need to go generate a complete doctest, the error
156 // will be better this way.
157 debug!("invalid AST:\n{test_code}");
158 return (test_code.to_string(), 0);
159 }
160 let mut line_offset = 0;
161 let mut prog = String::new();
162 let everything_else = self.everything_else.trim();
163 if opts.attrs.is_empty() {
164 // If there aren't any attributes supplied by #![doc(test(attr(...)))], then allow some
165 // lints that are commonly triggered in doctests. The crate-level test attributes are
166 // commonly used to make tests fail in case they trigger warnings, so having this there in
167 // that case may cause some tests to pass when they shouldn't have.
168 prog.push_str("#![allow(unused)]\n");
169 line_offset += 1;
170 }
171
172 // Next, any attributes that came from the crate root via #![doc(test(attr(...)))].
173 for attr in &opts.attrs {
174 prog.push_str(&format!("#![{attr}]\n"));
175 line_offset += 1;
176 }
177
178 // Now push any outer attributes from the example, assuming they
179 // are intended to be crate attributes.
180 if !self.crate_attrs.is_empty() {
181 prog.push_str(&self.crate_attrs);
182 if !self.crate_attrs.ends_with('\n') {
183 prog.push('\n');
184 }
185 }
186 if !self.maybe_crate_attrs.is_empty() {
187 prog.push_str(&self.maybe_crate_attrs);
188 if !self.maybe_crate_attrs.ends_with('\n') {
189 prog.push('\n');
190 }
191 }
192 if !self.crates.is_empty() {
193 prog.push_str(&self.crates);
194 if !self.crates.ends_with('\n') {
195 prog.push('\n');
196 }
197 }
198
199 // Don't inject `extern crate std` because it's already injected by the
200 // compiler.
201 if !self.already_has_extern_crate &&
202 !opts.no_crate_inject &&
203 let Some(crate_name) = crate_name &&
204 crate_name != "std" &&
205 // Don't inject `extern crate` if the crate is never used.
206 // NOTE: this is terribly inaccurate because it doesn't actually
207 // parse the source, but only has false positives, not false
208 // negatives.
209 test_code.contains(crate_name)
210 {
211 // rustdoc implicitly inserts an `extern crate` item for the own crate
212 // which may be unused, so we need to allow the lint.
213 prog.push_str("#[allow(unused_extern_crates)]\n");
214
215 prog.push_str(&format!("extern crate r#{crate_name};\n"));
216 line_offset += 1;
217 }
218
219 // FIXME: This code cannot yet handle no_std test cases yet
220 if dont_insert_main || self.has_main_fn || prog.contains("![no_std]") {
221 prog.push_str(everything_else);
222 } else {
223 let returns_result = everything_else.ends_with("(())");
224 // Give each doctest main function a unique name.
225 // This is for example needed for the tooling around `-C instrument-coverage`.
226 let inner_fn_name = if let Some(ref test_id) = self.test_id {
227 format!("_doctest_main_{test_id}")
228 } else {
229 "_inner".into()
230 };
231 let inner_attr = if self.test_id.is_some() { "#[allow(non_snake_case)] " } else { "" };
232 let (main_pre, main_post) = if returns_result {
233 (
234 format!(
235 "fn main() {{ {inner_attr}fn {inner_fn_name}() -> core::result::Result<(), impl core::fmt::Debug> {{\n",
236 ),
237 format!("\n}} {inner_fn_name}().unwrap() }}"),
238 )
239 } else if self.test_id.is_some() {
240 (
241 format!("fn main() {{ {inner_attr}fn {inner_fn_name}() {{\n",),
242 format!("\n}} {inner_fn_name}() }}"),
243 )
244 } else {
245 ("fn main() {\n".into(), "\n}".into())
246 };
247 // Note on newlines: We insert a line/newline *before*, and *after*
248 // the doctest and adjust the `line_offset` accordingly.
249 // In the case of `-C instrument-coverage`, this means that the generated
250 // inner `main` function spans from the doctest opening codeblock to the
251 // closing one. For example
252 // /// ``` <- start of the inner main
253 // /// <- code under doctest
254 // /// ``` <- end of the inner main
255 line_offset += 1;
256
257 prog.push_str(&main_pre);
258
259 // add extra 4 spaces for each line to offset the code block
260 if opts.insert_indent_space {
261 write!(
262 prog,
263 "{}",
264 fmt::from_fn(|f| everything_else
265 .lines()
266 .map(|line| fmt::from_fn(move |f| write!(f, " {line}")))
267 .joined("\n", f))
268 )
269 .unwrap();
270 } else {
271 prog.push_str(everything_else);
272 };
273 prog.push_str(&main_post);
274 }
275
276 debug!("final doctest:\n{prog}");
277
278 (prog, line_offset)
279 }
280}
281
282fn reset_error_count(psess: &ParseSess) {
283 // Reset errors so that they won't be reported as compiler bugs when dropping the
284 // dcx. Any errors in the tests will be reported when the test file is compiled,
285 // Note that we still need to cancel the errors above otherwise `Diag` will panic on
286 // drop.
287 psess.dcx().reset_err_count();
288}
289
290const DOCTEST_CODE_WRAPPER: &str = "fn f(){";
291
292fn parse_source(source: &str, crate_name: &Option<&str>) -> Result<ParseSourceInfo, ()> {
293 use rustc_errors::DiagCtxt;
294 use rustc_errors::emitter::{Emitter, HumanEmitter};
295 use rustc_span::source_map::FilePathMapping;
296
297 let mut info =
298 ParseSourceInfo { already_has_extern_crate: crate_name.is_none(), ..Default::default() };
299
300 let wrapped_source = format!("{DOCTEST_CODE_WRAPPER}{source}\n}}");
301
302 let filename = FileName::anon_source_code(&wrapped_source);
303
304 let sm = Arc::new(SourceMap::new(FilePathMapping::empty()));
305 let fallback_bundle = rustc_errors::fallback_fluent_bundle(
306 rustc_driver::DEFAULT_LOCALE_RESOURCES.to_vec(),
307 false,
308 );
309 info.supports_color =
310 HumanEmitter::new(stderr_destination(ColorConfig::Auto), fallback_bundle.clone())
311 .supports_color();
312 // Any errors in parsing should also appear when the doctest is compiled for real, so just
313 // send all the errors that the parser emits directly into a `Sink` instead of stderr.
314 let emitter = HumanEmitter::new(Box::new(io::sink()), fallback_bundle);
315
316 // FIXME(misdreavus): pass `-Z treat-err-as-bug` to the doctest parser
317 let dcx = DiagCtxt::new(Box::new(emitter)).disable_warnings();
318 let psess = ParseSess::with_dcx(dcx, sm);
319
320 let mut parser = match new_parser_from_source_str(&psess, filename, wrapped_source) {
321 Ok(p) => p,
322 Err(errs) => {
323 errs.into_iter().for_each(|err| err.cancel());
324 reset_error_count(&psess);
325 return Err(());
326 }
327 };
328
329 fn push_to_s(s: &mut String, source: &str, span: rustc_span::Span, prev_span_hi: &mut usize) {
330 let extra_len = DOCTEST_CODE_WRAPPER.len();
331 // We need to shift by the length of `DOCTEST_CODE_WRAPPER` because we
332 // added it at the beginning of the source we provided to the parser.
333 let mut hi = span.hi().0 as usize - extra_len;
334 if hi > source.len() {
335 hi = source.len();
336 }
337 s.push_str(&source[*prev_span_hi..hi]);
338 *prev_span_hi = hi;
339 }
340
341 fn check_item(item: &ast::Item, info: &mut ParseSourceInfo, crate_name: &Option<&str>) -> bool {
342 let mut is_extern_crate = false;
343 if !info.has_global_allocator
344 && item.attrs.iter().any(|attr| attr.name_or_empty() == sym::global_allocator)
345 {
346 info.has_global_allocator = true;
347 }
348 match item.kind {
349 ast::ItemKind::Fn(_) if !info.has_main_fn => {
350 if item.ident.name == sym::main {
351 info.has_main_fn = true;
352 }
353 }
354 ast::ItemKind::ExternCrate(original) => {
355 is_extern_crate = true;
356 if !info.already_has_extern_crate
357 && let Some(crate_name) = crate_name
358 {
359 info.already_has_extern_crate = match original {
360 Some(name) => name.as_str() == *crate_name,
361 None => item.ident.as_str() == *crate_name,
362 };
363 }
364 }
365 ast::ItemKind::MacroDef(..) => {
366 info.has_macro_def = true;
367 }
368 _ => {}
369 }
370 is_extern_crate
371 }
372
373 let mut prev_span_hi = 0;
374 let not_crate_attrs = [sym::forbid, sym::allow, sym::warn, sym::deny, sym::expect];
375 let parsed = parser.parse_item(rustc_parse::parser::ForceCollect::No);
376
377 let result = match parsed {
378 Ok(Some(ref item))
379 if let ast::ItemKind::Fn(ref fn_item) = item.kind
380 && let Some(ref body) = fn_item.body =>
381 {
382 for attr in &item.attrs {
383 let attr_name = attr.name_or_empty();
384
385 if attr.style == AttrStyle::Outer || not_crate_attrs.contains(&attr_name) {
386 // There is one exception to these attributes:
387 // `#![allow(internal_features)]`. If this attribute is used, we need to
388 // consider it only as a crate-level attribute.
389 if attr_name == sym::allow
390 && let Some(list) = attr.meta_item_list()
391 && list.iter().any(|sub_attr| {
392 sub_attr.name_or_empty().as_str() == "internal_features"
393 })
394 {
395 push_to_s(&mut info.crate_attrs, source, attr.span, &mut prev_span_hi);
396 } else {
397 push_to_s(
398 &mut info.maybe_crate_attrs,
399 source,
400 attr.span,
401 &mut prev_span_hi,
402 );
403 }
404 } else {
405 push_to_s(&mut info.crate_attrs, source, attr.span, &mut prev_span_hi);
406 }
407 }
408 let mut has_non_items = false;
409 for stmt in &body.stmts {
410 let mut is_extern_crate = false;
411 match stmt.kind {
412 StmtKind::Item(ref item) => {
413 is_extern_crate = check_item(item, &mut info, crate_name);
414 }
415 // We assume that the macro calls will expand to item(s) even though they could
416 // expand to statements and expressions.
417 StmtKind::MacCall(ref mac_call) => {
418 if !info.has_main_fn {
419 // For backward compatibility, we look for the token sequence `fn main(…)`
420 // in the macro input (!) to crudely detect main functions "masked by a
421 // wrapper macro". For the record, this is a horrible heuristic!
422 // See <https://github.com/rust-lang/rust/issues/56898>.
423 let mut iter = mac_call.mac.args.tokens.iter();
424 while let Some(token) = iter.next() {
425 if let TokenTree::Token(token, _) = token
426 && let TokenKind::Ident(kw::Fn, _) = token.kind
427 && let Some(TokenTree::Token(ident, _)) = iter.peek()
428 && let TokenKind::Ident(sym::main, _) = ident.kind
429 && let Some(TokenTree::Delimited(.., Delimiter::Parenthesis, _)) = {
430 iter.next();
431 iter.peek()
432 }
433 {
434 info.has_main_fn = true;
435 break;
436 }
437 }
438 }
439 }
440 StmtKind::Expr(ref expr) => {
441 if matches!(expr.kind, ast::ExprKind::Err(_)) {
442 reset_error_count(&psess);
443 return Err(());
444 }
445 has_non_items = true;
446 }
447 StmtKind::Let(_) | StmtKind::Semi(_) | StmtKind::Empty => has_non_items = true,
448 }
449
450 // Weirdly enough, the `Stmt` span doesn't include its attributes, so we need to
451 // tweak the span to include the attributes as well.
452 let mut span = stmt.span;
453 if let Some(attr) =
454 stmt.kind.attrs().iter().find(|attr| attr.style == AttrStyle::Outer)
455 {
456 span = span.with_lo(attr.span.lo());
457 }
458 if info.everything_else.is_empty()
459 && (!info.maybe_crate_attrs.is_empty() || !info.crate_attrs.is_empty())
460 {
461 // To keep the doctest code "as close as possible" to the original, we insert
462 // all the code located between this new span and the previous span which
463 // might contain code comments and backlines.
464 push_to_s(&mut info.crates, source, span.shrink_to_lo(), &mut prev_span_hi);
465 }
466 if !is_extern_crate {
467 push_to_s(&mut info.everything_else, source, span, &mut prev_span_hi);
468 } else {
469 push_to_s(&mut info.crates, source, span, &mut prev_span_hi);
470 }
471 }
472 if has_non_items {
473 // FIXME: if `info.has_main_fn` is `true`, emit a warning here to mention that
474 // this code will not be called.
475 info.has_main_fn = false;
476 }
477 Ok(info)
478 }
479 Err(e) => {
480 e.cancel();
481 Err(())
482 }
483 _ => Err(()),
484 };
485
486 reset_error_count(&psess);
487 result
488}
489