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