charon_lib/transform/
monomorphize.rs

1//! # Micro-pass: monomorphize all functions and types; at the end of this pass, all functions and types are monomorphic.
2use derive_generic_visitor::*;
3use indexmap::IndexMap;
4use std::collections::HashSet;
5
6use crate::ast::*;
7use crate::transform::TransformCtx;
8use std::fmt::Debug;
9
10use super::ctx::TransformPass;
11
12enum OptionHint<T, H> {
13    Some(T),
14    None,
15    Hint(H),
16}
17
18impl<T, H> OptionHint<T, H> {
19    fn is_some(&self) -> bool {
20        match self {
21            OptionHint::Some(_) => true,
22            OptionHint::None => false,
23            OptionHint::Hint(_) => false,
24        }
25    }
26
27    fn hint_or<'a>(&'a self, hint: &'a H) -> &'a H {
28        match self {
29            OptionHint::Some(_) => hint,
30            OptionHint::None => hint,
31            OptionHint::Hint(h) => h,
32        }
33    }
34}
35
36#[derive(Default)]
37struct PassData {
38    // Map of (poly item, generic args) -> mono item
39    // None indicates the item hasn't been monomorphized yet
40    items: IndexMap<(AnyTransId, GenericArgs), OptionHint<AnyTransId, (AnyTransId, BoxedArgs)>>,
41    worklist: Vec<AnyTransId>,
42    visited: HashSet<AnyTransId>,
43}
44
45impl PassData {
46    fn new() -> Self {
47        Self::default()
48    }
49}
50
51impl TranslatedCrate {
52    // FIXME(Nadrieril): implement type&tref normalization and use that instead
53    fn find_trait_impl_and_gargs(
54        self: &Self,
55        tref: &TraitRef,
56    ) -> Option<(&TraitImpl, GenericArgs)> {
57        match &tref.kind {
58            TraitRefKind::TraitImpl(impl_ref) => {
59                let trait_impl = self.trait_impls.get(impl_ref.id)?;
60                Some((trait_impl, impl_ref.generics.as_ref().clone()))
61            }
62            TraitRefKind::ParentClause(p, clause) => {
63                let (trait_impl, _) = self.find_trait_impl_and_gargs(p)?;
64                let t_ref = trait_impl.parent_trait_refs.get(*clause)?;
65                self.find_trait_impl_and_gargs(t_ref)
66            }
67            _ => None,
68        }
69    }
70}
71
72#[derive(Visitor)]
73struct UsageVisitor<'a> {
74    data: &'a mut PassData,
75    krate: &'a TranslatedCrate,
76}
77impl UsageVisitor<'_> {
78    fn found_use(
79        &mut self,
80        id: &AnyTransId,
81        gargs: &GenericArgs,
82        default: OptionHint<AnyTransId, (AnyTransId, BoxedArgs)>,
83    ) {
84        trace!("Mono: Found use: {:?} / {:?}", id, gargs);
85        self.data
86            .items
87            .entry((*id, gargs.clone()))
88            .or_insert(default);
89    }
90    fn found_use_ty(&mut self, tref: &TypeDeclRef) {
91        match tref.id {
92            TypeId::Adt(id) => {
93                self.found_use(&AnyTransId::Type(id), &tref.generics, OptionHint::None)
94            }
95            _ => {}
96        }
97    }
98    fn found_use_fn(&mut self, id: &FunDeclId, gargs: &GenericArgs) {
99        self.found_use(&AnyTransId::Fun(*id), gargs, OptionHint::None);
100    }
101    fn found_use_global_decl_ref(&mut self, id: &GlobalDeclId, gargs: &GenericArgs) {
102        self.found_use(&AnyTransId::Global(*id), gargs, OptionHint::None);
103    }
104    fn found_use_fn_hinted(
105        &mut self,
106        id: &FunDeclId,
107        gargs: &GenericArgs,
108        (h_id, h_args): (FunDeclId, BoxedArgs),
109    ) {
110        self.found_use(
111            &AnyTransId::Fun(*id),
112            gargs,
113            OptionHint::Hint((AnyTransId::Fun(h_id), h_args)),
114        );
115    }
116}
117impl VisitAst for UsageVisitor<'_> {
118    // we need to skip ItemMeta, as we don't want to collect the types in PathElem::Impl
119    fn visit_item_meta(&mut self, _: &ItemMeta) -> ControlFlow<Infallible> {
120        Continue(())
121    }
122
123    fn enter_aggregate_kind(&mut self, kind: &AggregateKind) {
124        match kind {
125            AggregateKind::Adt(tref, _, _) => self.found_use_ty(tref),
126            _ => {}
127        }
128    }
129
130    fn visit_ty_kind(&mut self, kind: &TyKind) -> ControlFlow<Infallible> {
131        match kind {
132            TyKind::Adt(tref) => {
133                self.found_use_ty(tref);
134            }
135            TyKind::FnDef(binder) => {
136                // we don't want to visit inside the binder, as it will have regions that
137                // haven't been erased; instead we visit the erased version, and skip
138                // the default "visit_inner" behaviour
139                let _ = self.visit(&binder.clone().erase());
140                return Continue(());
141            }
142            _ => {}
143        };
144        self.visit_inner(kind)
145    }
146
147    fn enter_fn_ptr(&mut self, fn_ptr: &FnPtr) {
148        match fn_ptr.func.as_ref() {
149            FunIdOrTraitMethodRef::Fun(FunId::Regular(id)) => {
150                self.found_use_fn(&id, &fn_ptr.generics)
151            }
152            FunIdOrTraitMethodRef::Trait(t_ref, name, id) => {
153                let Some((trait_impl, impl_gargs)) = self.krate.find_trait_impl_and_gargs(t_ref)
154                else {
155                    return;
156                };
157                let (_, bound_fn) = trait_impl.methods().find(|(n, _)| n == name).unwrap();
158                let fn_ref: Binder<Binder<FunDeclRef>> = Binder::new(
159                    BinderKind::Other,
160                    trait_impl.generics.clone(),
161                    bound_fn.clone(),
162                );
163                // This is the actual function we need to call!
164                // Whereas id is the trait method reference(?)
165                let fn_ref = fn_ref.apply(&impl_gargs).apply(&fn_ptr.generics);
166                let gargs_key = fn_ptr
167                    .generics
168                    .clone()
169                    .concat(&t_ref.trait_decl_ref.skip_binder.generics);
170                self.found_use_fn_hinted(&id, &gargs_key, (fn_ref.id, fn_ref.generics))
171            }
172            // These can't be monomorphized, since they're builtins
173            FunIdOrTraitMethodRef::Fun(FunId::Builtin(..)) => {}
174        }
175    }
176
177    fn enter_global_decl_ref(&mut self, glob: &GlobalDeclRef) {
178        self.found_use_global_decl_ref(&glob.id, &glob.generics);
179    }
180}
181
182// Akin to UsageVisitor, but substitutes all uses of generics with the monomorphized versions
183// This is a two-step process, because we can't mutate the translation context with new definitions
184// while also mutating the existing definitions.
185#[derive(Visitor)]
186struct SubstVisitor<'a> {
187    data: &'a PassData,
188}
189impl SubstVisitor<'_> {
190    fn subst_use<T, F>(&mut self, id: &mut T, gargs: &mut GenericArgs, of: F)
191    where
192        T: Into<AnyTransId> + Debug + Copy,
193        F: Fn(&AnyTransId) -> Option<&T>,
194    {
195        trace!("Mono: Subst use: {:?} / {:?}", id, gargs);
196        // Erase regions.
197        gargs.regions.iter_mut().for_each(|r| *r = Region::Erased);
198        let key = ((*id).into(), gargs.clone());
199        let subst = self.data.items.get(&key);
200        if let Some(OptionHint::Some(any_id)) = subst
201            && let Some(subst_id) = of(any_id)
202        {
203            *id = *subst_id;
204            *gargs = GenericArgs::empty();
205        } else {
206            warn!("Substitution missing for {:?} / {:?}", id, gargs);
207        }
208    }
209    fn subst_use_ty(&mut self, tref: &mut TypeDeclRef) {
210        match &mut tref.id {
211            TypeId::Adt(id) => {
212                self.subst_use(id, &mut tref.generics, AnyTransId::as_type);
213            }
214            _ => {}
215        }
216    }
217    fn subst_use_fun(&mut self, id: &mut FunDeclId, gargs: &mut GenericArgs) {
218        self.subst_use(id, gargs, AnyTransId::as_fun);
219    }
220    fn subst_use_glob(&mut self, id: &mut GlobalDeclId, gargs: &mut GenericArgs) {
221        self.subst_use(id, gargs, AnyTransId::as_global);
222    }
223}
224
225impl VisitAstMut for SubstVisitor<'_> {
226    fn enter_aggregate_kind(&mut self, kind: &mut AggregateKind) {
227        match kind {
228            AggregateKind::Adt(tref, _, _) => self.subst_use_ty(tref),
229            _ => {}
230        }
231    }
232
233    fn enter_ty_kind(&mut self, kind: &mut TyKind) {
234        match kind {
235            TyKind::Adt(tref) => self.subst_use_ty(tref),
236            TyKind::FnDef(binder) => {
237                // erase the FnPtr binder, as we'll monomorphise its content
238                if let FnPtr {
239                    func: box FunIdOrTraitMethodRef::Fun(FunId::Regular(id)),
240                    generics,
241                } = binder.clone().erase()
242                {
243                    *binder = RegionBinder::empty(FnPtr {
244                        func: Box::new(FunIdOrTraitMethodRef::Fun(FunId::Regular(id))),
245                        generics,
246                    });
247                }
248            }
249            _ => {}
250        }
251    }
252
253    fn enter_fn_ptr(&mut self, fn_ptr: &mut FnPtr) {
254        match fn_ptr.func.as_mut() {
255            FunIdOrTraitMethodRef::Fun(FunId::Regular(fun_id)) => {
256                self.subst_use_fun(fun_id, &mut fn_ptr.generics)
257            }
258            FunIdOrTraitMethodRef::Trait(t_ref, _, fun_id) => {
259                let mut gargs_key = fn_ptr
260                    .generics
261                    .clone()
262                    .concat(&t_ref.trait_decl_ref.skip_binder.generics);
263                self.subst_use_fun(fun_id, &mut gargs_key);
264                fn_ptr.generics = Box::new(gargs_key);
265            }
266            // These can't be monomorphized, since they're builtins
267            FunIdOrTraitMethodRef::Fun(FunId::Builtin(..)) => {}
268        }
269    }
270
271    fn exit_place(&mut self, place: &mut Place) {
272        match &mut place.kind {
273            // FIXME(Nadrieril): remove this id, replace with a helper fn
274            PlaceKind::Projection(inner, ProjectionElem::Field(FieldProjKind::Adt(id, _), _)) => {
275                // Trick, we don't know the generics but the projected place does, so
276                // we substitute it there, then update our current id.
277                let tref = inner.ty.as_adt().unwrap();
278                *id = *tref.id.as_adt().unwrap()
279            }
280            _ => {}
281        }
282    }
283
284    fn enter_global_decl_ref(&mut self, glob: &mut GlobalDeclRef) {
285        self.subst_use_glob(&mut glob.id, &mut glob.generics);
286    }
287}
288
289#[derive(Visitor)]
290#[allow(dead_code)]
291struct MissingIndexChecker<'a> {
292    krate: &'a TranslatedCrate,
293    current_item: Option<AnyTransItem<'a>>,
294}
295impl VisitAst for MissingIndexChecker<'_> {
296    fn enter_fun_decl_id(&mut self, id: &FunDeclId) {
297        if self.krate.fun_decls.get(*id).is_none() {
298            panic!(
299                "Missing function declaration for id: {:?}, in {:?}",
300                id, self.current_item
301            );
302        }
303    }
304
305    fn enter_trait_impl_id(&mut self, id: &TraitImplId) {
306        if self.krate.trait_impls.get(*id).is_none() {
307            panic!(
308                "Missing trait implementation for id: {:?}, in {:?}",
309                id, self.current_item
310            );
311        }
312    }
313
314    fn enter_trait_decl_id(&mut self, id: &TraitDeclId) {
315        if self.krate.trait_decls.get(*id).is_none() {
316            panic!(
317                "Missing trait declaration for id: {:?}, in {:?}",
318                id, self.current_item
319            );
320        }
321    }
322
323    fn enter_type_decl_id(&mut self, id: &TypeDeclId) {
324        if self.krate.type_decls.get(*id).is_none() {
325            panic!(
326                "Missing type declaration for id: {:?}, in {:?}",
327                id, self.current_item
328            );
329        }
330    }
331}
332
333fn find_uses(data: &mut PassData, krate: &TranslatedCrate, item: &AnyTransItem) {
334    let mut visitor = UsageVisitor { data, krate };
335    let _ = item.drive(&mut visitor);
336}
337
338fn subst_uses<T: AstVisitable + Debug>(data: &PassData, item: &mut T) {
339    let mut visitor = SubstVisitor { data };
340    let _ = item.drive_mut(&mut visitor);
341}
342
343// fn check_missing_indices(krate: &TranslatedCrate) {
344//     let mut visitor = MissingIndexChecker {
345//         krate,
346//         current_item: None,
347//     };
348//     for item in krate.all_items() {
349//         visitor.current_item = Some(item);
350//         item.drive(&mut visitor);
351//     }
352// }
353
354// fn path_for_generics(gargs: &GenericArgs) -> PathElem {
355//     PathElem::Ident(gargs.to_string(), Disambiguator::ZERO)
356// }
357
358pub struct Transform;
359impl TransformPass for Transform {
360    fn transform_ctx(&self, ctx: &mut TransformCtx) {
361        // Check the option which instructs to ignore this pass
362        if !ctx.options.monomorphize_as_pass {
363            return;
364        }
365
366        // From https://doc.rust-lang.org/nightly/nightly-rustc/rustc_monomorphize/collector/index.html#general-algorithm
367        //
368        // The purpose of the algorithm implemented in this module is to build the mono item
369        // graph for the current crate. It runs in two phases:
370        // 1. Discover the roots of the graph by traversing the HIR of the crate.
371        // 2. Starting from the roots, find uses by inspecting the MIR representation of the
372        //    item corresponding to a given node, until no more new nodes are found.
373        //
374        // The roots of the mono item graph correspond to the public non-generic syntactic
375        // items in the source code. We find them by walking the HIR of the crate, and whenever
376        // we hit upon a public function, method, or static item, we create a mono item
377        // consisting of the items DefId and, since we only consider non-generic items, an
378        // empty type-parameters set.
379        //
380        // Given a mono item node, we can discover uses by inspecting its MIR. We walk the MIR
381        // to find other mono items used by each mono item. Since the mono item we are
382        // currently at is always monomorphic, we also know the concrete type arguments of its
383        // used mono items. The specific forms a use can take in MIR are quite diverse: it
384        // includes calling functions/methods, taking a reference to a function/method, drop
385        // glue, and unsizing casts.
386
387        // In our version of the algorithm, we do the following:
388        // 1. Find all the roots, adding them to the worklist.
389        // 2. For each item in the worklist:
390        //    a. Find all the items it uses, adding them to the worklist and the generic
391        //      arguments to the item.
392        //    b. Mark the item as visited
393
394        // Final list of monomorphized items: { (poly item, generic args) -> mono item }
395        let mut data = PassData::new();
396
397        let empty_gargs = GenericArgs::empty();
398
399        // Find the roots of the mono item graph
400        for (id, item) in ctx.translated.all_items_with_ids() {
401            match item {
402                AnyTransItem::Fun(f) if f.signature.generics.is_empty() => {
403                    data.items
404                        .insert((id, empty_gargs.clone()), OptionHint::Some(id));
405                    data.worklist.push(id);
406                }
407                _ => {}
408            }
409        }
410
411        // Iterate over worklist -- these items are always monomorphic!
412        while let Some(id) = data.worklist.pop() {
413            if data.visited.contains(&id) {
414                continue;
415            }
416            data.visited.insert(id);
417
418            // 1. Find new uses
419            let Some(item) = ctx.translated.get_item(id) else {
420                trace!("Couldn't find item {:} in translated items?", id);
421                continue;
422            };
423            find_uses(&mut data, &ctx.translated, &item);
424
425            // 2. Iterate through all newly discovered uses
426            for ((id, gargs), mono) in data.items.iter_mut() {
427                if mono.is_some() {
428                    continue;
429                }
430
431                // a. Monomorphize the items if they're polymorphic, add them to the worklist
432                let new_mono = if gargs.is_empty() {
433                    *id
434                } else {
435                    match id {
436                        AnyTransId::Fun(_) => {
437                            let key_pair = (id.clone(), Box::new(gargs.clone()));
438                            let (AnyTransId::Fun(fun_id), gargs) = mono.hint_or(&key_pair) else {
439                                panic!("Unexpected ID type in hint_or");
440                            };
441                            let fun = ctx.translated.fun_decls.get(*fun_id).unwrap();
442                            let mut fun_sub = fun.clone().substitute(gargs);
443                            fun_sub.signature.generics = GenericParams::empty();
444                            fun_sub
445                                .item_meta
446                                .name
447                                .name
448                                .push(PathElem::Monomorphized(gargs.clone()));
449
450                            let fun_id_sub = ctx.translated.fun_decls.push_with(|id| {
451                                fun_sub.def_id = id;
452                                fun_sub
453                            });
454
455                            AnyTransId::Fun(fun_id_sub)
456                        }
457                        AnyTransId::Type(typ_id) => {
458                            let typ = ctx.translated.type_decls.get(*typ_id).unwrap();
459                            let mut typ_sub = typ.clone().substitute(gargs);
460                            typ_sub.generics = GenericParams::empty();
461                            typ_sub
462                                .item_meta
463                                .name
464                                .name
465                                .push(PathElem::Monomorphized(gargs.clone().into()));
466
467                            let typ_id_sub = ctx.translated.type_decls.push_with(|id| {
468                                typ_sub.def_id = id;
469                                typ_sub
470                            });
471
472                            AnyTransId::Type(typ_id_sub)
473                        }
474                        AnyTransId::Global(g_id) => {
475                            let Some(glob) = ctx.translated.global_decls.get(*g_id) else {
476                                // Something odd happened -- we ignore and move on
477                                *mono = OptionHint::Some(*id);
478                                warn!("Found a global that has no associated declaration");
479                                continue;
480                            };
481                            let mut glob_sub = glob.clone().substitute(gargs);
482                            glob_sub.generics = GenericParams::empty();
483                            glob_sub
484                                .item_meta
485                                .name
486                                .name
487                                .push(PathElem::Monomorphized(gargs.clone().into()));
488
489                            let init = ctx.translated.fun_decls.get(glob.init).unwrap();
490                            let mut init_sub = init.clone().substitute(gargs);
491                            init_sub.signature.generics = GenericParams::empty();
492                            init_sub
493                                .item_meta
494                                .name
495                                .name
496                                .push(PathElem::Monomorphized(gargs.clone().into()));
497
498                            let init_id_sub = ctx.translated.fun_decls.push_with(|id| {
499                                init_sub.def_id = id;
500                                glob_sub.init = id;
501                                init_sub
502                            });
503
504                            let g_id_sub = ctx.translated.global_decls.push_with(|id| {
505                                glob_sub.def_id = id;
506                                glob_sub
507                            });
508
509                            data.worklist.push(AnyTransId::Fun(init_id_sub));
510
511                            AnyTransId::Global(g_id_sub)
512                        }
513                        _ => todo!("Unhandled monomorphization target ID {:?}", id),
514                    }
515                };
516                trace!(
517                    "Mono: Monomorphized {:?} with {:?} to {:?}",
518                    id, gargs, new_mono
519                );
520                if id != &new_mono {
521                    trace!(" - From {:?}", ctx.translated.get_item(id.clone()));
522                    trace!(" - To {:?}", ctx.translated.get_item(new_mono.clone()));
523                }
524                *mono = OptionHint::Some(new_mono);
525                data.worklist.push(new_mono);
526
527                let Some(item) = ctx.translated.get_item(new_mono) else {
528                    trace!("Missing monomorphised item {new_mono:?}");
529                    continue;
530                };
531                ctx.translated
532                    .item_names
533                    .insert(new_mono, item.item_meta().name.clone());
534            }
535
536            // 3. Substitute all generics with the monomorphized versions
537            let Some(item) = ctx.translated.get_item_mut(id) else {
538                panic!("Couldn't find item {:} in translated items.", id)
539            };
540            match item {
541                AnyTransItemMut::Fun(f) => subst_uses(&data, f),
542                AnyTransItemMut::Type(t) => subst_uses(&data, t),
543                AnyTransItemMut::TraitImpl(t) => subst_uses(&data, t),
544                AnyTransItemMut::Global(g) => subst_uses(&data, g),
545                AnyTransItemMut::TraitDecl(t) => subst_uses(&data, t),
546            };
547        }
548
549        // Now, remove all polymorphic items from the translation context, as all their
550        // uses have been monomorphized and substituted
551        ctx.translated
552            .fun_decls
553            .retain(|f| data.visited.contains(&AnyTransId::Fun(f.def_id)));
554        ctx.translated
555            .type_decls
556            .retain(|t| data.visited.contains(&AnyTransId::Type(t.def_id)));
557        ctx.translated
558            .global_decls
559            .retain(|g| data.visited.contains(&AnyTransId::Global(g.def_id)));
560        // ctx.translated.trait_impls.retain(|t| t.generics.is_empty());
561
562        // TODO: Currently we don't update all TraitImpls/TraitDecls with the monomorphized versions
563        //       and removing the polymorphic ones, so this fails.
564        // Finally, ensure we didn't leave any IDs un-replaced
565        // check_missing_indices(&ctx.translated);
566    }
567}
charon_lib/transform/monomorphize.rs

charon_lib/transform/
monomorphize.rs
