1use std::cell::Cell;
4use std::env;
5use std::sync::Arc;
6
7use rustc_ast_pretty::pprust;
8use rustc_data_structures::fx::FxHashSet;
9use rustc_hir::def_id::{CRATE_DEF_ID, LocalDefId};
10use rustc_hir::{self as hir, CRATE_HIR_ID, intravisit};
11use rustc_middle::hir::nested_filter;
12use rustc_middle::ty::TyCtxt;
13use rustc_resolve::rustdoc::span_of_fragments;
14use rustc_span::source_map::SourceMap;
15use rustc_span::{BytePos, DUMMY_SP, FileName, Pos, Span, sym};
16
17use super::{DocTestVisitor, ScrapedDocTest};
18use crate::clean::{Attributes, extract_cfg_from_attrs};
19use crate::html::markdown::{self, ErrorCodes, LangString, MdRelLine};
20
21struct RustCollector {
22 source_map: Arc<SourceMap>,
23 tests: Vec<ScrapedDocTest>,
24 cur_path: Vec<String>,
25 position: Span,
26 global_crate_attrs: Vec<String>,
27}
28
29impl RustCollector {
30 fn get_filename(&self) -> FileName {
31 let filename = self.source_map.span_to_filename(self.position);
32 if let FileName::Real(ref filename) = filename {
33 let path = filename.remapped_path_if_available();
34 let path = env::current_dir()
37 .map(|cur_dir| path.strip_prefix(&cur_dir).unwrap_or(path))
38 .unwrap_or(path);
39 return path.to_owned().into();
40 }
41 filename
42 }
43
44 fn get_base_line(&self) -> usize {
45 let sp_lo = self.position.lo().to_usize();
46 let loc = self.source_map.lookup_char_pos(BytePos(sp_lo as u32));
47 loc.line
48 }
49}
50
51impl DocTestVisitor for RustCollector {
52 fn visit_test(&mut self, test: String, config: LangString, rel_line: MdRelLine) {
53 let base_line = self.get_base_line();
54 let line = base_line + rel_line.offset();
55 let count = Cell::new(base_line);
56 let span = if line > base_line {
57 match self.source_map.span_extend_while(self.position, |c| {
58 if c == '\n' {
59 let count_v = count.get();
60 count.set(count_v + 1);
61 if count_v >= line {
62 return false;
63 }
64 }
65 true
66 }) {
67 Ok(sp) => self.source_map.span_extend_to_line(sp.shrink_to_hi()),
68 _ => self.position,
69 }
70 } else {
71 self.position
72 };
73 self.tests.push(ScrapedDocTest::new(
74 self.get_filename(),
75 line,
76 self.cur_path.clone(),
77 config,
78 test,
79 span,
80 self.global_crate_attrs.clone(),
81 ));
82 }
83
84 fn visit_header(&mut self, _name: &str, _level: u32) {}
85}
86
87pub(super) struct HirCollector<'tcx> {
88 codes: ErrorCodes,
89 tcx: TyCtxt<'tcx>,
90 collector: RustCollector,
91}
92
93impl<'tcx> HirCollector<'tcx> {
94 pub fn new(codes: ErrorCodes, tcx: TyCtxt<'tcx>) -> Self {
95 let collector = RustCollector {
96 source_map: tcx.sess.psess.clone_source_map(),
97 cur_path: vec![],
98 position: DUMMY_SP,
99 tests: vec![],
100 global_crate_attrs: Vec::new(),
101 };
102 Self { codes, tcx, collector }
103 }
104
105 pub fn collect_crate(mut self) -> Vec<ScrapedDocTest> {
106 let tcx = self.tcx;
107 self.visit_testable(None, CRATE_DEF_ID, tcx.hir_span(CRATE_HIR_ID), |this| {
108 tcx.hir_walk_toplevel_module(this)
109 });
110 self.collector.tests
111 }
112}
113
114impl HirCollector<'_> {
115 fn visit_testable<F: FnOnce(&mut Self)>(
116 &mut self,
117 name: Option<String>,
118 def_id: LocalDefId,
119 sp: Span,
120 nested: F,
121 ) {
122 let ast_attrs = self.tcx.hir_attrs(self.tcx.local_def_id_to_hir_id(def_id));
123 if let Some(ref cfg) =
124 extract_cfg_from_attrs(ast_attrs.iter(), self.tcx, &FxHashSet::default())
125 && !cfg.matches(&self.tcx.sess.psess)
126 {
127 return;
128 }
129
130 let old_global_crate_attrs_len = self.collector.global_crate_attrs.len();
132 for doc_test_attrs in ast_attrs
133 .iter()
134 .filter(|a| a.has_name(sym::doc))
135 .flat_map(|a| a.meta_item_list().unwrap_or_default())
136 .filter(|a| a.has_name(sym::test))
137 {
138 let Some(doc_test_attrs) = doc_test_attrs.meta_item_list() else { continue };
139 for attr in doc_test_attrs
140 .iter()
141 .filter(|a| a.has_name(sym::attr))
142 .flat_map(|a| a.meta_item_list().unwrap_or_default())
143 .map(|i| pprust::meta_list_item_to_string(i))
144 {
145 self.collector.global_crate_attrs.push(attr);
147 }
148 }
149
150 let mut has_name = false;
151 if let Some(name) = name {
152 self.collector.cur_path.push(name);
153 has_name = true;
154 }
155
156 let attrs = Attributes::from_hir(ast_attrs);
159 if let Some(doc) = attrs.opt_doc_value() {
160 let span = span_of_fragments(&attrs.doc_strings).unwrap_or(sp);
161 self.collector.position = if span.edition().at_least_rust_2024() {
162 span
163 } else {
164 ast_attrs
167 .iter()
168 .find(|attr| attr.doc_str().is_some())
169 .map(|attr| {
170 attr.span().ctxt().outer_expn().expansion_cause().unwrap_or(attr.span())
171 })
172 .unwrap_or(DUMMY_SP)
173 };
174 markdown::find_testable_code(
175 &doc,
176 &mut self.collector,
177 self.codes,
178 Some(&crate::html::markdown::ExtraInfo::new(self.tcx, def_id, span)),
179 );
180 }
181
182 nested(self);
183
184 self.collector.global_crate_attrs.truncate(old_global_crate_attrs_len);
186
187 if has_name {
188 self.collector.cur_path.pop();
189 }
190 }
191}
192
193impl<'tcx> intravisit::Visitor<'tcx> for HirCollector<'tcx> {
194 type NestedFilter = nested_filter::All;
195
196 fn maybe_tcx(&mut self) -> Self::MaybeTyCtxt {
197 self.tcx
198 }
199
200 fn visit_item(&mut self, item: &'tcx hir::Item<'_>) {
201 let name = match &item.kind {
202 hir::ItemKind::Impl(impl_) => {
203 Some(rustc_hir_pretty::id_to_string(&self.tcx, impl_.self_ty.hir_id))
204 }
205 _ => item.kind.ident().map(|ident| ident.to_string()),
206 };
207
208 self.visit_testable(name, item.owner_id.def_id, item.span, |this| {
209 intravisit::walk_item(this, item);
210 });
211 }
212
213 fn visit_trait_item(&mut self, item: &'tcx hir::TraitItem<'_>) {
214 self.visit_testable(
215 Some(item.ident.to_string()),
216 item.owner_id.def_id,
217 item.span,
218 |this| {
219 intravisit::walk_trait_item(this, item);
220 },
221 );
222 }
223
224 fn visit_impl_item(&mut self, item: &'tcx hir::ImplItem<'_>) {
225 self.visit_testable(
226 Some(item.ident.to_string()),
227 item.owner_id.def_id,
228 item.span,
229 |this| {
230 intravisit::walk_impl_item(this, item);
231 },
232 );
233 }
234
235 fn visit_foreign_item(&mut self, item: &'tcx hir::ForeignItem<'_>) {
236 self.visit_testable(
237 Some(item.ident.to_string()),
238 item.owner_id.def_id,
239 item.span,
240 |this| {
241 intravisit::walk_foreign_item(this, item);
242 },
243 );
244 }
245
246 fn visit_variant(&mut self, v: &'tcx hir::Variant<'_>) {
247 self.visit_testable(Some(v.ident.to_string()), v.def_id, v.span, |this| {
248 intravisit::walk_variant(this, v);
249 });
250 }
251
252 fn visit_field_def(&mut self, f: &'tcx hir::FieldDef<'_>) {
253 self.visit_testable(Some(f.ident.to_string()), f.def_id, f.span, |this| {
254 intravisit::walk_field_def(this, f);
255 });
256 }
257}