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

charon_lib/transform/
monomorphize.rs
