use rustc_data_structures::fx::FxHashMap;
use rustc_hir::{BodyId, HirId, Impl, ItemKind, TraitItemKind, intravisit::Visitor};
use rustc_middle::ty::TyCtxt;
use rustc_span::Span;
pub type FnMap = FxHashMap<Option<HirId>, Vec<(BodyId, Span)>>;
pub struct FnCollector<'tcx> {
tcx: TyCtxt<'tcx>,
fnmap: FnMap,
}
impl<'tcx> FnCollector<'tcx> {
pub fn collect(tcx: TyCtxt<'tcx>) -> FnMap {
let mut collector = FnCollector {
tcx,
fnmap: FnMap::default(),
};
tcx.hir_visit_all_item_likes_in_crate(&mut collector);
collector.fnmap
}
}
impl<'tcx> Visitor<'tcx> for FnCollector<'tcx> {
fn visit_item(&mut self, item: &'tcx rustc_hir::Item<'tcx>) {
match &item.kind {
ItemKind::Impl(Impl {
generics: _generics,
self_ty,
items: impl_items,
..
}) => {
let key = Some(self_ty.hir_id);
let entry = self.fnmap.entry(key).or_default();
entry.extend(impl_items.iter().filter_map(|impl_item_id| {
self.tcx
.hir_maybe_body_owned_by(impl_item_id.owner_id.def_id)
.map(|body| (body.id(), self.tcx.hir_impl_item(*impl_item_id).span))
}));
}
ItemKind::Trait(
_,
_is_auto,
_safety,
_ident,
_generics,
_generic_bounds,
trait_item_ids,
) => {
let key = None;
let entry = self.fnmap.entry(key).or_default();
entry.extend(trait_item_ids.iter().filter_map(|trait_item_id| {
let trait_item = self.tcx.hir_trait_item(*trait_item_id);
if let TraitItemKind::Fn(_, _) = trait_item.kind {
self.tcx
.hir_maybe_body_owned_by(trait_item.owner_id.def_id)
.map(|body| (body.id(), trait_item.span))
} else {
None
}
}));
}
ItemKind::Fn {
sig: _,
generics: _,
body: body_id,
has_body: _,
ident: _,
} => {
let key = Some(body_id.hir_id);
let entry = self.fnmap.entry(key).or_default();
entry.push((*body_id, item.span));
}
_ => (),
}
}
fn visit_trait_item(&mut self, _trait_item: &'tcx rustc_hir::TraitItem<'tcx>) {
}
fn visit_impl_item(&mut self, _impl_item: &'tcx rustc_hir::ImplItem<'tcx>) {
}
fn visit_foreign_item(&mut self, _foreign_item: &'tcx rustc_hir::ForeignItem<'tcx>) {
}
}