Skip to main content

charon_driver/hax/
rustc_utils.rs

1use crate::hax::prelude::*;
2use rustc_hir::def::DefKind as RDefKind;
3use rustc_middle::{mir, ty};
4
5pub fn inst_binder<'tcx, T>(
6    tcx: ty::TyCtxt<'tcx>,
7    typing_env: ty::TypingEnv<'tcx>,
8    args: Option<ty::GenericArgsRef<'tcx>>,
9    x: ty::EarlyBinder<'tcx, T>,
10) -> T
11where
12    T: ty::TypeFoldable<ty::TyCtxt<'tcx>> + Clone,
13{
14    match args {
15        None => x.instantiate_identity(),
16        Some(args) => normalize(tcx, typing_env, x.instantiate(tcx, args)),
17    }
18}
19
20pub fn substitute<'tcx, T>(
21    tcx: ty::TyCtxt<'tcx>,
22    typing_env: ty::TypingEnv<'tcx>,
23    args: Option<ty::GenericArgsRef<'tcx>>,
24    x: T,
25) -> T
26where
27    T: ty::TypeFoldable<ty::TyCtxt<'tcx>>,
28{
29    inst_binder(tcx, typing_env, args, ty::EarlyBinder::bind(x))
30}
31
32#[extension_traits::extension(pub trait SubstBinder)]
33impl<'tcx, T: ty::TypeFoldable<ty::TyCtxt<'tcx>>> ty::Binder<'tcx, T> {
34    fn subst(
35        self,
36        tcx: ty::TyCtxt<'tcx>,
37        generics: &[ty::GenericArg<'tcx>],
38    ) -> ty::Binder<'tcx, T> {
39        ty::EarlyBinder::bind(self).instantiate(tcx, generics)
40    }
41}
42
43/// Whether the item can have generic parameters.
44pub(crate) fn can_have_generics<'tcx>(tcx: ty::TyCtxt<'tcx>, def_id: RDefId) -> bool {
45    use RDefKind::*;
46    !matches!(
47        get_def_kind(tcx, def_id),
48        ConstParam
49            | ExternCrate
50            | ForeignMod
51            | GlobalAsm
52            | LifetimeParam
53            | Macro(..)
54            | Mod
55            | TyParam
56            | Use
57    )
58}
59
60pub(crate) fn get_variant_kind<'s, S: UnderOwnerState<'s>>(
61    adt_def: &ty::AdtDef<'s>,
62    variant_index: rustc_abi::VariantIdx,
63    _s: &S,
64) -> VariantKind {
65    if adt_def.is_struct() {
66        VariantKind::Struct
67    } else if adt_def.is_union() {
68        VariantKind::Union
69    } else {
70        let index = variant_index;
71        VariantKind::Enum { index }
72    }
73}
74
75pub trait HasParamEnv<'tcx> {
76    fn param_env(&self) -> ty::ParamEnv<'tcx>;
77    fn typing_env(&self) -> ty::TypingEnv<'tcx>;
78}
79
80impl<'tcx, S: UnderOwnerState<'tcx>> HasParamEnv<'tcx> for S {
81    fn param_env(&self) -> ty::ParamEnv<'tcx> {
82        let tcx = self.base().tcx;
83        let def_id = self.owner_id();
84        if can_have_generics(tcx, def_id) {
85            tcx.param_env(def_id)
86        } else {
87            ty::ParamEnv::empty()
88        }
89    }
90    fn typing_env(&self) -> ty::TypingEnv<'tcx> {
91        ty::TypingEnv {
92            param_env: self.param_env(),
93            typing_mode: ty::TypingMode::PostAnalysis,
94        }
95    }
96}
97
98/// Gets the children of a module.
99pub fn get_mod_children<'tcx>(
100    tcx: ty::TyCtxt<'tcx>,
101    def_id: RDefId,
102) -> Vec<(Option<rustc_span::Ident>, RDefId)> {
103    match def_id.as_local() {
104        Some(ldid) => match tcx.hir_node_by_def_id(ldid) {
105            rustc_hir::Node::Crate(m)
106            | rustc_hir::Node::Item(&rustc_hir::Item {
107                kind: rustc_hir::ItemKind::Mod(_, m),
108                ..
109            }) => m
110                .item_ids
111                .iter()
112                .map(|&item_id| {
113                    let opt_ident = tcx.hir_item(item_id).kind.ident();
114                    let def_id = item_id.owner_id.to_def_id();
115                    (opt_ident, def_id)
116                })
117                .collect(),
118            node => panic!("DefKind::Module is an unexpected node: {node:?}"),
119        },
120        None => tcx
121            .module_children(def_id)
122            .iter()
123            .map(|child| (Some(child.ident), child.res.def_id()))
124            .collect(),
125    }
126}
127
128/// Gets the children of an `extern` block. Empty if the block is not defined in the current crate.
129pub fn get_foreign_mod_children<'tcx>(tcx: ty::TyCtxt<'tcx>, def_id: RDefId) -> Vec<RDefId> {
130    match def_id.as_local() {
131        Some(ldid) => tcx
132            .hir_node_by_def_id(ldid)
133            .expect_item()
134            .expect_foreign_mod()
135            .1
136            .iter()
137            .map(|foreign_item_ref| foreign_item_ref.owner_id.to_def_id())
138            .collect(),
139        None => vec![],
140    }
141}
142
143/// The signature of a method impl may be a subtype of the one expected from the trait decl, as in
144/// the example below. For correctness, we must be able to map from the method generics declared in
145/// the trait to the actual method generics. Because this would require type inference, we instead
146/// simply return the declared signature. This will cause issues if it is possible to use such a
147/// more-specific implementation with its more-specific type, but we have a few other issues with
148/// lifetime-generic function pointers anyway so this is unlikely to cause problems.
149///
150/// ```ignore
151/// trait MyCompare<Other>: Sized {
152///     fn compare(self, other: Other) -> bool;
153/// }
154/// impl<'a> MyCompare<&'a ()> for &'a () {
155///     // This implementation is more general because it works for non-`'a` refs. Note that only
156///     // late-bound vars may differ in this way.
157///     // `<&'a () as MyCompare<&'a ()>>::compare` has type `fn<'b>(&'a (), &'b ()) -> bool`,
158///     // but type `fn(&'a (), &'a ()) -> bool` was expected from the trait declaration.
159///     fn compare<'b>(self, _other: &'b ()) -> bool {
160///         true
161///     }
162/// }
163/// ```
164pub fn get_method_sig<'tcx>(
165    tcx: ty::TyCtxt<'tcx>,
166    typing_env: ty::TypingEnv<'tcx>,
167    def_id: RDefId,
168    method_args: Option<ty::GenericArgsRef<'tcx>>,
169) -> ty::PolyFnSig<'tcx> {
170    let real_sig = inst_binder(tcx, typing_env, method_args, tcx.fn_sig(def_id));
171    let item = tcx.associated_item(def_id);
172    let ty::AssocContainer::TraitImpl(Ok(decl_method_id)) = item.container else {
173        return real_sig;
174    };
175    let declared_sig = tcx.fn_sig(decl_method_id);
176
177    // TODO(Nadrieril): Temporary hack: if the signatures have the same number of bound vars, we
178    // keep the real signature. While the declared signature is more correct, it is also less
179    // normalized and we can't normalize without erasing regions but regions are crucial in
180    // function signatures. Hence we cheat here, until charon gains proper normalization
181    // capabilities.
182    if declared_sig.skip_binder().bound_vars().len() == real_sig.bound_vars().len() {
183        return real_sig;
184    }
185
186    let impl_def_id = item.container_id(tcx);
187    let method_args =
188        method_args.unwrap_or_else(|| ty::GenericArgs::identity_for_item(tcx, def_id));
189    // The trait predicate that is implemented by the surrounding impl block.
190    let implemented_trait_ref = tcx
191        .impl_trait_ref(impl_def_id)
192        .instantiate(tcx, method_args);
193    // Construct arguments for the declared method generics in the context of the implemented
194    // method generics.
195    let decl_args = method_args.rebase_onto(tcx, impl_def_id, implemented_trait_ref.args);
196    let sig = declared_sig.instantiate(tcx, decl_args);
197    // Avoids accidentally using the same lifetime name twice in the same scope
198    // (once in impl parameters, second in the method declaration late-bound vars).
199    let sig = tcx.anonymize_bound_vars(sig);
200    normalize(tcx, typing_env, sig)
201}
202
203/// Generates a list of `<trait_ref>::Ty` type aliases for each non-gat associated type of the
204/// given trait and its parents, in a specific order.
205pub fn assoc_tys_for_trait<'tcx>(
206    tcx: ty::TyCtxt<'tcx>,
207    typing_env: ty::TypingEnv<'tcx>,
208    tref: ty::TraitRef<'tcx>,
209) -> Vec<ty::AliasTy<'tcx>> {
210    fn gather_assoc_tys<'tcx>(
211        tcx: ty::TyCtxt<'tcx>,
212        typing_env: ty::TypingEnv<'tcx>,
213        assoc_tys: &mut Vec<ty::AliasTy<'tcx>>,
214        tref: ty::TraitRef<'tcx>,
215    ) {
216        assoc_tys.extend(
217            tcx.associated_items(tref.def_id)
218                .in_definition_order()
219                .filter(|assoc| matches!(assoc.kind, ty::AssocKind::Type { .. }))
220                .filter(|assoc| tcx.generics_of(assoc.def_id).own_params.is_empty())
221                .map(|assoc| ty::AliasTy::new(tcx, assoc.def_id, tref.args)),
222        );
223        for clause in tcx
224            .explicit_super_predicates_of(tref.def_id)
225            .map_bound(|clauses| clauses.iter().map(|(clause, _span)| *clause))
226            .iter_instantiated(tcx, tref.args)
227        {
228            if let Some(pred) = clause.as_trait_clause() {
229                let tref = erase_and_norm(tcx, typing_env, pred.skip_binder().trait_ref);
230                gather_assoc_tys(tcx, typing_env, assoc_tys, tref);
231            }
232        }
233    }
234    let mut ret = vec![];
235    gather_assoc_tys(tcx, typing_env, &mut ret, tref);
236    ret
237}
238
239/// Generates a `dyn Trait<Args.., Ty = <Self as Trait>::Ty..>` type for the given trait ref.
240pub fn dyn_self_ty<'tcx>(
241    tcx: ty::TyCtxt<'tcx>,
242    typing_env: ty::TypingEnv<'tcx>,
243    tref: ty::TraitRef<'tcx>,
244) -> Option<ty::Ty<'tcx>> {
245    let re_erased = tcx.lifetimes.re_erased;
246    if !tcx.is_dyn_compatible(tref.def_id) {
247        return None;
248    }
249
250    // The main `Trait<Args>` predicate.
251    let main_pred = ty::Binder::dummy(ty::ExistentialPredicate::Trait(
252        ty::ExistentialTraitRef::erase_self_ty(tcx, tref),
253    ));
254
255    let ty_constraints = assoc_tys_for_trait(tcx, typing_env, tref)
256        .into_iter()
257        .map(|alias_ty| {
258            let proj = ty::ProjectionPredicate {
259                projection_term: alias_ty.into(),
260                term: ty::Ty::new_alias(tcx, ty::Projection, alias_ty).into(),
261            };
262            let proj = ty::ExistentialProjection::erase_self_ty(tcx, proj);
263            ty::Binder::dummy(ty::ExistentialPredicate::Projection(proj))
264        });
265
266    let preds = {
267        // Stable sort predicates to prevent platform-specific ordering issues
268        let mut preds: Vec<_> = [main_pred].into_iter().chain(ty_constraints).collect();
269        preds.sort_by(|a, b| {
270            use rustc_middle::ty::ExistentialPredicateStableCmpExt;
271            a.skip_binder().stable_cmp(tcx, &b.skip_binder())
272        });
273        tcx.mk_poly_existential_predicates(&preds)
274    };
275    let ty = tcx.mk_ty_from_kind(ty::Dynamic(preds, re_erased));
276    let ty = normalize(tcx, typing_env, ty);
277    Some(ty)
278}
279
280pub fn closure_once_shim<'tcx>(
281    tcx: ty::TyCtxt<'tcx>,
282    closure_ty: ty::Ty<'tcx>,
283) -> Option<mir::Body<'tcx>> {
284    let ty::Closure(def_id, args) = closure_ty.kind() else {
285        unreachable!()
286    };
287    let instance = match args.as_closure().kind() {
288        ty::ClosureKind::Fn | ty::ClosureKind::FnMut => {
289            ty::Instance::fn_once_adapter_instance(tcx, *def_id, args)
290        }
291        ty::ClosureKind::FnOnce => return None,
292    };
293    let mir = tcx.instance_mir(instance.def).clone();
294    let mir = ty::EarlyBinder::bind(mir).instantiate(tcx, instance.args);
295    Some(mir)
296}
297
298pub fn drop_glue_shim<'tcx>(
299    tcx: ty::TyCtxt<'tcx>,
300    def_id: RDefId,
301    instantiate: Option<ty::GenericArgsRef<'tcx>>,
302) -> mir::Body<'tcx> {
303    let drop_in_place =
304        tcx.require_lang_item(rustc_hir::LangItem::DropInPlace, rustc_span::DUMMY_SP);
305    let ty = tcx.type_of(def_id);
306    let ty = match instantiate {
307        None => ty.instantiate_identity(),
308        Some(args) => ty.instantiate(tcx, args),
309    };
310    let instance_kind = ty::InstanceKind::DropGlue(drop_in_place, Some(ty));
311    tcx.instance_mir(instance_kind).clone()
312}