rapx/verify/
helpers.rs

1use rustc_abi::FieldIdx;
2use rustc_hir::{
3    ItemKind,
4    def_id::{DefId, LocalDefId},
5};
6use rustc_middle::{
7    mir::{BasicBlock, TerminatorKind},
8    ty::{Ty, TyCtxt, TyKind},
9};
10use rustc_span::Symbol;
11use syn::Expr;
12
13pub use crate::helpers::fn_info::parse_expr_into_number;
14pub use crate::helpers::mir_scan::{
15    Checkpoint, CheckpointKind, CheckpointLocation, collect_unsafe_callsites,
16};
17pub use crate::helpers::name::{
18    access_ident_recursive, get_cleaned_def_path_name, get_struct_self_ty, match_ty_with_ident,
19    parse_signature,
20};
21
22/// Collect all return basic block indices for a function body.
23pub fn collect_return_block_indices(tcx: TyCtxt<'_>, def_id: DefId) -> Vec<BasicBlock> {
24    let mut blocks = Vec::new();
25    if !tcx.is_mir_available(def_id) {
26        return blocks;
27    }
28    let body = tcx.optimized_mir(def_id);
29    for (bb, data) in body.basic_blocks.iter_enumerated() {
30        if matches!(data.terminator().kind, TerminatorKind::Return) {
31            blocks.push(bb);
32        }
33    }
34    blocks
35}
36
37pub fn parse_expr_into_local_and_ty<'tcx>(
38    tcx: TyCtxt<'tcx>,
39    def_id: DefId,
40    expr: &Expr,
41) -> Option<(usize, Vec<(usize, Ty<'tcx>)>, Ty<'tcx>)> {
42    if let Some((base_ident, fields)) = access_ident_recursive(expr) {
43        let (param_names, param_tys) = parse_signature(tcx, def_id);
44        if param_names[0] != "0" {
45            if let Some(param_index) = param_names.iter().position(|name| name == &base_ident) {
46                return resolve_projection_from_base_ident(
47                    tcx,
48                    base_ident,
49                    fields,
50                    param_index + 1,
51                    param_tys[param_index],
52                );
53            }
54        }
55
56        if let Some(struct_ty) = get_struct_self_ty(tcx, def_id) {
57            return resolve_projection_from_struct_ident(tcx, base_ident, fields, struct_ty);
58        }
59    }
60    None
61}
62
63/// Return the callee argument index represented by a MIR local.
64///
65/// Contract annotations written with parameter names are parsed in the callee's
66/// local namespace.  MIR local `_0` is the return place and argument locals are
67/// `_1..=_arg_count`, so callee local `_1` denotes checkpoint argument `0`.
68pub fn callee_param_index_for_local(tcx: TyCtxt<'_>, callee: DefId, local: usize) -> Option<usize> {
69    if local == 0 {
70        return None;
71    }
72
73    let arg_count = if tcx.is_mir_available(callee) {
74        tcx.optimized_mir(callee).arg_count
75    } else {
76        tcx.fn_sig(callee)
77            .skip_binder()
78            .inputs()
79            .skip_binder()
80            .len()
81    };
82
83    (local <= arg_count).then_some(local - 1)
84}
85
86pub fn is_std_crate_def_id(tcx: TyCtxt<'_>, def_id: DefId) -> bool {
87    matches!(
88        tcx.crate_name(def_id.krate).as_str(),
89        "core" | "std" | "alloc"
90    )
91}
92
93pub fn is_trait_unsafe(tcx: TyCtxt<'_>, trait_def_id: DefId) -> bool {
94    let Some(local_id) = trait_def_id.as_local() else {
95        return false;
96    };
97    let item = tcx.hir_expect_item(local_id);
98
99    #[cfg(not(rapx_rustc_ge_198))]
100    if let ItemKind::Trait(_, _, unsafety, _, _, _, _) = &item.kind {
101        return matches!(unsafety, rustc_hir::Safety::Unsafe);
102    }
103    #[cfg(rapx_rustc_ge_198)]
104    if let ItemKind::Trait { safety, .. } = &item.kind {
105        return matches!(safety, rustc_hir::Safety::Unsafe);
106    }
107
108    false
109}
110
111pub fn resolve_impl_self_ty_def_id(item: &rustc_hir::Item<'_>) -> Option<DefId> {
112    let ItemKind::Impl(rustc_hir::Impl { self_ty, .. }) = &item.kind else {
113        return None;
114    };
115    match &self_ty.kind {
116        rustc_hir::TyKind::Path(rustc_hir::QPath::Resolved(_, path)) => match path.res {
117            rustc_hir::def::Res::Def(
118                rustc_hir::def::DefKind::Struct
119                | rustc_hir::def::DefKind::Enum
120                | rustc_hir::def::DefKind::Union,
121                def_id,
122            ) => Some(def_id),
123            _ => None,
124        },
125        _ => None,
126    }
127}
128
129pub fn has_rapx_verify_attr(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
130    let hir_id = tcx.local_def_id_to_hir_id(def_id);
131
132    let rapx = Symbol::intern("rapx");
133    let verify = Symbol::intern("verify");
134
135    let attrs = tcx.hir_attrs(hir_id);
136
137    attrs.iter().any(|attr| {
138        #[cfg(rapx_rustc_ge_193)]
139        if attr.is_doc_comment().is_some() {
140            return false;
141        }
142        #[cfg(not(rapx_rustc_ge_193))]
143        if attr.is_doc_comment() {
144            return false;
145        }
146
147        let path = attr.path();
148
149        path.len() == 2 && path[0] == rapx && path[1] == verify
150    })
151}
152
153pub fn get_owner_struct_def_id(tcx: TyCtxt<'_>, def_id: DefId) -> Option<DefId> {
154    let assoc_item = tcx.opt_associated_item(def_id)?;
155    let impl_id = assoc_item.impl_container(tcx)?;
156    let self_ty = tcx.type_of(impl_id).skip_binder();
157
158    match self_ty.kind() {
159        TyKind::Adt(adt_def, _) => Some(adt_def.did()),
160        _ => None,
161    }
162}
163
164fn resolve_projection_from_base_ident<'tcx>(
165    tcx: TyCtxt<'tcx>,
166    base_ident: String,
167    fields: Vec<String>,
168    base_local: usize,
169    base_ty: Ty<'tcx>,
170) -> Option<(usize, Vec<(usize, Ty<'tcx>)>, Ty<'tcx>)> {
171    let mut current_ty = base_ty;
172    let mut field_indices = Vec::new();
173    for field_name in fields {
174        let Some((field_idx, field_ty)) = resolve_next_field(tcx, current_ty, &field_name) else {
175            return if field_indices.is_empty() && base_ident.is_empty() {
176                None
177            } else {
178                None
179            };
180        };
181        current_ty = field_ty;
182        field_indices.push((field_idx, current_ty));
183    }
184    Some((base_local, field_indices, current_ty))
185}
186
187fn resolve_projection_from_struct_ident<'tcx>(
188    tcx: TyCtxt<'tcx>,
189    base_ident: String,
190    fields: Vec<String>,
191    struct_ty: Ty<'tcx>,
192) -> Option<(usize, Vec<(usize, Ty<'tcx>)>, Ty<'tcx>)> {
193    let Some((field_idx, field_ty)) = resolve_next_field(tcx, struct_ty, &base_ident) else {
194        return None;
195    };
196
197    let mut current_ty = field_ty;
198    let mut field_indices = vec![(field_idx, current_ty)];
199    for field_name in fields {
200        let Some((next_field_idx, next_field_ty)) =
201            resolve_next_field(tcx, current_ty, &field_name)
202        else {
203            return None;
204        };
205        current_ty = next_field_ty;
206        field_indices.push((next_field_idx, current_ty));
207    }
208
209    Some((1, field_indices, current_ty))
210}
211
212fn resolve_next_field<'tcx>(
213    tcx: TyCtxt<'tcx>,
214    base_ty: Ty<'tcx>,
215    field_name: &str,
216) -> Option<(usize, Ty<'tcx>)> {
217    let peeled_ty = base_ty.peel_refs();
218    if let TyKind::Adt(adt_def, arg_list) = *peeled_ty.kind() {
219        if !adt_def.is_struct() && !adt_def.is_union() {
220            return None;
221        }
222        let variant = adt_def.non_enum_variant();
223        if let Ok(field_idx) = field_name.parse::<usize>() {
224            if field_idx < variant.fields.len() {
225                #[cfg(not(rapx_rustc_ge_198))]
226                let field_ty = variant.fields[FieldIdx::from_usize(field_idx)].ty(tcx, arg_list);
227                #[cfg(rapx_rustc_ge_198)]
228                let field_ty = variant.fields[FieldIdx::from_usize(field_idx)]
229                    .ty(tcx, arg_list)
230                    .skip_norm_wip();
231                return Some((field_idx, field_ty));
232            }
233        }
234        if let Some((idx, _)) = variant
235            .fields
236            .iter()
237            .enumerate()
238            .find(|(_, f)| f.ident(tcx).name.to_string() == field_name)
239        {
240            #[cfg(not(rapx_rustc_ge_198))]
241            let field_ty = variant.fields[FieldIdx::from_usize(idx)].ty(tcx, arg_list);
242            #[cfg(rapx_rustc_ge_198)]
243            let field_ty = variant.fields[FieldIdx::from_usize(idx)]
244                .ty(tcx, arg_list)
245                .skip_norm_wip();
246            return Some((idx, field_ty));
247        }
248    }
249    None
250}