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.kind.ident().unwrap(), name };
139 self.tests.push(test);
140 }
141
142 if let ast::ItemKind::Mod(
145 _,
146 _,
147 ModKind::Loaded(.., ast::ModSpans { inner_span: span, .. }, _),
148 ) = item.kind
149 {
150 let prev_tests = mem::take(&mut self.tests);
151 walk_item_kind(&mut item.kind, item.span, item.id, &mut item.vis, (), self);
152 self.add_test_cases(item.id, span, prev_tests);
153 } else {
154 walk_item(&mut InnerItemLinter { sess: self.cx.ext_cx.sess }, &item);
156 }
157 }
158}
159
160struct InnerItemLinter<'a> {
161 sess: &'a Session,
162}
163
164impl<'a> Visitor<'a> for InnerItemLinter<'_> {
165 fn visit_item(&mut self, i: &'a ast::Item) {
166 if let Some(attr) = attr::find_by_name(&i.attrs, sym::rustc_test_marker) {
167 self.sess.psess.buffer_lint(
168 UNNAMEABLE_TEST_ITEMS,
169 attr.span,
170 i.id,
171 BuiltinLintDiag::UnnameableTestItems,
172 );
173 }
174 }
175}
176
177fn entry_point_type(item: &ast::Item, at_root: bool) -> EntryPointType {
178 match &item.kind {
179 ast::ItemKind::Fn(fn_) => {
180 rustc_ast::entry::entry_point_type(&item.attrs, at_root, Some(fn_.ident.name))
181 }
182 _ => EntryPointType::None,
183 }
184}
185
186struct EntryPointCleaner<'a> {
189 sess: &'a Session,
191 depth: usize,
192 def_site: Span,
193}
194
195impl<'a> MutVisitor for EntryPointCleaner<'a> {
196 fn visit_item(&mut self, item: &mut P<ast::Item>) {
197 self.depth += 1;
198 ast::mut_visit::walk_item(self, item);
199 self.depth -= 1;
200
201 match entry_point_type(&item, self.depth == 0) {
205 EntryPointType::MainNamed | EntryPointType::RustcMainAttr => {
206 let allow_dead_code = attr::mk_attr_nested_word(
207 &self.sess.psess.attr_id_generator,
208 ast::AttrStyle::Outer,
209 ast::Safety::Default,
210 sym::allow,
211 sym::dead_code,
212 self.def_site,
213 );
214 item.attrs.retain(|attr| !attr.has_name(sym::rustc_main));
215 item.attrs.push(allow_dead_code);
216 }
217 EntryPointType::None | EntryPointType::OtherMain => {}
218 };
219 }
220}
221
222fn generate_test_harness(
224 sess: &Session,
225 resolver: &mut dyn ResolverExpand,
226 reexport_test_harness_main: Option<Symbol>,
227 krate: &mut ast::Crate,
228 features: &Features,
229 panic_strategy: PanicStrategy,
230 test_runner: Option<ast::Path>,
231) {
232 let econfig = ExpansionConfig::default("test".to_string(), features);
233 let ext_cx = ExtCtxt::new(sess, econfig, resolver, None);
234
235 let expn_id = ext_cx.resolver.expansion_for_ast_pass(
236 DUMMY_SP,
237 AstPass::TestHarness,
238 &[sym::test, sym::rustc_attrs, sym::coverage_attribute],
239 None,
240 );
241 let def_site = DUMMY_SP.with_def_site_ctxt(expn_id.to_expn_id());
242
243 let mut cleaner = EntryPointCleaner { sess, depth: 0, def_site };
245 cleaner.visit_crate(krate);
246
247 let cx = TestCtxt {
248 ext_cx,
249 panic_strategy,
250 def_site,
251 test_cases: Vec::new(),
252 reexport_test_harness_main,
253 test_runner,
254 };
255
256 TestHarnessGenerator { cx, tests: Vec::new() }.visit_crate(krate);
257}
258
259fn mk_main(cx: &mut TestCtxt<'_>) -> P<ast::Item> {
290 let sp = cx.def_site;
291 let ecx = &cx.ext_cx;
292 let test_ident = Ident::new(sym::test, sp);
293
294 let runner_name = match cx.panic_strategy {
295 PanicStrategy::Unwind => "test_main_static",
296 PanicStrategy::Abort => "test_main_static_abort",
297 };
298
299 let mut test_runner = cx.test_runner.clone().unwrap_or_else(|| {
301 ecx.path(sp, vec![test_ident, Ident::from_str_and_span(runner_name, sp)])
302 });
303
304 test_runner.span = sp;
305
306 let test_main_path_expr = ecx.expr_path(test_runner);
307 let call_test_main = ecx.expr_call(sp, test_main_path_expr, thin_vec![mk_tests_slice(cx, sp)]);
308 let call_test_main = ecx.stmt_expr(call_test_main);
309
310 let test_extern_stmt = ecx.stmt_item(
312 sp,
313 ecx.item(sp, ast::AttrVec::new(), ast::ItemKind::ExternCrate(None, test_ident)),
314 );
315
316 let main_attr = ecx.attr_word(sym::rustc_main, sp);
318 let coverage_attr = ecx.attr_nested_word(sym::coverage, sym::off, sp);
320 let doc_hidden_attr = ecx.attr_nested_word(sym::doc, sym::hidden, sp);
322
323 let main_ret_ty = ecx.ty(sp, ast::TyKind::Tup(ThinVec::new()));
325
326 let main_body = if cx.test_runner.is_none() {
328 ecx.block(sp, thin_vec![test_extern_stmt, call_test_main])
329 } else {
330 ecx.block(sp, thin_vec![call_test_main])
331 };
332
333 let decl = ecx.fn_decl(ThinVec::new(), ast::FnRetTy::Ty(main_ret_ty));
334 let sig = ast::FnSig { decl, header: ast::FnHeader::default(), span: sp };
335 let defaultness = ast::Defaultness::Final;
336
337 let main_ident = match cx.reexport_test_harness_main {
339 Some(sym) => Ident::new(sym, sp.with_ctxt(SyntaxContext::root())),
340 None => Ident::new(sym::main, sp),
341 };
342
343 let main = ast::ItemKind::Fn(Box::new(ast::Fn {
344 defaultness,
345 sig,
346 ident: main_ident,
347 generics: ast::Generics::default(),
348 contract: None,
349 body: Some(main_body),
350 define_opaque: None,
351 }));
352
353 let main = P(ast::Item {
354 attrs: thin_vec![main_attr, coverage_attr, doc_hidden_attr],
355 id: ast::DUMMY_NODE_ID,
356 kind: main,
357 vis: ast::Visibility { span: sp, kind: ast::VisibilityKind::Public, tokens: None },
358 span: sp,
359 tokens: None,
360 });
361
362 let main = AstFragment::Items(smallvec![main]);
364 cx.ext_cx.monotonic_expander().fully_expand_fragment(main).make_items().pop().unwrap()
365}
366
367fn mk_tests_slice(cx: &TestCtxt<'_>, sp: Span) -> P<ast::Expr> {
370 debug!("building test vector from {} tests", cx.test_cases.len());
371 let ecx = &cx.ext_cx;
372
373 let mut tests = cx.test_cases.clone();
374 tests.sort_by(|a, b| a.name.as_str().cmp(b.name.as_str()));
375
376 ecx.expr_array_ref(
377 sp,
378 tests
379 .iter()
380 .map(|test| {
381 ecx.expr_addr_of(test.span, ecx.expr_path(ecx.path(test.span, vec![test.ident])))
382 })
383 .collect(),
384 )
385}
386
387fn get_test_name(i: &ast::Item) -> Option<Symbol> {
388 attr::first_attr_value_str_by_name(&i.attrs, sym::rustc_test_marker)
389}
390
391fn get_test_runner(dcx: DiagCtxtHandle<'_>, krate: &ast::Crate) -> Option<ast::Path> {
392 let test_attr = attr::find_by_name(&krate.attrs, sym::test_runner)?;
393 let meta_list = test_attr.meta_item_list()?;
394 let span = test_attr.span;
395 match &*meta_list {
396 [single] => match single.meta_item() {
397 Some(meta_item) if meta_item.is_word() => return Some(meta_item.path.clone()),
398 _ => {
399 dcx.emit_err(errors::TestRunnerInvalid { span });
400 }
401 },
402 _ => {
403 dcx.emit_err(errors::TestRunnerNargs { span });
404 }
405 }
406 None
407}