Skip to main content

rapx/helpers/
fn_info.rs

1use rustc_hir::{Safety, def::DefKind, def_id::DefId};
2use rustc_middle::{
3    mir::Local,
4    ty,
5    ty::{AssocKind, Mutability, TyCtxt, TyKind},
6};
7use rustc_span::{kw, sym};
8use std::{
9    collections::{HashMap, HashSet},
10    fmt::Debug,
11    hash::Hash,
12};
13use syn::Expr;
14
15pub use super::mir_scan::check_safety;
16pub use super::name::get_cleaned_def_path_name;
17
18#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
19pub enum FnKind {
20    Fn,
21    Method,
22    Constructor,
23    Intrinsic,
24}
25
26#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
27pub struct FnInfo {
28    pub def_id: DefId,
29    pub fn_safety: Safety,
30    pub fn_kind: FnKind,
31}
32
33impl FnInfo {
34    pub fn new(def_id: DefId, fn_safety: Safety, fn_kind: FnKind) -> Self {
35        FnInfo {
36            def_id,
37            fn_safety,
38            fn_kind,
39        }
40    }
41}
42
43#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
44pub struct AdtInfo {
45    pub def_id: DefId,
46    pub literal_cons_enabled: bool,
47}
48
49impl AdtInfo {
50    pub fn new(def_id: DefId, literal_cons_enabled: bool) -> Self {
51        AdtInfo {
52            def_id,
53            literal_cons_enabled,
54        }
55    }
56}
57
58pub fn check_visibility(tcx: TyCtxt, func_defid: DefId) -> bool {
59    if !tcx.visibility(func_defid).is_public() {
60        return false;
61    }
62    true
63}
64
65pub fn get_type(tcx: TyCtxt<'_>, def_id: DefId) -> FnKind {
66    if let Some(assoc_item) = tcx.opt_associated_item(def_id) {
67        match assoc_item.kind {
68            AssocKind::Fn { has_self, .. } => {
69                if has_self {
70                    return FnKind::Method;
71                } else {
72                    let fn_sig = tcx.fn_sig(def_id).skip_binder();
73                    let output = fn_sig.output().skip_binder();
74                    // return type is 'Self'
75                    if output.is_param(0) {
76                        return FnKind::Constructor;
77                    }
78                    // return type is struct's name
79                    if let Some(impl_id) = assoc_item.impl_container(tcx) {
80                        let ty = tcx.type_of(impl_id).skip_binder();
81                        if output == ty {
82                            return FnKind::Constructor;
83                        }
84                    }
85                    match output.kind() {
86                        TyKind::Ref(_, ref_ty, _) => {
87                            if ref_ty.is_param(0) {
88                                return FnKind::Constructor;
89                            }
90                            if let Some(impl_id) = assoc_item.impl_container(tcx) {
91                                let ty = tcx.type_of(impl_id).skip_binder();
92                                if *ref_ty == ty {
93                                    return FnKind::Constructor;
94                                }
95                            }
96                        }
97                        TyKind::Adt(adt_def, substs) => {
98                            if adt_def.is_enum()
99                                && (tcx.is_diagnostic_item(sym::Option, adt_def.did())
100                                    || tcx.is_diagnostic_item(sym::Result, adt_def.did())
101                                    || tcx.is_diagnostic_item(kw::Box, adt_def.did()))
102                            {
103                                let inner_ty = substs.type_at(0);
104                                if inner_ty.is_param(0) {
105                                    return FnKind::Constructor;
106                                }
107                                if let Some(impl_id) = assoc_item.impl_container(tcx) {
108                                    let ty_impl = tcx.type_of(impl_id).skip_binder();
109                                    if inner_ty == ty_impl {
110                                        return FnKind::Constructor;
111                                    }
112                                }
113                            }
114                        }
115                        _ => {}
116                    }
117                }
118            }
119            _ => todo!(),
120        }
121    }
122    return FnKind::Fn;
123}
124
125/// Returns true when the function is a "wrapped" constructor that returns
126/// `Option<Self>` / `Result<Self, _>` rather than a bare `Self`.
127///
128/// `get_type` classifies these as [`FnKind::Constructor`], but for the wrapped
129/// forms the `None`/`Err` paths do not produce a `Self`, so a struct invariant
130/// can only be meaningfully discharged on the `Some`/`Ok` paths. This helper
131/// lets `verify_struct_invariants` skip the benign `Unknown` results on the
132/// non-`Self` paths. (`Box<Self>` is intentionally *not* included: every path
133/// still produces a `Self` behind the pointer.)
134pub fn returns_wrapped_self(tcx: TyCtxt<'_>, def_id: DefId) -> bool {
135    let Some(assoc_item) = tcx.opt_associated_item(def_id) else {
136        return false;
137    };
138    if !matches!(assoc_item.kind, AssocKind::Fn { has_self: false, .. }) {
139        return false;
140    }
141    let fn_sig = tcx.fn_sig(def_id).skip_binder();
142    let output = fn_sig.output().skip_binder();
143    let TyKind::Adt(adt_def, substs) = output.kind() else {
144        return false;
145    };
146    if !(adt_def.is_enum()
147        && (tcx.is_diagnostic_item(sym::Option, adt_def.did())
148            || tcx.is_diagnostic_item(sym::Result, adt_def.did())))
149    {
150        return false;
151    }
152    let inner_ty = substs.type_at(0);
153    if inner_ty.is_param(0) {
154        return true;
155    }
156    if let Some(impl_id) = assoc_item.impl_container(tcx) {
157        let ty_impl = tcx.type_of(impl_id).skip_binder();
158        if inner_ty == ty_impl {
159            return true;
160        }
161    }
162    false
163}
164
165// result: adt_def_id, is_literal
166pub fn get_adt_via_method(tcx: TyCtxt<'_>, method_def_id: DefId) -> Option<AdtInfo> {
167    let assoc_item = tcx.opt_associated_item(method_def_id)?;
168    let impl_id = assoc_item.impl_container(tcx)?;
169    let ty = tcx.type_of(impl_id).skip_binder();
170    let adt_def = ty.ty_adt_def()?;
171    let adt_def_id = adt_def.did();
172
173    let all_fields: Vec<_> = adt_def.all_fields().collect();
174    let total_count = all_fields.len();
175
176    if total_count == 0 {
177        return Some(AdtInfo::new(adt_def_id, true));
178    }
179
180    let pub_count = all_fields
181        .iter()
182        .filter(|field| tcx.visibility(field.did).is_public())
183        .count();
184
185    if pub_count == 0 {
186        return None;
187    }
188    Some(AdtInfo::new(adt_def_id, pub_count == total_count))
189}
190// return all the impls def id of corresponding struct
191fn get_impls_for_struct(tcx: TyCtxt<'_>, struct_def_id: DefId) -> Vec<DefId> {
192    let mut impls = Vec::new();
193    for item_id in tcx.hir_crate_items(()).free_items() {
194        let item = tcx.hir_item(item_id);
195        if let rustc_hir::ItemKind::Impl(impl_details) = &item.kind {
196            if let rustc_hir::TyKind::Path(rustc_hir::QPath::Resolved(_, path)) =
197                &impl_details.self_ty.kind
198            {
199                if let rustc_hir::def::Res::Def(_, def_id) = path.res {
200                    if def_id == struct_def_id {
201                        impls.push(item_id.owner_id.to_def_id());
202                    }
203                }
204            }
205        }
206    }
207    impls
208}
209
210pub fn get_adt_def_id_by_adt_method(tcx: TyCtxt<'_>, def_id: DefId) -> Option<DefId> {
211    if let Some(assoc_item) = tcx.opt_associated_item(def_id) {
212        if let Some(impl_id) = assoc_item.impl_container(tcx) {
213            // get struct ty
214            let ty = tcx.type_of(impl_id).skip_binder();
215            if let Some(adt_def) = ty.ty_adt_def() {
216                return Some(adt_def.did());
217            }
218        }
219    }
220    None
221}
222
223fn has_mut_self_param(tcx: TyCtxt, def_id: DefId) -> bool {
224    if let Some(assoc_item) = tcx.opt_associated_item(def_id) {
225        match assoc_item.kind {
226            AssocKind::Fn { has_self, .. } => {
227                if has_self && tcx.is_mir_available(def_id) {
228                    let body = tcx.optimized_mir(def_id);
229                    let fst_arg = body.local_decls[Local::from_usize(1)].clone();
230                    let ty = fst_arg.ty;
231                    let is_mut_ref = matches!(ty.kind(), ty::Ref(_, _, Mutability::Mut));
232                    return fst_arg.mutability.is_mut() || is_mut_ref;
233                }
234            }
235            _ => (),
236        }
237    }
238    false
239}
240
241// Check each field's visibility, return the public fields vec
242fn get_public_fields(tcx: TyCtxt, def_id: DefId) -> HashSet<usize> {
243    let adt_def = tcx.adt_def(def_id);
244    adt_def
245        .all_fields()
246        .enumerate()
247        .filter_map(|(index, field_def)| tcx.visibility(field_def.did).is_public().then_some(index))
248        .collect()
249}
250
251/// parse expr into number.
252pub fn parse_expr_into_number(expr: &Expr) -> Option<usize> {
253    if let Expr::Lit(expr_lit) = expr {
254        if let syn::Lit::Int(lit_int) = &expr_lit.lit {
255            return lit_int.base10_parse::<usize>().ok();
256        }
257    }
258    None
259}
260
261pub fn get_all_std_fns_by_rustc_public(tcx: TyCtxt) -> Vec<DefId> {
262    let mut all_std_fn_def = Vec::new();
263    let mut results = Vec::new();
264    let mut core_fn_def: Vec<_> = rustc_public::find_crates("core")
265        .iter()
266        .flat_map(|krate| krate.fn_defs())
267        .collect();
268    let mut std_fn_def: Vec<_> = rustc_public::find_crates("std")
269        .iter()
270        .flat_map(|krate| krate.fn_defs())
271        .collect();
272    let mut alloc_fn_def: Vec<_> = rustc_public::find_crates("alloc")
273        .iter()
274        .flat_map(|krate| krate.fn_defs())
275        .collect();
276    all_std_fn_def.append(&mut core_fn_def);
277    all_std_fn_def.append(&mut std_fn_def);
278    all_std_fn_def.append(&mut alloc_fn_def);
279
280    for fn_def in &all_std_fn_def {
281        let def_id = crate::def_id::to_internal(fn_def, tcx);
282        results.push(def_id);
283    }
284    results
285}
286
287// Input the adt def id
288// Return set of (mutable method def_id, fields can be modified)
289pub fn get_all_mutable_methods(tcx: TyCtxt, src_def_id: DefId) -> HashMap<DefId, HashSet<usize>> {
290    let mut std_results = HashMap::new();
291    if get_type(tcx, src_def_id) == FnKind::Constructor {
292        return std_results;
293    }
294    let all_std_fn_def = get_all_std_fns_by_rustc_public(tcx);
295    let target_adt_def = get_adt_def_id_by_adt_method(tcx, src_def_id);
296    let mut is_std = false;
297    for &def_id in &all_std_fn_def {
298        let adt_def = get_adt_def_id_by_adt_method(tcx, def_id);
299        if adt_def.is_some() && adt_def == target_adt_def && src_def_id != def_id {
300            if has_mut_self_param(tcx, def_id) {
301                std_results.insert(def_id, HashSet::new());
302            }
303            is_std = true;
304        }
305    }
306    if is_std {
307        return std_results;
308    }
309    let mut results = HashMap::new();
310    let public_fields = target_adt_def.map_or_else(HashSet::new, |def| get_public_fields(tcx, def));
311    let impl_vec = target_adt_def.map_or_else(Vec::new, |def| get_impls_for_struct(tcx, def));
312    for impl_id in impl_vec {
313        if !matches!(tcx.def_kind(impl_id), rustc_hir::def::DefKind::Impl { .. }) {
314            continue;
315        }
316        let associated_items = tcx.associated_items(impl_id);
317        for item in associated_items.in_definition_order() {
318            if let ty::AssocKind::Fn {
319                name: _,
320                has_self: _,
321            } = item.kind
322            {
323                let item_def_id = item.def_id;
324                if has_mut_self_param(tcx, item_def_id) {
325                    let modified_fields = public_fields.clone();
326                    results.insert(item_def_id, modified_fields);
327                }
328            }
329        }
330    }
331    results
332}
333
334pub fn get_cons(tcx: TyCtxt<'_>, def_id: DefId) -> Vec<DefId> {
335    let mut cons = Vec::new();
336    if tcx.def_kind(def_id) == DefKind::Fn || get_type(tcx, def_id) == FnKind::Constructor {
337        return cons;
338    }
339    if let Some(assoc_item) = tcx.opt_associated_item(def_id) {
340        if let Some(impl_id) = assoc_item.impl_container(tcx) {
341            let ty = tcx.type_of(impl_id).skip_binder();
342            if let Some(adt_def) = ty.ty_adt_def() {
343                let adt_def_id = adt_def.did();
344                let impls = tcx.inherent_impls(adt_def_id);
345                for impl_def_id in impls {
346                    for item in tcx.associated_item_def_ids(*impl_def_id) {
347                        if (tcx.def_kind(*item) == DefKind::Fn
348                            || tcx.def_kind(*item) == DefKind::AssocFn)
349                            && get_type(tcx, *item) == FnKind::Constructor
350                        {
351                            cons.push(*item);
352                        }
353                    }
354                }
355            }
356        }
357    }
358    cons
359}
360
361/// Find `&mut self` methods (mutators) on the same struct as `def_id`.
362///
363/// A mutator is a method whose first parameter is a mutable reference to Self.
364/// These methods can change struct fields and affect subsequent invariant checks.
365pub fn get_muts(tcx: TyCtxt<'_>, def_id: DefId) -> Vec<DefId> {
366    let mut muts = Vec::new();
367    if let Some(assoc_item) = tcx.opt_associated_item(def_id) {
368        if let Some(impl_id) = assoc_item.impl_container(tcx) {
369            let ty = tcx.type_of(impl_id).skip_binder();
370            if let Some(adt_def) = ty.ty_adt_def() {
371                let adt_def_id = adt_def.did();
372                let impls = tcx.inherent_impls(adt_def_id);
373                for impl_def_id in impls {
374                    for item in tcx.associated_item_def_ids(*impl_def_id) {
375                        if !matches!(tcx.def_kind(*item), DefKind::Fn | DefKind::AssocFn) {
376                            continue;
377                        }
378                        if get_type(tcx, *item) != FnKind::Method {
379                            continue;
380                        }
381                        let Some(assoc) = tcx.opt_associated_item(*item) else {
382                            continue;
383                        };
384                        if !matches!(assoc.kind, AssocKind::Fn { has_self: true, .. }) {
385                            continue;
386                        }
387                        let fn_sig = tcx.fn_sig(*item).instantiate_identity().skip_binder();
388                        let all = fn_sig.inputs_and_output;
389                        let first_param = all.first().copied();
390                        if let Some(TyKind::Ref(_, _, Mutability::Mut)) =
391                            first_param.map(|ty| ty.kind())
392                        {
393                            muts.push(*item);
394                        }
395                    }
396                }
397            }
398        }
399    }
400    muts
401}
402
403pub fn append_fn_with_types(tcx: TyCtxt, def_id: DefId) -> FnInfo {
404    FnInfo::new(def_id, check_safety(tcx, def_id), get_type(tcx, def_id))
405}
406
407pub fn get_ptr_deref_dummy_def_id(tcx: TyCtxt<'_>) -> Option<DefId> {
408    tcx.hir_crate_items(()).free_items().find_map(|item_id| {
409        let def_id = item_id.owner_id.to_def_id();
410        let name = tcx.opt_item_name(def_id)?;
411
412        (name.as_str() == "__raw_ptr_deref_dummy").then_some(def_id)
413    })
414}
415
416/// Return field indices that a `&mut self` method writes to.
417///
418/// Scans the MIR body for assignments to `(*self).field_n` and returns the
419/// set of field indices that are modified.  Used by --skip-invariant mode to know which
420/// constructor-inherited invariants are invalidated by a mutator.
421pub fn get_mutated_fields(tcx: TyCtxt<'_>, def_id: DefId) -> Vec<usize> {
422    use rustc_middle::mir::{ProjectionElem, StatementKind};
423
424    let body = tcx.optimized_mir(def_id);
425    let mut fields = Vec::new();
426
427    for (_, data) in body.basic_blocks.iter().enumerate() {
428        for statement in &data.statements {
429            if let StatementKind::Assign(assign) = &statement.kind {
430                let (place, _) = &**assign;
431                if place.local.as_usize() != 1 {
432                    continue;
433                }
434                let mut saw_deref = false;
435                for proj in place.projection.iter() {
436                    match proj {
437                        ProjectionElem::Deref => {
438                            saw_deref = true;
439                        }
440                        ProjectionElem::Field(index, _) if saw_deref => {
441                            let idx = index.as_usize();
442                            if !fields.contains(&idx) {
443                                fields.push(idx);
444                            }
445                        }
446                        _ => {}
447                    }
448                }
449            }
450        }
451    }
452
453    fields
454}
455
456pub fn is_externally_reachable(tcx: TyCtxt<'_>, def_id: DefId) -> bool {
457    let Some(local) = def_id.as_local() else {
458        return true;
459    };
460    tcx.effective_visibilities(()).is_reachable(local)
461}