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