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