Skip to main content

rustdoc/doctest/
rust.rs

1//! Doctest functionality used only for doctests in `.rs` source files.
2
3use 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        // Try collecting `#[doc(test(attr(...)))]`
131        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                // FIXME: This is ugly, remove when `test_attrs` has been ported to new attribute API.
136                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                            // From this ident, we can have things like:
145                            //
146                            // * Group: `allow(...)`
147                            // * Name/value: `crate_name = "..."`
148                            // * Tokens: `html_no_url`
149                            //
150                            // So we peek next element to know what case we are in.
151                            match peek {
152                                Some(TokenTree::Group(g)) => {
153                                    let g = g.to_string();
154                                    iter.next();
155                                    // Add the additional attributes to the global_crate_attrs vector
156                                    self.collector.global_crate_attrs.push(format!("{i}{g}"));
157                                }
158                                // If next item is `=`, it means it's a name value so we will need
159                                // to get the value as well.
160                                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                                        // Add the additional attributes to the global_crate_attrs vector
165                                        self.collector
166                                            .global_crate_attrs
167                                            .push(format!("{i}{p}{last}"));
168                                    }
169                                }
170                                _ => {
171                                    // Add the additional attributes to the global_crate_attrs vector
172                                    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        // The collapse-docs pass won't combine sugared/raw doc attributes, or included files with
188        // anything else, this will combine them for us.
189        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                // this span affects filesystem path resolution,
196                // so we need to keep it the same as it was previously
197                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        // Restore global_crate_attrs to it's previous size/content
221        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}