charon_driver/hax/
rustc_utils.rs

1use std::collections::{HashMap, HashSet};
2
3use rustc_hir::def::DefKind as RDefKind;
4use rustc_middle::infer::canonical::{CanonicalVarKinds, CanonicalVarValues};
5use rustc_middle::ty::relate::{
6    Relate, RelateResult, TypeRelation, relate_args_with_variances, structurally_relate_consts,
7    structurally_relate_tys,
8};
9use rustc_middle::{mir, ty};
10use rustc_span::kw;
11use rustc_type_ir::Interner;
12
13use crate::hax::prelude::*;
14
15/// Structurally pairs a canonicalized type with its inferred type to recover the values rustc
16/// replaced with canonical variables. Inference regions are deliberately erased instead.
17pub fn match_canonical_var_values<'tcx>(
18    tcx: ty::TyCtxt<'tcx>,
19    var_kinds: CanonicalVarKinds<'tcx>,
20    canonical_ty: ty::Ty<'tcx>,
21    inferred_ty: ty::Ty<'tcx>,
22) -> Option<CanonicalVarValues<'tcx>> {
23    let values = var_kinds
24        .iter()
25        .map(|kind| match kind {
26            ty::CanonicalVarKind::Region(_) => Some(tcx.lifetimes.re_erased.into()),
27            _ => None,
28        })
29        .collect();
30    let mut matcher = CanonicalVarMatcher { tcx, values };
31    matcher.relate(canonical_ty, inferred_ty).ok()?;
32
33    // A canonical type variable may be equated to an earlier canonical root without occurring
34    // independently in the canonicalized type.
35    for (index, kind) in var_kinds.iter().enumerate() {
36        if matcher.values[index].is_none()
37            && let ty::CanonicalVarKind::Ty { sub_root, .. } = kind
38        {
39            matcher.values[index] = matcher.values[sub_root.as_usize()];
40        }
41    }
42    let values = matcher.values.into_iter().collect::<Option<Vec<_>>>()?;
43    Some(CanonicalVarValues {
44        var_values: tcx.mk_args(&values),
45    })
46}
47
48struct CanonicalVarMatcher<'tcx> {
49    tcx: ty::TyCtxt<'tcx>,
50    values: Vec<Option<ty::GenericArg<'tcx>>>,
51}
52
53impl<'tcx> CanonicalVarMatcher<'tcx> {
54    fn record<T>(
55        &mut self,
56        var: ty::BoundVar,
57        value: ty::GenericArg<'tcx>,
58        result: T,
59    ) -> RelateResult<'tcx, T> {
60        let slot = &mut self.values[var.as_usize()];
61        if slot.is_some_and(|old| old != value) {
62            Err(ty::error::TypeError::Mismatch)
63        } else {
64            *slot = Some(value);
65            Ok(result)
66        }
67    }
68}
69
70impl<'tcx> TypeRelation<ty::TyCtxt<'tcx>> for CanonicalVarMatcher<'tcx> {
71    fn cx(&self) -> ty::TyCtxt<'tcx> {
72        self.tcx
73    }
74
75    fn relate_ty_args(
76        &mut self,
77        a_ty: ty::Ty<'tcx>,
78        _: ty::Ty<'tcx>,
79        _: rustc_hir::def_id::DefId,
80        a_args: ty::GenericArgsRef<'tcx>,
81        b_args: ty::GenericArgsRef<'tcx>,
82        _: impl FnOnce(ty::GenericArgsRef<'tcx>) -> ty::Ty<'tcx>,
83    ) -> RelateResult<'tcx, ty::Ty<'tcx>> {
84        ty::relate::relate_args_invariantly(self, a_args, b_args)?;
85        Ok(a_ty)
86    }
87
88    fn relate_with_variance<T: Relate<ty::TyCtxt<'tcx>>>(
89        &mut self,
90        _: ty::Variance,
91        _: ty::VarianceDiagInfo<ty::TyCtxt<'tcx>>,
92        a: T,
93        b: T,
94    ) -> RelateResult<'tcx, T> {
95        self.relate(a, b)
96    }
97
98    fn tys(&mut self, a: ty::Ty<'tcx>, b: ty::Ty<'tcx>) -> RelateResult<'tcx, ty::Ty<'tcx>> {
99        if let ty::Bound(ty::BoundVarIndexKind::Canonical, bound) = *a.kind() {
100            self.record(bound.var, b.into(), a)
101        } else {
102            structurally_relate_tys(self, a, b)
103        }
104    }
105
106    fn regions(
107        &mut self,
108        a: ty::Region<'tcx>,
109        _: ty::Region<'tcx>,
110    ) -> RelateResult<'tcx, ty::Region<'tcx>> {
111        Ok(a)
112    }
113
114    fn consts(
115        &mut self,
116        a: ty::Const<'tcx>,
117        b: ty::Const<'tcx>,
118    ) -> RelateResult<'tcx, ty::Const<'tcx>> {
119        if let ty::ConstKind::Bound(ty::BoundVarIndexKind::Canonical, bound) = a.kind() {
120            self.record(bound.var, b.into(), a)
121        } else {
122            structurally_relate_consts(self, a, b)
123        }
124    }
125
126    fn binders<T>(
127        &mut self,
128        a: ty::Binder<'tcx, T>,
129        b: ty::Binder<'tcx, T>,
130    ) -> RelateResult<'tcx, ty::Binder<'tcx, T>>
131    where
132        T: Relate<ty::TyCtxt<'tcx>>,
133    {
134        self.relate(a.skip_binder(), b.skip_binder())?;
135        Ok(a)
136    }
137}
138
139/// Computes the variance of each region bound by `sig`.
140///
141/// This follows rustc's `FunctionalVariances` helper from the `impl_trait_overcaptures` lint.
142pub fn fn_sig_bound_region_variances<'tcx>(
143    tcx: ty::TyCtxt<'tcx>,
144    sig: ty::PolyFnSig<'tcx>,
145) -> HashMap<ty::BoundVar, ty::Variance> {
146    struct FunctionalVariances<'tcx> {
147        tcx: ty::TyCtxt<'tcx>,
148        variances: HashMap<ty::BoundVar, ty::Variance>,
149        ambient_variance: ty::Variance,
150        target_binder: ty::DebruijnIndex,
151    }
152
153    impl<'tcx> TypeRelation<ty::TyCtxt<'tcx>> for FunctionalVariances<'tcx> {
154        fn cx(&self) -> ty::TyCtxt<'tcx> {
155            self.tcx
156        }
157
158        fn relate_ty_args(
159            &mut self,
160            a_ty: ty::Ty<'tcx>,
161            _: ty::Ty<'tcx>,
162            def_id: rustc_hir::def_id::DefId,
163            a_args: ty::GenericArgsRef<'tcx>,
164            b_args: ty::GenericArgsRef<'tcx>,
165            _: impl FnOnce(ty::GenericArgsRef<'tcx>) -> ty::Ty<'tcx>,
166        ) -> RelateResult<'tcx, ty::Ty<'tcx>> {
167            relate_args_with_variances(self, self.tcx.variances_of(def_id), a_args, b_args)?;
168            Ok(a_ty)
169        }
170
171        fn relate_with_variance<T: Relate<ty::TyCtxt<'tcx>>>(
172            &mut self,
173            variance: ty::Variance,
174            _: ty::VarianceDiagInfo<ty::TyCtxt<'tcx>>,
175            a: T,
176            b: T,
177        ) -> RelateResult<'tcx, T> {
178            let old_variance = self.ambient_variance;
179            self.ambient_variance = self.ambient_variance.xform(variance);
180            let result = self.relate(a, b);
181            self.ambient_variance = old_variance;
182            result
183        }
184
185        fn tys(&mut self, a: ty::Ty<'tcx>, b: ty::Ty<'tcx>) -> RelateResult<'tcx, ty::Ty<'tcx>> {
186            structurally_relate_tys(self, a, b)
187        }
188
189        fn regions(
190            &mut self,
191            a: ty::Region<'tcx>,
192            _: ty::Region<'tcx>,
193        ) -> RelateResult<'tcx, ty::Region<'tcx>> {
194            if let ty::ReBound(ty::BoundVarIndexKind::Bound(binder), ty::BoundRegion { var, .. }) =
195                a.kind()
196                && binder == self.target_binder
197            {
198                self.variances
199                    .entry(var)
200                    .and_modify(|old| *old = unify_variances(*old, self.ambient_variance))
201                    .or_insert(self.ambient_variance);
202            }
203            Ok(a)
204        }
205
206        fn consts(
207            &mut self,
208            a: ty::Const<'tcx>,
209            b: ty::Const<'tcx>,
210        ) -> RelateResult<'tcx, ty::Const<'tcx>> {
211            structurally_relate_consts(self, a, b)
212        }
213
214        fn binders<T>(
215            &mut self,
216            a: ty::Binder<'tcx, T>,
217            b: ty::Binder<'tcx, T>,
218        ) -> RelateResult<'tcx, ty::Binder<'tcx, T>>
219        where
220            T: Relate<ty::TyCtxt<'tcx>>,
221        {
222            let old_target_binder = self.target_binder;
223            self.target_binder = self.target_binder.shifted_in(1);
224            let result = self.relate(a.skip_binder(), b.skip_binder());
225            self.target_binder = old_target_binder;
226            result?;
227            Ok(a)
228        }
229    }
230
231    fn unify_variances(a: ty::Variance, b: ty::Variance) -> ty::Variance {
232        match (a, b) {
233            (ty::Bivariant, other) | (other, ty::Bivariant) => other,
234            (ty::Invariant, _) | (_, ty::Invariant) => ty::Invariant,
235            (ty::Contravariant, ty::Covariant) | (ty::Covariant, ty::Contravariant) => {
236                ty::Invariant
237            }
238            (ty::Contravariant, ty::Contravariant) => ty::Contravariant,
239            (ty::Covariant, ty::Covariant) => ty::Covariant,
240        }
241    }
242
243    let mut relation = FunctionalVariances {
244        tcx,
245        variances: HashMap::new(),
246        // `Relate for FnSig` makes inputs contravariant. We are computing the variance of the
247        // signature's own bound parameters, so cancel that outer function-type variance. Nested
248        // function types still introduce their own contravariance as usual.
249        ambient_variance: ty::Contravariant,
250        target_binder: ty::INNERMOST,
251    };
252    relation
253        .relate(sig.skip_binder(), sig.skip_binder())
254        .expect("a signature must relate to itself");
255
256    // A bound region that does not occur in the signature is bivariant.
257    for (index, var) in sig.bound_vars().iter().enumerate() {
258        if matches!(var, ty::BoundVariableKind::Region(_)) {
259            relation
260                .variances
261                .entry(ty::BoundVar::from_usize(index))
262                .or_insert(ty::Bivariant);
263        }
264    }
265    relation.variances
266}
267
268pub fn inst_binder<'tcx, T>(
269    tcx: ty::TyCtxt<'tcx>,
270    typing_env: ty::TypingEnv<'tcx>,
271    args: Option<ty::GenericArgsRef<'tcx>>,
272    x: ty::EarlyBinder<'tcx, T>,
273) -> T
274where
275    T: ty::TypeFoldable<ty::TyCtxt<'tcx>> + Clone,
276{
277    match args {
278        None => x.instantiate_identity().skip_normalization(),
279        Some(args) => normalize(tcx, typing_env, x.instantiate(tcx, args)),
280    }
281}
282
283pub fn substitute<'tcx, T>(
284    tcx: ty::TyCtxt<'tcx>,
285    typing_env: ty::TypingEnv<'tcx>,
286    args: Option<ty::GenericArgsRef<'tcx>>,
287    x: T,
288) -> T
289where
290    T: ty::TypeFoldable<ty::TyCtxt<'tcx>>,
291{
292    inst_binder(tcx, typing_env, args, ty::EarlyBinder::bind(x))
293}
294
295/// Make a new `ParamEnv` from a list of clauses.
296pub fn param_env_from_clauses<'tcx>(
297    tcx: ty::TyCtxt<'tcx>,
298    predicates: impl Iterator<Item = ty::Clause<'tcx>>,
299) -> ty::ParamEnv<'tcx> {
300    let cause = rustc_trait_selection::traits::ObligationCause::dummy();
301    let param_env = ty::ParamEnv::new(tcx.mk_clauses_from_iter(predicates));
302    rustc_trait_selection::traits::normalize_param_env_or_error(tcx, param_env, cause)
303}
304
305#[extension_traits::extension(pub trait SubstBinder)]
306impl<'tcx, T: ty::TypeFoldable<ty::TyCtxt<'tcx>>> ty::Binder<'tcx, T> {
307    fn subst(
308        self,
309        tcx: ty::TyCtxt<'tcx>,
310        generics: &[ty::GenericArg<'tcx>],
311    ) -> ty::Binder<'tcx, T> {
312        ty::EarlyBinder::bind(self)
313            .instantiate(tcx, generics)
314            .skip_normalization()
315    }
316}
317
318/// Whether the item can have generic parameters.
319pub(crate) fn can_have_generics<'tcx>(tcx: ty::TyCtxt<'tcx>, def_id: RDefId) -> bool {
320    use RDefKind::*;
321    !matches!(
322        get_def_kind(tcx, def_id),
323        ConstParam
324            | ExternCrate
325            | ForeignMod
326            | GlobalAsm
327            | LifetimeParam
328            | Macro(..)
329            | Mod
330            | TyParam
331            | Use
332    )
333}
334
335pub(crate) fn get_variant_kind<'s, S: UnderOwnerState<'s>>(
336    adt_def: &ty::AdtDef<'s>,
337    variant_index: rustc_abi::VariantIdx,
338    _s: &S,
339) -> VariantKind {
340    if adt_def.is_struct() {
341        VariantKind::Struct
342    } else if adt_def.is_union() {
343        VariantKind::Union
344    } else {
345        let index = variant_index;
346        VariantKind::Enum { index }
347    }
348}
349
350/// Gets the children of a module.
351pub fn get_mod_children<'tcx>(
352    tcx: ty::TyCtxt<'tcx>,
353    def_id: RDefId,
354) -> Vec<(Option<rustc_span::Ident>, RDefId)> {
355    match def_id.as_local() {
356        Some(ldid) => match tcx.hir_node_by_def_id(ldid) {
357            rustc_hir::Node::Crate(m)
358            | rustc_hir::Node::Item(&rustc_hir::Item {
359                kind: rustc_hir::ItemKind::Mod(_, m),
360                ..
361            }) => m
362                .item_ids
363                .iter()
364                .map(|&item_id| {
365                    let opt_ident = tcx.hir_item(item_id).kind.ident();
366                    let def_id = item_id.owner_id.to_def_id();
367                    (opt_ident, def_id)
368                })
369                .collect(),
370            node => panic!("DefKind::Module is an unexpected node: {node:?}"),
371        },
372        None => tcx
373            .module_children(def_id)
374            .iter()
375            .filter_map(|child| Some((Some(child.ident), child.res.opt_def_id()?)))
376            .collect(),
377    }
378}
379
380/// Gets the children of an `extern` block. Empty if the block is not defined in the current crate.
381pub fn get_foreign_mod_children<'tcx>(tcx: ty::TyCtxt<'tcx>, def_id: RDefId) -> Vec<RDefId> {
382    match def_id.as_local() {
383        Some(ldid) => tcx
384            .hir_node_by_def_id(ldid)
385            .expect_item()
386            .expect_foreign_mod()
387            .1
388            .iter()
389            .map(|foreign_item_ref| foreign_item_ref.owner_id.to_def_id())
390            .collect(),
391        None => vec![],
392    }
393}
394
395/// The signature of a method impl may be a subtype of the one expected from the trait decl, as in
396/// the example below. For correctness, we must be able to map from the method generics declared in
397/// the trait to the actual method generics. Because this would require type inference, we instead
398/// simply return the declared signature. This will cause issues if it is possible to use such a
399/// more-specific implementation with its more-specific type, but we have a few other issues with
400/// lifetime-generic function pointers anyway so this is unlikely to cause problems.
401///
402/// ```ignore
403/// trait MyCompare<Other>: Sized {
404///     fn compare(self, other: Other) -> bool;
405/// }
406/// impl<'a> MyCompare<&'a ()> for &'a () {
407///     // This implementation is more general because it works for non-`'a` refs. Note that only
408///     // late-bound vars may differ in this way.
409///     // `<&'a () as MyCompare<&'a ()>>::compare` has type `fn<'b>(&'a (), &'b ()) -> bool`,
410///     // but type `fn(&'a (), &'a ()) -> bool` was expected from the trait declaration.
411///     fn compare<'b>(self, _other: &'b ()) -> bool {
412///         true
413///     }
414/// }
415/// ```
416pub fn get_method_sig<'tcx>(
417    tcx: ty::TyCtxt<'tcx>,
418    typing_env: ty::TypingEnv<'tcx>,
419    def_id: RDefId,
420    method_args: Option<ty::GenericArgsRef<'tcx>>,
421) -> ty::PolyFnSig<'tcx> {
422    let real_sig = inst_binder(tcx, typing_env, method_args, tcx.fn_sig(def_id));
423    let item = tcx.associated_item(def_id);
424    let ty::AssocContainer::TraitImpl(Ok(decl_method_id)) = item.container else {
425        return real_sig;
426    };
427    let declared_sig = tcx.fn_sig(decl_method_id);
428
429    let impl_def_id = item.container_id(tcx);
430    let method_args =
431        method_args.unwrap_or_else(|| ty::GenericArgs::identity_for_item(tcx, def_id));
432    // The trait predicate that is implemented by the surrounding impl block.
433    let implemented_trait_ref = tcx
434        .impl_trait_ref(impl_def_id)
435        .instantiate(tcx, method_args);
436    let implemented_trait_ref = normalize(tcx, typing_env, implemented_trait_ref);
437    // Construct arguments for the declared method generics in the context of the implemented
438    // method generics.
439    let decl_args = method_args.rebase_onto(tcx, impl_def_id, implemented_trait_ref.args);
440    let sig = declared_sig.instantiate(tcx, decl_args);
441    let sig = normalize(tcx, typing_env, sig);
442
443    if let container_named_lts = tcx
444        .generics_of(impl_def_id)
445        .own_params
446        .iter()
447        .filter(|p| matches!(p.kind, ty::GenericParamDefKind::Lifetime))
448        .filter(|p| p.name != kw::UnderscoreLifetime)
449        .map(|p| p.name)
450        .collect::<HashSet<_>>()
451        && sig
452            .bound_vars()
453            .iter()
454            .map(|v| v.expect_region())
455            .filter_map(|v| v.get_name(tcx))
456            .any(|lt| container_named_lts.contains(&lt))
457    {
458        // Avoids using the same lifetime name twice in the same scope (once in impl parameters,
459        // second in the method declaration late-bound vars).
460        tcx.anonymize_bound_vars(sig)
461    } else {
462        sig
463    }
464}
465
466/// Generates a list of `<trait_ref>::Ty` type aliases for each non-gat associated type of the
467/// given trait and its parents, in a specific order.
468pub fn assoc_tys_for_trait<'tcx>(
469    tcx: ty::TyCtxt<'tcx>,
470    typing_env: ty::TypingEnv<'tcx>,
471    tref: ty::TraitRef<'tcx>,
472) -> Vec<ty::AliasTy<'tcx>> {
473    fn gather_assoc_tys<'tcx>(
474        tcx: ty::TyCtxt<'tcx>,
475        typing_env: ty::TypingEnv<'tcx>,
476        assoc_tys: &mut Vec<ty::AliasTy<'tcx>>,
477        tref: ty::TraitRef<'tcx>,
478    ) {
479        assoc_tys.extend(
480            tcx.associated_items(tref.def_id)
481                .in_definition_order()
482                .filter(|assoc| matches!(assoc.kind, ty::AssocKind::Type { .. }))
483                .filter(|assoc| {
484                    tcx.generics_of(assoc.def_id).own_params.is_empty()
485                        && tcx.predicates_of(assoc.def_id).predicates.is_empty()
486                })
487                .map(|assoc| {
488                    ty::AliasTy::new(tcx, tcx.alias_ty_kind_from_def_id(assoc.def_id), tref.args)
489                }),
490        );
491        for clause in tcx
492            .explicit_super_predicates_of(tref.def_id)
493            .map_bound(|clauses| clauses.iter().map(|(clause, _span)| *clause))
494            .iter_instantiated(tcx, tref.args)
495        {
496            if let Some(pred) = clause.as_trait_clause() {
497                let tref = erase_and_norm(tcx, typing_env, pred.map(|b| b.skip_binder().trait_ref));
498                gather_assoc_tys(tcx, typing_env, assoc_tys, tref);
499            }
500        }
501    }
502    let mut ret = vec![];
503    gather_assoc_tys(tcx, typing_env, &mut ret, tref);
504    ret
505}
506
507/// Generates a `dyn Trait<Args.., Ty = <Self as Trait>::Ty..>` type for the given trait ref.
508pub fn dyn_self_ty<'tcx>(
509    tcx: ty::TyCtxt<'tcx>,
510    typing_env: ty::TypingEnv<'tcx>,
511    tref: ty::TraitRef<'tcx>,
512) -> Option<ty::Ty<'tcx>> {
513    let re_erased = tcx.lifetimes.re_erased;
514    if !tcx.is_dyn_compatible(tref.def_id) {
515        return None;
516    }
517
518    // The main `Trait<Args>` predicate.
519    let main_pred = ty::Binder::dummy(ty::ExistentialPredicate::Trait(
520        ty::ExistentialTraitRef::erase_self_ty(tcx, tref),
521    ));
522
523    let ty_constraints = assoc_tys_for_trait(tcx, typing_env, tref)
524        .into_iter()
525        .map(|alias_ty| {
526            let proj = ty::ProjectionPredicate {
527                projection_term: alias_ty.into(),
528                term: ty::Ty::new_alias(tcx, alias_ty).into(),
529            };
530            let proj = ty::ExistentialProjection::erase_self_ty(tcx, proj);
531            ty::Binder::dummy(ty::ExistentialPredicate::Projection(proj))
532        });
533
534    let preds = {
535        // Stable sort predicates to prevent platform-specific ordering issues
536        let mut preds: Vec<_> = [main_pred].into_iter().chain(ty_constraints).collect();
537        preds.sort_by(|a, b| {
538            use rustc_middle::ty::ExistentialPredicateStableCmpExt;
539            a.skip_binder().stable_cmp(tcx, &b.skip_binder())
540        });
541        tcx.mk_poly_existential_predicates(&preds)
542    };
543    let ty = tcx.mk_ty_from_kind(ty::Dynamic(preds, re_erased));
544    let ty = normalize(tcx, typing_env, ty::Unnormalized::new_wip(ty));
545    Some(ty)
546}
547
548pub fn closure_once_shim<'tcx>(
549    tcx: ty::TyCtxt<'tcx>,
550    closure_ty: ty::Ty<'tcx>,
551) -> Option<mir::Body<'tcx>> {
552    let ty::Closure(def_id, args) = closure_ty.kind() else {
553        unreachable!()
554    };
555    let instance = match args.as_closure().kind() {
556        ty::ClosureKind::Fn | ty::ClosureKind::FnMut => {
557            ty::Instance::fn_once_adapter_instance(tcx, *def_id, args)
558        }
559        ty::ClosureKind::FnOnce => return None,
560    };
561    let mir = tcx.instance_mir(instance.def).clone();
562    let mir = ty::EarlyBinder::bind(mir)
563        .instantiate(tcx, instance.args)
564        .skip_normalization();
565    Some(mir)
566}
567
568pub fn drop_glue_shim<'tcx>(
569    s: &impl UnderOwnerState<'tcx>,
570    def_id: &DefId,
571    instantiate: Option<ty::GenericArgsRef<'tcx>>,
572) -> mir::Body<'tcx> {
573    let tcx = s.base().tcx;
574    let drop_glue = tcx.require_lang_item(rustc_hir::LangItem::DropGlue, rustc_span::DUMMY_SP);
575    let ty = inst_binder(tcx, s.typing_env(), instantiate, def_id.type_of(s));
576    let mut body = rustc_mir_transform::build_drop_shim(tcx, drop_glue, Some(ty), s.typing_env());
577    // Set the mir phase so that charon knows the contained drops are precise.
578    body.phase = mir::MirPhase::Runtime(mir::RuntimePhase::Optimized);
579    body
580}
charon_driver/hax/rustc_utils.rs

charon_driver/hax/
rustc_utils.rs
