rapx/analysis/safetyflow_analysis/
fn_collector.rs

1use crate::compat::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            #[cfg(not(rapx_rustc_ge_198))]
44            ItemKind::Trait(
45                _,
46                _is_auto,
47                _safety,
48                _ident,
49                _generics,
50                _generic_bounds,
51                trait_item_ids,
52            ) => {
53                let key = None;
54                let entry = self.fnmap.entry(key).or_default();
55                entry.extend(trait_item_ids.iter().filter_map(|trait_item_id| {
56                    let trait_item = self.tcx.hir_trait_item(*trait_item_id);
57                    if let TraitItemKind::Fn(_, _) = trait_item.kind {
58                        self.tcx
59                            .hir_maybe_body_owned_by(trait_item.owner_id.def_id)
60                            .map(|body| (body.id(), trait_item.span))
61                    } else {
62                        None
63                    }
64                }));
65            }
66            #[cfg(rapx_rustc_ge_198)]
67            ItemKind::Trait {
68                is_auto: _is_auto,
69                safety: _safety,
70                ident: _ident,
71                generics: _generics,
72                bounds: _generic_bounds,
73                items: trait_item_ids,
74                impl_restriction: _,
75                constness: _,
76            } => {
77                let key = None;
78                let entry = self.fnmap.entry(key).or_default();
79                entry.extend(trait_item_ids.iter().filter_map(|trait_item_id| {
80                    let trait_item = self.tcx.hir_trait_item(*trait_item_id);
81                    if let TraitItemKind::Fn(_, _) = trait_item.kind {
82                        self.tcx
83                            .hir_maybe_body_owned_by(trait_item.owner_id.def_id)
84                            .map(|body| (body.id(), trait_item.span))
85                    } else {
86                        None
87                    }
88                }));
89            }
90            ItemKind::Fn {
91                sig: _,
92                generics: _,
93                body: body_id,
94                has_body: _,
95                ident: _,
96            } => {
97                let key = Some(body_id.hir_id);
98                let entry = self.fnmap.entry(key).or_default();
99                entry.push((*body_id, item.span));
100            }
101            _ => (),
102        }
103    }
104
105    fn visit_trait_item(&mut self, _trait_item: &'tcx rustc_hir::TraitItem<'tcx>) {
106        // We don't process items inside trait blocks
107    }
108
109    fn visit_impl_item(&mut self, _impl_item: &'tcx rustc_hir::ImplItem<'tcx>) {
110        // We don't process items inside impl blocks
111    }
112
113    fn visit_foreign_item(&mut self, _foreign_item: &'tcx rustc_hir::ForeignItem<'tcx>) {
114        // We don't process foreign items
115    }
116}