charon_lib/transform/normalize/
skip_trait_refs_when_known.rs

1use rustc_hash::FxHashMap as HashMap;
2
3use derive_generic_visitor::*;
4
5use crate::transform::ctx::UllbcPass;
6use crate::{transform::TransformCtx, ullbc_ast::*};
7
8#[derive(Visitor)]
9struct NormalizeFnPtr<'a> {
10    ctx: &'a TransformCtx,
11    /// Types are hash-consed and bodies mention the same types many times; remember what each
12    /// visited type was rewritten to instead of exploring it again.
13    visited_tys: HashMap<Ty, Ty>,
14}
15
16impl VisitAstMut for NormalizeFnPtr<'_> {
17    fn enter_fn_ptr(&mut self, fn_ptr: &mut FnPtr) {
18        if let Some(new_fn_ptr) = normalize_method_call(self.ctx, fn_ptr) {
19            *fn_ptr = new_fn_ptr;
20        }
21    }
22    fn visit_ty(&mut self, ty: &mut Ty) -> ControlFlow<Self::Break> {
23        if let Some(new_ty) = self.visited_tys.get(ty) {
24            *ty = new_ty.clone();
25            return Continue(());
26        }
27        let old_ty = ty.clone();
28        self.visit_inner(ty)?;
29        self.visited_tys.insert(old_ty, ty.clone());
30        Continue(())
31    }
32}
33
34fn normalize_method_call(ctx: &TransformCtx, fn_ptr: &FnPtr) -> Option<FnPtr> {
35    let impl_trait_ref;
36    let method_generics;
37    let (impl_ref, method_id, method_generics) = match fn_ptr.kind.as_ref() {
38        // Transform `Trait::default_method<X>[impl_trait_for_X]` to a direct method call.
39        FnPtrKind::Fun(fun_id) => {
40            let fun_decl = ctx.translated.fun_decls.get(*fun_id)?;
41            let FunSource::TraitDefault {
42                trait_ref,
43                item_id: method_id,
44            } = &fun_decl.src
45            else {
46                return None;
47            };
48            // If the first trait proof (for the self clause) is a known impl.
49            impl_trait_ref = fn_ptr
50                .generics
51                .trait_refs
52                .get(TraitClauseId::ZERO)?
53                .clone()
54                .no_bound_vars();
55            let impl_ref = impl_trait_ref.kind.as_trait_impl()?;
56            method_generics = {
57                let generics = &fn_ptr.generics;
58                let trait_generics = trait_ref.generics.as_ref();
59                GenericArgs {
60                    regions: generics
61                        .regions
62                        .clone()
63                        .split_off(trait_generics.regions.len()),
64                    types: generics.types.clone().split_off(trait_generics.types.len()),
65                    const_generics: generics
66                        .const_generics
67                        .clone()
68                        .split_off(trait_generics.const_generics.len()),
69                    // The `+ 1` is for the self clause.
70                    trait_refs: generics
71                        .trait_refs
72                        .clone()
73                        .split_off(trait_generics.trait_refs.len() + 1),
74                }
75            };
76            (impl_ref, method_id, &method_generics)
77        }
78        // Transform `impl_trait_for_X::method` to a direct method call.
79        FnPtrKind::Trait(trait_ref, method_id) => {
80            let TraitRefKind::TraitImpl(impl_ref) = &trait_ref.kind else {
81                return None;
82            };
83            (impl_ref, method_id, fn_ptr.generics.as_ref())
84        }
85    };
86
87    let trait_impl = &ctx.translated.trait_impls.get(impl_ref.id)?;
88    // Find the function declaration corresponding to this impl.
89    let bound_fn = trait_impl.methods.get(*method_id)?;
90    if !method_generics.matches(&bound_fn.params) {
91        return None;
92    }
93
94    // Make the two levels of binding explicit: outer binder for the impl block, inner binder for
95    // the method.
96    let fn_ref: Binder<Binder<FunDeclRef>> = Binder::new(
97        BinderKind::Other,
98        trait_impl.generics.clone(),
99        bound_fn.clone(),
100    );
101    // Substitute the appropriate generics into the function call.
102    let fn_ref = fn_ref.apply(&impl_ref.generics).apply(method_generics);
103    Some(FnPtr::new(FnPtrKind::Fun(fn_ref.id), fn_ref.generics))
104}
105
106pub struct Transform;
107impl UllbcPass for Transform {
108    fn transform_item(&self, ctx: &mut TransformCtx, mut item: ItemRefMut<'_>) {
109        let _ = item.drive_mut(&mut NormalizeFnPtr {
110            ctx,
111            visited_tys: Default::default(),
112        });
113    }
114}
charon_lib/transform/normalize/skip_trait_refs_when_known.rs

charon_lib/transform/normalize/
skip_trait_refs_when_known.rs
