charon_lib/transform/
monomorphize.rs

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

charon_lib/transform/
monomorphize.rs
