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

charon_lib/transform/normalize/
monomorphize.rs
