equi_ty/extractor/visitors/
struct_enum_visitor.rs

1#![allow(clippy::doc_markdown)]
2
3use std::collections::{BTreeSet, 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::{EnumInfo, FieldInfo, StructInfo, VariantInfo, VariantKind, get_struct_kind};
14use crate::extractor::visitors::{GenericInfo, TypKind};
15use crate::{utils_crate, utils_misc};
16
17pub struct StructEnumVisitor<'tcx> {
18    tcx: TyCtxt<'tcx>,
19    pub structs: HashMap<DefId, StructInfo>,
20    pub enums: HashMap<DefId, EnumInfo>,
21}
22
23impl fmt::Debug for StructEnumVisitor<'_> {
24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25        f.debug_struct("StructVisitor")
26            .field("tcx", &"<tcx>")
27            .field("structs", &self.structs)
28            .field("enums", &self.enums)
29            .finish()
30    }
31}
32
33impl<'tcx> StructEnumVisitor<'tcx> {
34    pub(crate) fn new(tcx: TyCtxt<'tcx>) -> Self {
35        Self {
36            tcx,
37            structs: HashMap::new(),
38            enums: HashMap::new(),
39        }
40    }
41
42    fn field_info_from_field_def(&self, field: &rustc_hir::FieldDef<'_>) -> FieldInfo {
43        let is_unnamed =
44            field.ident.span.is_dummy() || field.ident.as_str().parse::<usize>().is_ok();
45
46        let name = if is_unnamed {
47            None
48        } else {
49            Some(field.ident.to_string())
50        };
51
52        // NOTE: tcx.type_of() wont help us because it gives us a rustc_middle::ty::Ty.
53        // The rustc_hir::Ty contains the Path and the resolution, so use that.
54        // There should be a way to get the final resolved type from rustc_middle::ty::Ty but I
55        // dont have time to explore that right now.
56
57        eprintln!("tmp:: field:\n{field:#?}");
58
59        let (type_kind, type_def_path) = match field.ty.kind {
60            rustc_hir::TyKind::Path(rustc_hir::QPath::Resolved(_, rustc_hir::Path { res, .. })) => {
61                eprintln!("tmp:: res:\n{res:#?}");
62                match res {
63                    rustc_hir::def::Res::Def(rustc_hir::def::DefKind::TyParam, def_id) => {
64                        // generic param
65                        (TypKind::Generic, self.tcx.item_name(*def_id).to_string())
66                    }
67                    rustc_hir::def::Res::Def(_def_kind, def_id) => {
68                        (TypKind::Ty, self.tcx.def_path_debug_str(*def_id))
69                    }
70                    rustc_hir::def::Res::PrimTy(prim_ty) => {
71                        (TypKind::Ty, prim_ty.name().to_string())
72                    }
73                    _ => {
74                        todo!("unhandled res variant:\n{res:#?}");
75                    }
76                }
77            }
78
79            ty_kind => {
80                todo!("unhandled field.ty.kind:\n{ty_kind:#?}");
81            }
82        };
83
84        let type_def_path = utils_misc::remove_cratenum(&type_def_path)
85            .expect("failed while removing cratenum from type");
86
87        eprintln!("tmp:: type_kind: {type_kind:?}, type_def_path: {type_def_path}");
88
89        FieldInfo {
90            name,
91            type_def_path,
92            type_kind,
93        }
94    }
95
96    fn variant_info_from_variant(variant: &rustc_hir::Variant<'_>) -> VariantInfo {
97        let kind = match variant.data {
98            VariantData::Struct { .. } => VariantKind::Struct,
99            VariantData::Tuple(_, _, _) => VariantKind::Tuple,
100            VariantData::Unit(_, _) => VariantKind::Unit,
101        };
102
103        VariantInfo {
104            name: variant.ident.to_string(),
105            kind,
106        }
107    }
108
109    fn handle_generics(generics: &rustc_hir::Generics<'_>) -> Vec<GenericInfo> {
110        generics
111            .params
112            .iter()
113            .filter_map(|p| match p.name {
114                rustc_hir::ParamName::Plain(ident) => Some(GenericInfo {
115                    name: ident.to_string(),
116                    bounds: BTreeSet::new(),
117                }),
118                rustc_hir::ParamName::Error(_) | rustc_hir::ParamName::Fresh => None,
119            })
120            .collect()
121    }
122}
123
124impl<'tcx> HVisitor<'tcx> for StructEnumVisitor<'tcx> {
125    type NestedFilter = All;
126
127    fn maybe_tcx(&mut self) -> Self::MaybeTyCtxt {
128        self.tcx
129    }
130
131    #[allow(clippy::too_many_lines)]
132    fn visit_item(&mut self, item: &'tcx rustc_hir::Item<'tcx>) -> Self::Result {
133        // trace!("item: \n{item:#?}");
134
135        let item_def_id = item.owner_id.to_def_id();
136        trace!("item_def_id: {item_def_id:?}");
137
138        if !item.is_adt() {
139            trace!("skipping as item is not an adt");
140            intravisit::walk_item(self, item);
141            return;
142        }
143
144        // NOTE: since a private item could be reexported by a higher level module, cannot skip all not public items.
145        // let vis = self.tcx.visibility(item_def_id);
146        // if !matches!(vis, ty::Visibility::Public) {
147        //     trace!("skipping as visibility is not pub");
148        //     intravisit::walk_item(self, item);
149        //     return;
150        // }
151
152        let item_def_path_str = self.tcx.def_path_str(item_def_id);
153        trace!("item_def_path_str: {item_def_path_str}");
154        let item_def_path_debug_str = self.tcx.def_path_debug_str(item_def_id);
155        trace!("item_def_path_debug_str: {item_def_path_debug_str}");
156
157        let crate_ = utils_crate::get_crate_name(self.tcx).to_string();
158        // let def_path = format!("{crate_}::{item_def_path_str}");
159        let def_path = utils_misc::remove_cratenum(&item_def_path_debug_str).unwrap();
160        // keeping the default import_path same as def_path till we find a way to identify the
161        // "pub" export path
162        let import_path = def_path.clone();
163
164        match item.kind {
165            rustc_hir::ItemKind::Enum(enum_def, generics) => {
166                if item.span.from_expansion() {
167                    trace!("skipping as item's span is from an expansion");
168                    return;
169                }
170                // trace!("enum item: \n{item:#?}");
171                assert!(!self.enums.contains_key(&item_def_id));
172
173                let variants = enum_def
174                    .variants
175                    .iter()
176                    .map(|v| Self::variant_info_from_variant(v))
177                    .collect();
178
179                let generics = Self::handle_generics(generics);
180
181                let info = EnumInfo {
182                    name: item.ident.name.to_string(),
183                    crate_,
184                    def_path,
185                    import_path,
186                    variants,
187                    methods: BTreeSet::new(),
188                    associated_methods: BTreeSet::new(),
189                    traits: BTreeSet::new(),
190                    generics,
191                };
192                self.enums.insert(item_def_id, info);
193            }
194            rustc_hir::ItemKind::Struct(variant_data, generics) => {
195                if item.span.from_expansion() {
196                    trace!("skipping as item's span is from an expansion");
197                    return;
198                }
199                // trace!("struct item: \n{item:#?}");
200                assert!(!self.structs.contains_key(&item_def_id));
201
202                let fields = match variant_data {
203                    VariantData::Struct { fields, .. } | VariantData::Tuple(fields, ..) => fields
204                        .iter()
205                        .map(|f| self.field_info_from_field_def(f))
206                        .collect(),
207                    VariantData::Unit(hir_id, local_def_id) => Vec::new(),
208                };
209
210                let generics = Self::handle_generics(generics);
211
212                let name = item.ident.name.to_string();
213                let kind = get_struct_kind(item).expect("failed to get struct kind");
214                let info = StructInfo {
215                    name,
216                    kind,
217                    crate_,
218                    def_path,
219                    import_path,
220                    fields,
221                    methods: BTreeSet::new(),
222                    associated_methods: BTreeSet::new(),
223                    traits: BTreeSet::new(),
224                    generics,
225                };
226                self.structs.insert(item_def_id, info);
227            }
228
229            rustc_hir::ItemKind::Impl(_)
230            | rustc_hir::ItemKind::ExternCrate(_)
231            | rustc_hir::ItemKind::Use(_, _)
232            | rustc_hir::ItemKind::Static(_, _, _)
233            | rustc_hir::ItemKind::Const(_, _, _)
234            | rustc_hir::ItemKind::Fn { .. }
235            | rustc_hir::ItemKind::Macro(_, _)
236            | rustc_hir::ItemKind::Mod(_)
237            | rustc_hir::ItemKind::ForeignMod { .. }
238            | rustc_hir::ItemKind::GlobalAsm { .. }
239            | rustc_hir::ItemKind::TyAlias(_, _)
240            | rustc_hir::ItemKind::Union(_, _)
241            | rustc_hir::ItemKind::Trait(_, _, _, _, _)
242            | rustc_hir::ItemKind::TraitAlias(_, _) => {
243                // do nothing for these for now
244                if item.span.from_expansion() {
245                    trace!("skipping as item's span is from an expansion");
246                    return;
247                }
248            }
249        }
250
251        intravisit::walk_item(self, item);
252    }
253}