rapx/analysis/upg/
fn_collector.rs

1use rustc_data_structures::fx::FxHashMap;
2use rustc_hir::{BodyId, HirId, Impl, ItemKind, TraitItemKind, intravisit::Visitor};
3use rustc_middle::ty::TyCtxt;
4use rustc_span::Span;
5
6/// Maps `HirId` of a type to `BodyId` of related impls.
7pub type FnMap = FxHashMap<Option<HirId>, Vec<(BodyId, Span)>>;
8
9///This structs is used to collect all functions and implementations of structs and traits.
10pub struct FnCollector<'tcx> {
11    tcx: TyCtxt<'tcx>,
12    fnmap: FnMap,
13}
14
15impl<'tcx> FnCollector<'tcx> {
16    pub fn collect(tcx: TyCtxt<'tcx>) -> FnMap {
17        let mut collector = FnCollector {
18            tcx,
19            fnmap: FnMap::default(),
20        };
21        tcx.hir_visit_all_item_likes_in_crate(&mut collector);
22        collector.fnmap
23    }
24}
25
26impl<'tcx> Visitor<'tcx> for FnCollector<'tcx> {
27    fn visit_item(&mut self, item: &'tcx rustc_hir::Item<'tcx>) {
28        match &item.kind {
29            ItemKind::Impl(Impl {
30                generics: _generics,
31                self_ty,
32                items: impl_items,
33                ..
34            }) => {
35                let key = Some(self_ty.hir_id);
36                let entry = self.fnmap.entry(key).or_default();
37                entry.extend(impl_items.iter().filter_map(|impl_item_id| {
38                    self.tcx
39                        .hir_maybe_body_owned_by(impl_item_id.owner_id.def_id)
40                        .map(|body| (body.id(), self.tcx.hir_impl_item(*impl_item_id).span))
41                }));
42            }
43            ItemKind::Trait(
44                _,
45                _is_auto,
46                _safety,
47                _ident,
48                _generics,
49                _generic_bounds,
50                trait_item_ids,
51            ) => {
52                let key = None;
53                let entry = self.fnmap.entry(key).or_default();
54                entry.extend(trait_item_ids.iter().filter_map(|trait_item_id| {
55                    let trait_item = self.tcx.hir_trait_item(*trait_item_id);
56                    if let TraitItemKind::Fn(_, _) = trait_item.kind {
57                        self.tcx
58                            .hir_maybe_body_owned_by(trait_item.owner_id.def_id)
59                            .map(|body| (body.id(), trait_item.span))
60                    } else {
61                        None
62                    }
63                }));
64            }
65            ItemKind::Fn {
66                sig: _,
67                generics: _,
68                body: body_id,
69                has_body: _,
70                ident: _,
71            } => {
72                let key = Some(body_id.hir_id);
73                let entry = self.fnmap.entry(key).or_default();
74                entry.push((*body_id, item.span));
75            }
76            _ => (),
77        }
78    }
79
80    fn visit_trait_item(&mut self, _trait_item: &'tcx rustc_hir::TraitItem<'tcx>) {
81        // We don't process items inside trait blocks
82    }
83
84    fn visit_impl_item(&mut self, _impl_item: &'tcx rustc_hir::ImplItem<'tcx>) {
85        // We don't process items inside impl blocks
86    }
87
88    fn visit_foreign_item(&mut self, _foreign_item: &'tcx rustc_hir::ForeignItem<'tcx>) {
89        // We don't process foreign items
90    }
91}
rapx/analysis/upg/fn_collector.rs

rapx/analysis/upg/
fn_collector.rs
