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

charon_driver/hax/
rustc_utils.rs
