1use std::cell::Cell;
4use std::str::FromStr;
5use std::sync::Arc;
6
7use proc_macro2::{TokenStream, TokenTree};
8use rustc_attr_parsing::eval_config_entry;
9use rustc_hir::attrs::AttributeKind;
10use rustc_hir::def_id::{CRATE_DEF_ID, LocalDefId};
11use rustc_hir::{self as hir, Attribute, CRATE_HIR_ID, intravisit};
12use rustc_middle::hir::nested_filter;
13use rustc_middle::ty::TyCtxt;
14use rustc_resolve::rustdoc::span_of_fragments;
15use rustc_span::source_map::SourceMap;
16use rustc_span::{BytePos, DUMMY_SP, FileName, Pos, Span};
17
18use super::{DocTestVisitor, ScrapedDocTest};
19use crate::clean::{Attributes, CfgInfo, extract_cfg_from_attrs};
20use crate::html::markdown::{self, CodeLineMapping, ErrorCodes, LangString, MdRelLine};
21
22struct RustCollector {
23 source_map: Arc<SourceMap>,
24 tests: Vec<ScrapedDocTest>,
25 cur_path: Vec<String>,
26 position: Span,
27 global_crate_attrs: Vec<String>,
28}
29
30impl RustCollector {
31 fn get_filename(&self) -> FileName {
32 let filename = self.source_map.span_to_filename(self.position);
33 filename
34 }
35
36 fn get_base_line(&self) -> usize {
37 let sp_lo = self.position.lo().to_usize();
38 let loc = self.source_map.lookup_char_pos(BytePos(sp_lo as u32));
39 loc.line
40 }
41}
42
43impl DocTestVisitor for RustCollector {
44 fn visit_test(
45 &mut self,
46 test: String,
47 config: LangString,
48 rel_line: MdRelLine,
49 code_mappings: Vec<CodeLineMapping>,
50 ) {
51 let base_line = self.get_base_line();
52 let line = base_line + rel_line.offset();
53 let count = Cell::new(base_line);
54 let span = if line > base_line {
55 match self.source_map.span_extend_while(self.position, |c| {
56 if c == '\n' {
57 let count_v = count.get();
58 count.set(count_v + 1);
59 if count_v >= line {
60 return false;
61 }
62 }
63 true
64 }) {
65 Ok(sp) => self.source_map.span_extend_to_line(sp.shrink_to_hi()),
66 _ => self.position,
67 }
68 } else {
69 self.position
70 };
71 self.tests.push(ScrapedDocTest::new(
72 self.get_filename(),
73 line,
74 self.cur_path.clone(),
75 config,
76 test,
77 span,
78 code_mappings,
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 && !eval_config_entry(&self.tcx.sess, cfg.inner()).as_bool()
125 {
126 return;
127 }
128
129 let source_map = self.tcx.sess.source_map();
130 let old_global_crate_attrs_len = self.collector.global_crate_attrs.len();
132 for attr in ast_attrs {
133 let Attribute::Parsed(AttributeKind::Doc(d)) = attr else { continue };
134 for attr_span in &d.test_attrs {
135 if let Ok(snippet) = source_map.span_to_snippet(*attr_span)
137 && let Ok(stream) = TokenStream::from_str(&snippet)
138 {
139 let mut iter = stream.into_iter().peekable();
140 while let Some(token) = iter.next() {
141 if let TokenTree::Ident(i) = token {
142 let i = i.to_string();
143 let peek = iter.peek();
144 match peek {
152 Some(TokenTree::Group(g)) => {
153 let g = g.to_string();
154 iter.next();
155 self.collector.global_crate_attrs.push(format!("{i}{g}"));
157 }
158 Some(TokenTree::Punct(p)) if p.as_char() == '=' => {
161 let p = p.to_string();
162 iter.next();
163 if let Some(last) = iter.next() {
164 self.collector
166 .global_crate_attrs
167 .push(format!("{i}{p}{last}"));
168 }
169 }
170 _ => {
171 self.collector.global_crate_attrs.push(i.to_string());
173 }
174 }
175 }
176 }
177 }
178 }
179 }
180
181 let mut has_name = false;
182 if let Some(name) = name {
183 self.collector.cur_path.push(name);
184 has_name = true;
185 }
186
187 let attrs = Attributes::from_hir(ast_attrs);
190 if let Some(doc) = attrs.opt_doc_value() {
191 let span = span_of_fragments(&attrs.doc_strings).unwrap_or(sp);
192 self.collector.position = if span.edition().at_least_rust_2024() {
193 span
194 } else {
195 ast_attrs
198 .iter()
199 .find(|attr| attr.doc_str().is_some())
200 .map(|attr| {
201 attr.span().ctxt().outer_expn().expansion_cause().unwrap_or(attr.span())
202 })
203 .unwrap_or(DUMMY_SP)
204 };
205 markdown::find_testable_code(
206 &doc,
207 &mut self.collector,
208 self.codes,
209 Some(&crate::html::markdown::ExtraInfo::new(
210 self.tcx,
211 def_id,
212 span,
213 Some(&attrs.doc_strings),
214 )),
215 );
216 }
217
218 nested(self);
219
220 self.collector.global_crate_attrs.truncate(old_global_crate_attrs_len);
222
223 if has_name {
224 self.collector.cur_path.pop();
225 }
226 }
227}
228
229impl<'tcx> intravisit::Visitor<'tcx> for HirCollector<'tcx> {
230 type NestedFilter = nested_filter::All;
231
232 fn maybe_tcx(&mut self) -> Self::MaybeTyCtxt {
233 self.tcx
234 }
235
236 fn visit_item(&mut self, item: &'tcx hir::Item<'_>) {
237 let name = match &item.kind {
238 hir::ItemKind::Impl(impl_) => {
239 Some(rustc_hir_pretty::id_to_string(&self.tcx, impl_.self_ty.hir_id))
240 }
241 _ => item.kind.ident().map(|ident| ident.to_string()),
242 };
243
244 self.visit_testable(name, item.owner_id.def_id, item.span, |this| {
245 intravisit::walk_item(this, item);
246 });
247 }
248
249 fn visit_trait_item(&mut self, item: &'tcx hir::TraitItem<'_>) {
250 self.visit_testable(
251 Some(item.ident.to_string()),
252 item.owner_id.def_id,
253 item.span,
254 |this| {
255 intravisit::walk_trait_item(this, item);
256 },
257 );
258 }
259
260 fn visit_impl_item(&mut self, item: &'tcx hir::ImplItem<'_>) {
261 self.visit_testable(
262 Some(item.ident.to_string()),
263 item.owner_id.def_id,
264 item.span,
265 |this| {
266 intravisit::walk_impl_item(this, item);
267 },
268 );
269 }
270
271 fn visit_foreign_item(&mut self, item: &'tcx hir::ForeignItem<'_>) {
272 self.visit_testable(
273 Some(item.ident.to_string()),
274 item.owner_id.def_id,
275 item.span,
276 |this| {
277 intravisit::walk_foreign_item(this, item);
278 },
279 );
280 }
281
282 fn visit_variant(&mut self, v: &'tcx hir::Variant<'_>) {
283 self.visit_testable(Some(v.ident.to_string()), v.def_id, v.span, |this| {
284 intravisit::walk_variant(this, v);
285 });
286 }
287
288 fn visit_field_def(&mut self, f: &'tcx hir::FieldDef<'_>) {
289 self.visit_testable(Some(f.ident.to_string()), f.def_id, f.span, |this| {
290 intravisit::walk_field_def(this, f);
291 });
292 }
293}