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