equi_ty/extractor/visitors/
impl_visitor.rs

1#![allow(clippy::doc_markdown)]
2
3use std::collections::HashMap;
4use std::fmt;
5
6use rustc_hir::VariantData;
7use rustc_hir::def_id::DefId;
8use rustc_hir::intravisit::{self, Visitor as HVisitor};
9use rustc_middle::hir::nested_filter::All;
10use rustc_middle::ty::TyCtxt;
11use tracing::trace;
12
13use super::{DefPathStr, EnumInfo, FieldInfo, StructInfo, TraitInfo, VariantInfo, VariantKind};
14use crate::{utils_crate, utils_misc};
15
16pub struct ImplVisitor<'tcx> {
17    tcx: TyCtxt<'tcx>,
18    pub structs: HashMap<DefId, StructInfo>,
19    pub enums: HashMap<DefId, EnumInfo>,
20
21    all_traits: Vec<DefId>,
22
23    pub foreign_structs: HashMap<DefId, DefPathStr>,
24    pub foreign_enums: HashMap<DefId, DefPathStr>,
25}
26
27impl fmt::Debug for ImplVisitor<'_> {
28    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
29        f.debug_struct("StructVisitor")
30            .field("tcx", &"<tcx>")
31            .field("structs", &self.structs)
32            .field("enums", &self.enums)
33            .field("foreign_structs", &self.foreign_structs)
34            .field("foreign_enums", &self.foreign_enums)
35            .finish()
36    }
37}
38
39impl<'tcx> ImplVisitor<'tcx> {
40    pub(crate) fn new(
41        tcx: TyCtxt<'tcx>,
42        structs: HashMap<DefId, StructInfo>,
43        enums: HashMap<DefId, EnumInfo>,
44    ) -> Self {
45        Self {
46            tcx,
47            structs,
48            enums,
49            foreign_structs: HashMap::new(),
50            foreign_enums: HashMap::new(),
51            all_traits: tcx.all_traits().collect(),
52        }
53    }
54
55    fn variant_info_from_variant(variant: &rustc_hir::Variant<'_>) -> VariantInfo {
56        let kind = match variant.data {
57            VariantData::Struct { .. } => VariantKind::Struct,
58            VariantData::Tuple(_, _, _) => VariantKind::Tuple,
59            VariantData::Unit(_, _) => VariantKind::Unit,
60        };
61
62        VariantInfo {
63            name: variant.ident.to_string(),
64            kind,
65        }
66    }
67
68    fn handle_impl(&mut self, impl_: &rustc_hir::Impl<'_>) {
69        let rustc_hir::Impl {
70            constness,
71            safety,
72            polarity,
73            defaultness,
74            defaultness_span,
75            generics,
76            of_trait,
77            self_ty,
78            items,
79        } = impl_;
80
81        // parent item typ? (enum/struct)
82        if let rustc_hir::TyKind::Path(rustc_hir::QPath::Resolved(_, path)) = self_ty.kind {
83            if let Some(selfty_def_id) = path.res.opt_def_id() {
84                trace!("tmp:: selfty_def_id: {selfty_def_id:?}");
85                let def_kind = self.tcx.def_kind(selfty_def_id);
86                let is_selfty_struct = matches!(def_kind, rustc_hir::def::DefKind::Struct);
87                let is_selfty_enum = matches!(def_kind, rustc_hir::def::DefKind::Enum);
88
89                if !is_selfty_struct && !is_selfty_enum {
90                    trace!("skipping impl as it is not for a struct/enum");
91                    return;
92                }
93
94                #[allow(clippy::single_match_else)]
95                match of_trait {
96                    Some(trait_) => {
97                        eprintln!("tmp:: trait_: {trait_:#?}");
98                        // check if this trait is a local trait
99                        let trait_def_id = trait_.path.res.def_id();
100                        eprintln!("tmp:: trait_def_id: {trait_def_id:?}");
101                        let mut trait_def_path = self.tcx.def_path_str(trait_def_id);
102                        let trait_def_path_debug_str = self.tcx.def_path_debug_str(trait_def_id);
103                        eprintln!("tmp:: trait_def_path_debug_str: {trait_def_path_debug_str}");
104                        if trait_def_id.is_local() {
105                            let crate_ = utils_crate::get_crate_name(self.tcx).to_string();
106                            trait_def_path.insert_str(0, &crate_);
107                        }
108                        eprintln!("tmp:: trait_def_path: {trait_def_path}");
109
110                        if !selfty_def_id.is_local() {
111                            todo!("handle case when a trait is implemented for a foreign item");
112                        }
113
114                        assert!(
115                            self.enums.contains_key(&selfty_def_id)
116                                || self.structs.contains_key(&selfty_def_id),
117                            "selfty_def_id not found in enums/structs ({:?})",
118                            selfty_def_id
119                        );
120
121                        let def_path =
122                            utils_misc::remove_cratenum(&trait_def_path_debug_str).unwrap();
123
124                        let ti = TraitInfo {
125                            name: self.tcx.item_name(trait_def_id).to_string(),
126                            def_path: def_path.clone(),
127                            import_path: def_path,
128                        };
129
130                        if let Some(ei) = self.enums.get_mut(&selfty_def_id) {
131                            ei.traits.insert(ti);
132                        } else if let Some(si) = self.structs.get_mut(&selfty_def_id) {
133                            si.traits.insert(ti);
134                        } else {
135                            todo!(
136                                "selfty is neither local struct nor local enum [PLEASE INVESTIGATE"
137                            );
138                        }
139                    }
140                    None => {
141                        // regular impl
142                        for impl_item in *items {
143                            // we only care about fns for now.
144                            if let rustc_hir::AssocItemKind::Fn { has_self } = impl_item.kind {
145                                let name = impl_item.ident.to_string();
146                                if is_selfty_struct {
147                                    assert!(self.structs.contains_key(&selfty_def_id));
148                                    self.structs.entry(selfty_def_id).and_modify(|si| {
149                                        if has_self {
150                                            si.associated_methods.insert(name);
151                                        } else {
152                                            si.methods.insert(name);
153                                        }
154                                    });
155                                } else if is_selfty_enum {
156                                    assert!(self.enums.contains_key(&selfty_def_id));
157                                    self.enums.entry(selfty_def_id).and_modify(|ei| {
158                                        if has_self {
159                                            ei.associated_methods.insert(name);
160                                        } else {
161                                            ei.methods.insert(name);
162                                        }
163                                    });
164                                }
165                            }
166                        }
167                    }
168                }
169            }
170        }
171    }
172}
173
174impl<'tcx> HVisitor<'tcx> for ImplVisitor<'tcx> {
175    type NestedFilter = All;
176
177    fn maybe_tcx(&mut self) -> Self::MaybeTyCtxt {
178        self.tcx
179    }
180
181    #[allow(clippy::too_many_lines)]
182    fn visit_item(&mut self, item: &'tcx rustc_hir::Item<'tcx>) -> Self::Result {
183        // trace!("item: \n{item:#?}");
184
185        let item_def_id = item.owner_id.to_def_id();
186        trace!("item_def_id: {item_def_id:?}");
187
188        // // NOTE: since impls are items too, we cannot skip *all* non-adts.
189        // if !item.is_adt() {
190        //     trace!("skipping as item is not an adt");
191        //     intravisit::walk_item(self, item);
192        //     return;
193        // }
194
195        // // NOTE: since an "impl" does not have a visibility, cannot skip all not public items.
196        // let vis = self.tcx.visibility(item_def_id);
197        // if !matches!(vis, ty::Visibility::Public) {
198        //     trace!("skipping as visibility is not pub");
199        //     intravisit::walk_item(self, item);
200        //     return;
201        // }
202
203        let item_def_path_str = self.tcx.def_path_str(item_def_id);
204        trace!("item_def_path_str: {item_def_path_str}");
205        let item_def_path_debug_str = self.tcx.def_path_debug_str(item_def_id);
206        trace!("item_def_path_debug_str: {item_def_path_debug_str}");
207
208        let crate_ = utils_crate::get_crate_name(self.tcx).to_string();
209        // let def_path = format!("{crate_}::{item_def_path_str}");
210        let def_path = utils_misc::remove_cratenum(&item_def_path_debug_str).unwrap();
211        // keeping the default import_path same as def_path till we find a way to identify the
212        // "pub" export path
213        let import_path = def_path.clone();
214
215        match item.kind {
216            rustc_hir::ItemKind::Impl(impl_) => {
217                eprintln!("tmp:: impl: {item:#?}");
218                self.handle_impl(impl_);
219            }
220
221            rustc_hir::ItemKind::Enum(_, _)
222            | rustc_hir::ItemKind::Struct(_, _)
223            | rustc_hir::ItemKind::ExternCrate(_)
224            | rustc_hir::ItemKind::Use(_, _)
225            | rustc_hir::ItemKind::Static(_, _, _)
226            | rustc_hir::ItemKind::Const(_, _, _)
227            | rustc_hir::ItemKind::Fn { .. }
228            | rustc_hir::ItemKind::Macro(_, _)
229            | rustc_hir::ItemKind::Mod(_)
230            | rustc_hir::ItemKind::ForeignMod { .. }
231            | rustc_hir::ItemKind::GlobalAsm { .. }
232            | rustc_hir::ItemKind::TyAlias(_, _)
233            | rustc_hir::ItemKind::Union(_, _)
234            | rustc_hir::ItemKind::Trait(_, _, _, _, _)
235            | rustc_hir::ItemKind::TraitAlias(_, _) => {
236                // do nothing for these for now
237                if item.span.from_expansion() {
238                    trace!("skipping as item's span is from an expansion");
239                    return;
240                }
241            }
242        }
243
244        intravisit::walk_item(self, item);
245    }
246}
equi_ty/extractor/visitors/impl_visitor.rs

equi_ty/extractor/visitors/
impl_visitor.rs
