rustc_utils/source_map/
find_bodies.rs

1use log::trace;
2use rustc_hir::{BodyId, intravisit::Visitor};
3use rustc_middle::{hir::nested_filter::OnlyBodies, ty::TyCtxt};
4use rustc_span::Span;
5
6use crate::{SpanExt, block_timer};
7
8struct BodyFinder<'tcx> {
9  tcx: TyCtxt<'tcx>,
10  bodies: Vec<(Span, BodyId)>,
11}
12
13impl<'tcx> Visitor<'tcx> for BodyFinder<'tcx> {
14  type NestedFilter = OnlyBodies;
15
16  fn visit_nested_body(&mut self, id: BodyId) {
17    let tcx = self.tcx;
18
19    // const/static items are considered to have bodies, so we want to exclude
20    // them from our search for functions
21    if !tcx
22      .hir_body_owner_kind(tcx.hir_body_owner_def_id(id))
23      .is_fn_or_closure()
24    {
25      return;
26    }
27
28    let body = tcx.hir_body(id);
29    self.visit_body(body);
30
31    let span = tcx.hir_span_with_body(tcx.hir_body_owner(id));
32    trace!(
33      "Searching body for {:?} with span {span:?} (local {:?})",
34      self
35        .tcx
36        .def_path_debug_str(tcx.hir_body_owner_def_id(id).to_def_id()),
37      span.as_local(body.value.span)
38    );
39
40    if !span.from_expansion() {
41      self.bodies.push((span, id));
42    }
43  }
44}
45
46/// Finds all bodies in the current crate
47pub fn find_bodies(tcx: TyCtxt) -> Vec<(Span, BodyId)> {
48  block_timer!("find_bodies");
49  let mut finder = BodyFinder {
50    tcx,
51    bodies: Vec::new(),
52  };
53  tcx.hir_visit_all_item_likes_in_crate(&mut finder);
54  finder.bodies
55}
56
57/// Finds all the bodies that enclose the given span, from innermost to outermost
58pub fn find_enclosing_bodies(tcx: TyCtxt, sp: Span) -> impl Iterator<Item = BodyId> {
59  let mut bodies = find_bodies(tcx);
60  bodies.retain(|(other, _)| other.contains(sp));
61  bodies.sort_by_key(|(span, _)| span.size());
62  bodies.into_iter().map(|(_, id)| id)
63}
64
65#[cfg(test)]
66mod test {
67  use super::*;
68  use crate::test_utils::{self, CompileResult};
69
70  #[test]
71  fn test_find_bodies() {
72    let input = r"
73// Ignore constants
74const C: usize = 0;
75
76fn a() {
77  // Catch nested bodies
78  fn b() {}
79}
80
81fn c() {}
82
83macro_rules! m {
84  () => { fn d() {} }
85}
86
87// Ignore macro-generated bodies
88m!{}
89";
90    test_utils::CompileBuilder::new(input).compile(|CompileResult { tcx }| {
91      assert_eq!(find_bodies(tcx).len(), 3);
92    });
93  }
94}
rustc_utils/source_map/find_bodies.rs

rustc_utils/source_map/
find_bodies.rs
