1use std::mem;
4
5use rustc_ast as ast;
6use rustc_ast::entry::EntryPointType;
7use rustc_ast::mut_visit::*;
8use rustc_ast::ptr::P;
9use rustc_ast::visit::{Visitor, walk_item};
10use rustc_ast::{ModKind, attr};
11use rustc_errors::DiagCtxtHandle;
12use rustc_expand::base::{ExtCtxt, ResolverExpand};
13use rustc_expand::expand::{AstFragment, ExpansionConfig};
14use rustc_feature::Features;
15use rustc_lint_defs::BuiltinLintDiag;
16use rustc_session::Session;
17use rustc_session::lint::builtin::UNNAMEABLE_TEST_ITEMS;
18use rustc_span::hygiene::{AstPass, SyntaxContext, Transparency};
19use rustc_span::{DUMMY_SP, Ident, Span, Symbol, sym};
20use rustc_target::spec::PanicStrategy;
21use smallvec::smallvec;
22use thin_vec::{ThinVec, thin_vec};
23use tracing::debug;
24
25use crate::errors;
26
27#[derive(Clone)]
28struct Test {
29 span: Span,
30 ident: Ident,
31 name: Symbol,
32}
33
34struct TestCtxt<'a> {
35 ext_cx: ExtCtxt<'a>,
36 panic_strategy: PanicStrategy,
37 def_site: Span,
38 test_cases: Vec<Test>,
39 reexport_test_harness_main: Option<Symbol>,
40 test_runner: Option<ast::Path>,
41}
42
43pub fn inject(
46 krate: &mut ast::Crate,
47 sess: &Session,
48 features: &Features,
49 resolver: &mut dyn ResolverExpand,
50) {
51 let dcx = sess.dcx();
52 let panic_strategy = sess.panic_strategy();
53 let platform_panic_strategy = sess.target.panic_strategy;
54
55 let reexport_test_harness_main =
60 attr::first_attr_value_str_by_name(&krate.attrs, sym::reexport_test_harness_main);
61
62 let test_runner = get_test_runner(dcx, krate);
65
66 if sess.is_test_crate() {
67 let panic_strategy = match (panic_strategy, sess.opts.unstable_opts.panic_abort_tests) {
68 (PanicStrategy::Abort, true) => PanicStrategy::Abort,
69 (PanicStrategy::Abort, false) => {
70 if panic_strategy == platform_panic_strategy {
71 } else {
74 dcx.emit_err(errors::TestsNotSupport {});
75 }
76 PanicStrategy::Unwind
77 }
78 (PanicStrategy::Unwind, _) => PanicStrategy::Unwind,
79 };
80 generate_test_harness(
81 sess,
82 resolver,
83 reexport_test_harness_main,
84 krate,
85 features,
86 panic_strategy,
87 test_runner,
88 )
89 }
90}
91
92struct TestHarnessGenerator<'a> {
93 cx: TestCtxt<'a>,
94 tests: Vec<Test>,
95}
96
97impl TestHarnessGenerator<'_> {
98 fn add_test_cases(&mut self, node_id: ast::NodeId, span: Span, prev_tests: Vec<Test>) {
99 let mut tests = mem::replace(&mut self.tests, prev_tests);
100
101 if !tests.is_empty() {
102 let expn_id = self.cx.ext_cx.resolver.expansion_for_ast_pass(
105 span,
106 AstPass::TestHarness,
107 &[],
108 Some(node_id),
109 );
110 for test in &mut tests {
111 test.ident.span =
114 test.ident.span.apply_mark(expn_id.to_expn_id(), Transparency::Opaque);
115 }
116 self.cx.test_cases.extend(tests);
117 }
118 }
119}
120
121impl<'a> MutVisitor for TestHarnessGenerator<'a> {
122 fn visit_crate(&mut self, c: &mut ast::Crate) {
123 let prev_tests = mem::take(&mut self.tests);
124 walk_crate(self, c);
125 self.add_test_cases(ast::CRATE_NODE_ID, c.spans.inner_span, prev_tests);
126
127 c.items.push(mk_main(&mut self.cx));
129 }
130
131 fn visit_item(&mut self, item: &mut P<ast::Item>) {
132 let item = &mut **item;
133
134 if let Some(name) = get_test_name(&item) {
135 debug!("this is a test item");
136
137 let test = Test { span: item.span, ident: item.ident, name };
138 self.tests.push(test);
139 }
140
141 if let ast::ItemKind::Mod(
144 _,
145 ModKind::Loaded(.., ast::ModSpans { inner_span: span, .. }, _),
146 ) = item.kind
147 {
148 let prev_tests = mem::take(&mut self.tests);
149 walk_item_kind(
150 &mut item.kind,
151 item.span,
152 item.id,
153 &mut item.ident,
154 &mut item.vis,
155 (),
156 self,
157 );
158 self.add_test_cases(item.id, span, prev_tests);
159 } else {
160 walk_item(&mut InnerItemLinter { sess: self.cx.ext_cx.sess }, &item);
162 }
163 }
164}
165
166struct InnerItemLinter<'a> {
167 sess: &'a Session,
168}
169
170impl<'a> Visitor<'a> for InnerItemLinter<'_> {
171 fn visit_item(&mut self, i: &'a ast::Item) {
172 if let Some(attr) = attr::find_by_name(&i.attrs, sym::rustc_test_marker) {
173 self.sess.psess.buffer_lint(
174 UNNAMEABLE_TEST_ITEMS,
175 attr.span,
176 i.id,
177 BuiltinLintDiag::UnnameableTestItems,
178 );
179 }
180 }
181}
182
183fn entry_point_type(item: &ast::Item, at_root: bool) -> EntryPointType {
184 match item.kind {
185 ast::ItemKind::Fn(..) => {
186 rustc_ast::entry::entry_point_type(&item.attrs, at_root, Some(item.ident.name))
187 }
188 _ => EntryPointType::None,
189 }
190}
191
192struct EntryPointCleaner<'a> {
195 sess: &'a Session,
197 depth: usize,
198 def_site: Span,
199}
200
201impl<'a> MutVisitor for EntryPointCleaner<'a> {
202 fn visit_item(&mut self, item: &mut P<ast::Item>) {
203 self.depth += 1;
204 ast::mut_visit::walk_item(self, item);
205 self.depth -= 1;
206
207 match entry_point_type(&item, self.depth == 0) {
211 EntryPointType::MainNamed | EntryPointType::RustcMainAttr => {
212 let allow_dead_code = attr::mk_attr_nested_word(
213 &self.sess.psess.attr_id_generator,
214 ast::AttrStyle::Outer,
215 ast::Safety::Default,
216 sym::allow,
217 sym::dead_code,
218 self.def_site,
219 );
220 item.attrs.retain(|attr| !attr.has_name(sym::rustc_main));
221 item.attrs.push(allow_dead_code);
222 }
223 EntryPointType::None | EntryPointType::OtherMain => {}
224 };
225 }
226}
227
228fn generate_test_harness(
230 sess: &Session,
231 resolver: &mut dyn ResolverExpand,
232 reexport_test_harness_main: Option<Symbol>,
233 krate: &mut ast::Crate,
234 features: &Features,
235 panic_strategy: PanicStrategy,
236 test_runner: Option<ast::Path>,
237) {
238 let econfig = ExpansionConfig::default("test".to_string(), features);
239 let ext_cx = ExtCtxt::new(sess, econfig, resolver, None);
240
241 let expn_id = ext_cx.resolver.expansion_for_ast_pass(
242 DUMMY_SP,
243 AstPass::TestHarness,
244 &[sym::test, sym::rustc_attrs, sym::coverage_attribute],
245 None,
246 );
247 let def_site = DUMMY_SP.with_def_site_ctxt(expn_id.to_expn_id());
248
249 let mut cleaner = EntryPointCleaner { sess, depth: 0, def_site };
251 cleaner.visit_crate(krate);
252
253 let cx = TestCtxt {
254 ext_cx,
255 panic_strategy,
256 def_site,
257 test_cases: Vec::new(),
258 reexport_test_harness_main,
259 test_runner,
260 };
261
262 TestHarnessGenerator { cx, tests: Vec::new() }.visit_crate(krate);
263}
264
265fn mk_main(cx: &mut TestCtxt<'_>) -> P<ast::Item> {
296 let sp = cx.def_site;
297 let ecx = &cx.ext_cx;
298 let test_id = Ident::new(sym::test, sp);
299
300 let runner_name = match cx.panic_strategy {
301 PanicStrategy::Unwind => "test_main_static",
302 PanicStrategy::Abort => "test_main_static_abort",
303 };
304
305 let mut test_runner = cx
307 .test_runner
308 .clone()
309 .unwrap_or_else(|| ecx.path(sp, vec![test_id, Ident::from_str_and_span(runner_name, sp)]));
310
311 test_runner.span = sp;
312
313 let test_main_path_expr = ecx.expr_path(test_runner);
314 let call_test_main = ecx.expr_call(sp, test_main_path_expr, thin_vec![mk_tests_slice(cx, sp)]);
315 let call_test_main = ecx.stmt_expr(call_test_main);
316
317 let test_extern_stmt = ecx.stmt_item(
319 sp,
320 ecx.item(sp, test_id, ast::AttrVec::new(), ast::ItemKind::ExternCrate(None)),
321 );
322
323 let main_attr = ecx.attr_word(sym::rustc_main, sp);
325 let coverage_attr = ecx.attr_nested_word(sym::coverage, sym::off, sp);
327 let doc_hidden_attr = ecx.attr_nested_word(sym::doc, sym::hidden, sp);
329
330 let main_ret_ty = ecx.ty(sp, ast::TyKind::Tup(ThinVec::new()));
332
333 let main_body = if cx.test_runner.is_none() {
335 ecx.block(sp, thin_vec![test_extern_stmt, call_test_main])
336 } else {
337 ecx.block(sp, thin_vec![call_test_main])
338 };
339
340 let decl = ecx.fn_decl(ThinVec::new(), ast::FnRetTy::Ty(main_ret_ty));
341 let sig = ast::FnSig { decl, header: ast::FnHeader::default(), span: sp };
342 let defaultness = ast::Defaultness::Final;
343 let main = ast::ItemKind::Fn(Box::new(ast::Fn {
344 defaultness,
345 sig,
346 generics: ast::Generics::default(),
347 contract: None,
348 body: Some(main_body),
349 }));
350
351 let main_id = match cx.reexport_test_harness_main {
353 Some(sym) => Ident::new(sym, sp.with_ctxt(SyntaxContext::root())),
354 None => Ident::new(sym::main, sp),
355 };
356
357 let main = P(ast::Item {
358 ident: main_id,
359 attrs: thin_vec![main_attr, coverage_attr, doc_hidden_attr],
360 id: ast::DUMMY_NODE_ID,
361 kind: main,
362 vis: ast::Visibility { span: sp, kind: ast::VisibilityKind::Public, tokens: None },
363 span: sp,
364 tokens: None,
365 });
366
367 let main = AstFragment::Items(smallvec![main]);
369 cx.ext_cx.monotonic_expander().fully_expand_fragment(main).make_items().pop().unwrap()
370}
371
372fn mk_tests_slice(cx: &TestCtxt<'_>, sp: Span) -> P<ast::Expr> {
375 debug!("building test vector from {} tests", cx.test_cases.len());
376 let ecx = &cx.ext_cx;
377
378 let mut tests = cx.test_cases.clone();
379 tests.sort_by(|a, b| a.name.as_str().cmp(b.name.as_str()));
380
381 ecx.expr_array_ref(
382 sp,
383 tests
384 .iter()
385 .map(|test| {
386 ecx.expr_addr_of(test.span, ecx.expr_path(ecx.path(test.span, vec![test.ident])))
387 })
388 .collect(),
389 )
390}
391
392fn get_test_name(i: &ast::Item) -> Option<Symbol> {
393 attr::first_attr_value_str_by_name(&i.attrs, sym::rustc_test_marker)
394}
395
396fn get_test_runner(dcx: DiagCtxtHandle<'_>, krate: &ast::Crate) -> Option<ast::Path> {
397 let test_attr = attr::find_by_name(&krate.attrs, sym::test_runner)?;
398 let meta_list = test_attr.meta_item_list()?;
399 let span = test_attr.span;
400 match &*meta_list {
401 [single] => match single.meta_item() {
402 Some(meta_item) if meta_item.is_word() => return Some(meta_item.path.clone()),
403 _ => {
404 dcx.emit_err(errors::TestRunnerInvalid { span });
405 }
406 },
407 _ => {
408 dcx.emit_err(errors::TestRunnerNargs { span });
409 }
410 }
411 None
412}