Skip to main content

charon_driver/hax/
rustc_utils.rs

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