charon_lib/transform/normalize/
partial_monomorphization.rs

1//! This module implements partial monomorphization, which allows specializing generic items on
2//! some specific instanciation patterns. This is used by Aeneas to avoid nested mutable borrows:
3//! we transform `Iter<'a, &'b mut T>` to `{Iter::<_, &mut U>}<'a, 'b, T>`, where
4//! ```ignore
5//! struct {Iter::<'a, &'b mut U>}<'a, 'b, U> {
6//!   // the field of `Iter` but instantiated with `T -> &'b mut U`.
7//! }
8//! ```
9//!
10//! Note: We may need to partial-mono the same item multiple times: `Foo::<&mut A, B>`, `Foo::<A,
11//! &mut B>`. Note also that partial-mono is infectious: `Foo<Bar<&mut A>>` generates `Bar::<&mut
12//! A>` then `Foo::<Bar::<&mut A>>``.
13use std::collections::{HashMap, HashSet, VecDeque};
14use std::fmt::Display;
15use std::mem;
16
17use derive_generic_visitor::Visitor;
18use index_vec::Idx;
19
20use crate::ast::visitor::{VisitWithBinderDepth, VisitorWithBinderDepth};
21use crate::formatter::IntoFormatter;
22use crate::options::MonomorphizeMut;
23use crate::pretty::FmtWithCtx;
24use crate::register_error;
25use crate::transform::ctx::TransformPass;
26use crate::{transform::TransformCtx, ullbc_ast::*};
27
28type MutabilityShape = Binder<GenericArgs>;
29
30/// See the docs of `MutabilityShapeBuilder::compute_shape`.
31#[derive(Visitor)]
32struct MutabilityShapeBuilder<'pm, 'ctx> {
33    pm: &'pm PartialMonomorphizer<'ctx>,
34    /// The parameters that will constitute the final binder.
35    params: GenericParams,
36    /// The arguments to pass to the final binder to recover the input arguments.
37    extracted: GenericArgs,
38    /// Current depth under which we're visiting.
39    binder_depth: DeBruijnId,
40}
41
42impl<'pm, 'ctx> MutabilityShapeBuilder<'pm, 'ctx> {
43    /// Compute the mutability "shape" of a set of generic arguments by factoring out the minimal
44    /// amount of information that still allows reconstructing the original arguments while keeping the
45    /// "shape arguments" free of mutable borrows.
46    ///
47    /// For example, for input:
48    ///   <u32, &'a mut &'b A, Option::<&'a mut bool>>
49    /// we want to build:
50    ///   binder<'a, 'b, A, B, C> <A, &'a mut B, Option::<&'b mut C>>
51    /// from which we can recover the original arguments by instantiating it with:
52    ///   <'a, 'a, u32, &'b A, bool>
53    ///
54    /// Formally, given `let (shape, shape_args) = get_mutability_shape(args);`, we have the following:
55    /// - `shape.substitute(shape_args) == args`;
56    /// - `shape_args` contains no infected types;
57    /// - `shape` is as shallow as possible (i.e. takes just enough to get all the infected types
58    ///   and not more).
59    ///
60    /// Note: the input arguments are assumed to have been already partially monomorphized, in the
61    /// sense that we won't recurse inside ADT args because we assume any ADT applied to infected
62    /// args to have been replaced with a fresh infected ADT.
63    fn compute_shape(
64        pm: &'pm PartialMonomorphizer<'ctx>,
65        target_params: &GenericParams,
66        args: &GenericArgs,
67    ) -> (MutabilityShape, GenericArgs) {
68        // We start with the implicit parameters from the original item. We'll need to substitute
69        // them once we've figured out the mapping of explicit parameters, but we'll also be adding
70        // new trait clauses potentially so we can't leave the vector empty (the ids would be
71        // wrong).
72        let mut shape_contents = args.clone();
73        let mut builder = Self {
74            pm,
75            params: GenericParams {
76                regions: IndexVec::new(),
77                types: IndexVec::new(),
78                const_generics: IndexVec::new(),
79                ..target_params.clone()
80            },
81            extracted: GenericArgs {
82                regions: IndexVec::new(),
83                types: IndexVec::new(),
84                const_generics: IndexVec::new(),
85                trait_refs: mem::take(&mut shape_contents.trait_refs),
86            },
87            binder_depth: DeBruijnId::zero(),
88        };
89
90        // Traverse the generics and replace any non-infected type, region or const generic with a
91        // fresh variable.
92        let _ = VisitWithBinderDepth::new(&mut builder).visit(&mut shape_contents);
93
94        let shape_params = {
95            let mut shape_params = builder.params;
96            // Now the explicit params in `shape_params` are correct, and the implicit params are a mix
97            // of the old params and new trait clauses. The old params may refer to the old explicit
98            // params which is wrong and must be fixed up.
99            shape_params.trait_clauses = shape_params.trait_clauses.map_indexed(|i, x| {
100                if i.index() < target_params.trait_clauses.len() {
101                    x.substitute_explicits(&shape_contents)
102                } else {
103                    x
104                }
105            });
106            shape_params.trait_type_constraints =
107                shape_params.trait_type_constraints.map_indexed(|i, x| {
108                    if i.index() < target_params.trait_type_constraints.len() {
109                        x.substitute_explicits(&shape_contents)
110                    } else {
111                        x
112                    }
113                });
114            shape_params.regions_outlive = shape_params
115                .regions_outlive
116                .into_iter()
117                .enumerate()
118                .map(|(i, x)| {
119                    if i < target_params.regions_outlive.len() {
120                        x.substitute_explicits(&shape_contents)
121                    } else {
122                        x
123                    }
124                })
125                .collect();
126            shape_params.types_outlive = shape_params
127                .types_outlive
128                .into_iter()
129                .enumerate()
130                .map(|(i, x)| {
131                    if i < target_params.types_outlive.len() {
132                        x.substitute_explicits(&shape_contents)
133                    } else {
134                        x
135                    }
136                })
137                .collect();
138            shape_params
139        };
140
141        // The first half of the trait params correspond to the original item clauses so we can
142        // pass them unmodified.
143        shape_contents.trait_refs = shape_params.identity_args().trait_refs;
144        shape_contents
145            .trait_refs
146            .truncate(target_params.trait_clauses.len());
147
148        let shape_args = builder.extracted;
149        let shape = Binder::new(BinderKind::Other, shape_params, shape_contents);
150        (shape, shape_args)
151    }
152
153    /// Replace this value with a fresh variable, and record that we did so.
154    fn replace_with_fresh_var<Id, Param, Arg>(
155        &mut self,
156        val: &mut Arg,
157        mk_param: impl FnOnce(Id) -> Param,
158        mk_value: impl FnOnce(DeBruijnVar<Id>) -> Arg,
159    ) where
160        Id: Idx + Display,
161        Arg: TyVisitable + Clone,
162        GenericParams: HasIdxVecOf<Id, Output = Param>,
163        GenericArgs: HasIdxVecOf<Id, Output = Arg>,
164    {
165        let Some(shifted_val) = val.clone().move_from_under_binders(self.binder_depth) else {
166            // Give up on this value.
167            return;
168        };
169        // Record the mapping in the output `GenericArgs`.
170        self.extracted.get_idx_vec_mut().push(shifted_val);
171        // Put a fresh param in place of `val`.
172        let id = self.params.get_idx_vec_mut().push_with(mk_param);
173        *val = mk_value(DeBruijnVar::bound(self.binder_depth, id));
174    }
175}
176
177impl<'pm, 'ctx> VisitorWithBinderDepth for MutabilityShapeBuilder<'pm, 'ctx> {
178    fn binder_depth_mut(&mut self) -> &mut DeBruijnId {
179        &mut self.binder_depth
180    }
181}
182
183impl<'pm, 'ctx> VisitAstMut for MutabilityShapeBuilder<'pm, 'ctx> {
184    fn visit<T: AstVisitable>(&mut self, x: &mut T) -> ControlFlow<Self::Break> {
185        VisitWithBinderDepth::new(self).visit(x)
186    }
187
188    fn enter_ty(&mut self, ty: &mut Ty) {
189        if !self.pm.is_infected(ty) {
190            self.replace_with_fresh_var(
191                ty,
192                |id| TypeParam::new(id, format!("T{id}"), Variance::Unknown),
193                |v| v.into(),
194            );
195        }
196    }
197    fn exit_ty_kind(&mut self, kind: &mut TyKind) {
198        if let TyKind::Adt(TypeDeclRef {
199            id: TypeId::Adt(id),
200            generics,
201        }) = kind
202        {
203            // Since the type was not replaced with a type var, it's an infected type. We've
204            // traversed it so we have its final explicit arguments. Now we need to satisfy its
205            // predicates. For that we add all its predicates to the new item, and pass those new
206            // trait clauses to it.
207            let Some(target_params) = self.pm.generic_params.get(&(*id).into()) else {
208                return;
209            };
210            let Some(shifted_generics) =
211                generics.clone().move_from_under_binders(self.binder_depth)
212            else {
213                // Give up on this value.
214                return;
215            };
216
217            // Add the target predicates (properly substituted) to the new item params.
218            let num_clauses_before_merge = self.params.trait_clauses.len();
219            self.params.merge_predicates_from(
220                target_params
221                    .clone()
222                    .substitute_explicits(&shifted_generics),
223            );
224
225            // Record the trait arguments in the output `GenericArgs`.
226            self.extracted
227                .trait_refs
228                .extend(shifted_generics.trait_refs);
229
230            // Replace each trait ref with a clause var.
231            for (target_clause_id, tref) in generics.trait_refs.iter_mut_enumerated() {
232                let clause_id = target_clause_id + num_clauses_before_merge;
233                *tref =
234                    self.params.trait_clauses[clause_id].identity_tref_at_depth(self.binder_depth);
235            }
236        }
237    }
238    fn enter_region(&mut self, r: &mut Region) {
239        self.replace_with_fresh_var(
240            r,
241            |id| RegionParam::new(id, None, Variance::Unknown),
242            |v| v.into(),
243        );
244    }
245    // TODO: we're missing type info for this
246    // fn enter_const_generic(&mut self, cg: &mut ConstGeneric) {
247    //     self.replace_with_fresh_var(cg, |id| {
248    //         ConstGenericParam::new(id, format!("N{id}"), cg.ty().clone())
249    //     });
250    // }
251    fn visit_trait_ref(&mut self, _tref: &mut TraitRef) -> ControlFlow<Self::Break> {
252        // We don't touch trait refs or we'd risk adding duplicated extra params. Instead, we fix
253        // them up in `exit_ty_kind` and `compute_shape`.
254        ControlFlow::Continue(())
255    }
256
257    fn visit_constant_expr(
258        &mut self,
259        _: &mut ConstantExpr,
260    ) -> ::std::ops::ControlFlow<Self::Break> {
261        ControlFlow::Continue(())
262    }
263}
264
265#[derive(Visitor)]
266struct PartialMonomorphizer<'a> {
267    ctx: &'a mut TransformCtx,
268    /// Tracks the closest span to emit useful errors.
269    span: Span,
270    /// Whether we partially-monomorphize type declarations.
271    specialize_adts: bool,
272    /// Types that contain mutable references.
273    infected_types: HashSet<TypeDeclId>,
274    /// Map of generic params for each item. We can't use `ctx.translated` because while iterating
275    /// over items the current item isn't available anymore, which would break recursive types.
276    /// This also makes it possible to record the generics of our to-be-added items without adding
277    /// them.
278    generic_params: HashMap<ItemId, GenericParams>,
279    /// Map of partial monomorphizations. The source item applied with the generic params gives the
280    /// target item. The resulting partially-monomorphized item will have the binder params as
281    /// generic params.
282    partial_mono_shapes: SeqHashMap<(ItemId, MutabilityShape), ItemId>,
283    /// Reverse of `partial_mono_shapes`.
284    reverse_shape_map: HashMap<ItemId, (ItemId, MutabilityShape)>,
285    /// Items that need to be processed.
286    to_process: VecDeque<ItemId>,
287}
288
289impl<'a> PartialMonomorphizer<'a> {
290    pub fn new(ctx: &'a mut TransformCtx, specialize_adts: bool) -> Self {
291        // Compute the types that contain `&mut` (even indirectly). We actually can ignore
292        // `&'static mut`, so we simply rely on our "lifetime mutability" computation.
293        let infected_types: HashSet<_> = ctx
294            .translated
295            .type_decls
296            .iter()
297            .filter(|tdecl| {
298                tdecl
299                    .generics
300                    .regions
301                    .iter()
302                    .any(|r| r.mutability.is_mutable())
303            })
304            .map(|tdecl| tdecl.def_id)
305            .collect();
306
307        // Record the generic params of all items.
308        let generic_params: HashMap<ItemId, GenericParams> = ctx
309            .translated
310            .all_items()
311            .map(|item| (item.id(), item.generic_params().clone()))
312            .collect();
313
314        // Enqueue all items to be processed.
315        let to_process = ctx.translated.all_ids().collect();
316        PartialMonomorphizer {
317            ctx,
318            span: Span::dummy(),
319            specialize_adts,
320            infected_types,
321            generic_params,
322            to_process,
323            partial_mono_shapes: SeqHashMap::default(),
324            reverse_shape_map: Default::default(),
325        }
326    }
327
328    /// Whether this type is or contains a `&mut`. This assumes that we've already visited this
329    /// type and partially monomorphized any ADT references.
330    fn is_infected(&self, ty: &Ty) -> bool {
331        match ty.kind() {
332            TyKind::Ref(_, _, RefKind::Mut) => true,
333            TyKind::Ref(_, ty, _)
334            | TyKind::RawPtr(ty, _)
335            | TyKind::Array(ty, _)
336            | TyKind::Pattern(ty, _)
337            | TyKind::Slice(ty) => self.is_infected(ty),
338            TyKind::Adt(tref) => match tref.as_adt() {
339                Some(id) => {
340                    let ty_infected = self.infected_types.contains(&id);
341                    let args_infected = if self.specialize_adts {
342                        // Since we make sure to only call the method on a processed type, any type
343                        // with infected arguments would have been replaced with a fresh instantiated
344                        // (and infected type). Hence we don't need to check the arguments here, only
345                        // the type id.
346                        false
347                    } else {
348                        tref.generics.types.iter().any(|ty| self.is_infected(ty))
349                    };
350                    ty_infected || args_infected
351                }
352                None => {
353                    // Builtin types have no declaration to specialize, so infected arguments stay
354                    // visible inside them.
355                    tref.generics.types.iter().any(|ty| self.is_infected(ty))
356                }
357            },
358            // A function pointer/item by itself doesn't carry any mutable reference, even if it
359            // uses some in its signature. Compare with closures: a closure without captures
360            // doesn't trigger partial mono regardless of its signature.
361            TyKind::FnDef(..) | TyKind::FnPtr(..) => false,
362            TyKind::DynTrait(_) => {
363                register_error!(
364                    self.ctx,
365                    self.span,
366                    "`dyn Trait` is unsupported with `--monomorphize-mut`"
367                );
368                false
369            }
370            TyKind::TypeVar(..)
371            | TyKind::Literal(..)
372            | TyKind::Never
373            | TyKind::TraitType(..)
374            | TyKind::PtrMetadata(..)
375            | TyKind::Error(_) => false,
376        }
377    }
378
379    /// Given that `generics` apply to item `id`, if any of the generics is infected we generate a
380    /// reference to a new item obtained by partially instantiating item `id`. (That new item isn't
381    /// added immediately but is added to the `to_process` queue to be created later).
382    fn process_generics(&mut self, id: ItemId, generics: &GenericArgs) -> Option<DeclRef<ItemId>> {
383        if !generics.types.iter().any(|ty| self.is_infected(ty)) {
384            return None;
385        }
386
387        // If the type is already an instantiation, transform this reference into a reference to
388        // the original type so we don't instantiate the instantiation.
389        let mut new_generics;
390        let (id, generics) = if let Some(&(base_id, ref shape)) = self.reverse_shape_map.get(&id) {
391            new_generics = shape.clone().apply(generics);
392            let _ = self.visit(&mut new_generics); // New instantiation may require cleanup.
393            (base_id, &new_generics)
394        } else {
395            (id, generics)
396        };
397
398        // Split the args between the infected part and the non-infected part.
399        let item_params = self.generic_params.get(&id)?;
400        let (shape, shape_args) =
401            MutabilityShapeBuilder::compute_shape(self, item_params, generics);
402
403        // Create a new type id.
404        let new_params = shape.params.clone();
405        let key: (ItemId, MutabilityShape) = (id, shape);
406        let new_id = *self
407            .partial_mono_shapes
408            .entry(key.clone())
409            .or_insert_with(|| {
410                let new_id = match id {
411                    ItemId::Type(_) => {
412                        let new_id = self.ctx.translated.type_decls.reserve_slot();
413                        self.infected_types.insert(new_id);
414                        new_id.into()
415                    }
416                    ItemId::Fun(_) => self.ctx.translated.fun_decls.reserve_slot().into(),
417                    ItemId::Global(_) => self.ctx.translated.global_decls.reserve_slot().into(),
418                    ItemId::TraitDecl(_) => self.ctx.translated.trait_decls.reserve_slot().into(),
419                    ItemId::TraitImpl(_) => self.ctx.translated.trait_impls.reserve_slot().into(),
420                };
421                self.generic_params.insert(new_id, new_params);
422                self.reverse_shape_map.insert(new_id, key);
423                self.to_process.push_back(new_id);
424                new_id
425            });
426
427        let fmt_ctx = self.ctx.into_fmt();
428        trace!(
429            "processing {}{}\n output: {}{}",
430            id.with_ctx(&fmt_ctx),
431            generics.with_ctx(&fmt_ctx),
432            new_id.with_ctx(&fmt_ctx),
433            shape_args.with_ctx(&fmt_ctx),
434        );
435        Some(DeclRef {
436            id: new_id,
437            generics: Box::new(shape_args),
438            trait_ref: None,
439        })
440    }
441
442    /// Traverse the item, replacing any type instantiations we don't want with references to
443    /// soon-to-be-created partially-monomorphized types. This does not access the items in
444    /// `self.translated`, which may be missing since we took `item` out for processing.
445    pub fn process_item(&mut self, item: &mut ItemRefMut<'_>) {
446        let _ = item.drive_mut(self);
447    }
448
449    /// Creates the item corresponding to this id by instantiating the item it is based on.
450    ///
451    /// This accesses the items in `self.translated`, which must therefore all be there.
452    /// That's why items are created outside of `process_item`.
453    pub fn create_pending_instantiation(&mut self, new_id: ItemId) -> ItemByVal {
454        let (orig_id, shape) = &self.reverse_shape_map[&new_id];
455        let mut decl = self
456            .ctx
457            .translated
458            .get_item(*orig_id)
459            .unwrap()
460            .to_owned()
461            .substitute_with_self(&shape.skip_binder, &TraitRefKind::SelfId);
462
463        let mut decl_mut = decl.as_mut();
464        decl_mut.set_id(new_id);
465        *decl_mut.generic_params() = shape.params.clone();
466
467        let name_ref = &mut decl_mut.item_meta().name;
468        *name_ref = mem::take::<crate::ast::Name>(name_ref).instantiate(shape.clone());
469        self.ctx
470            .translated
471            .item_names
472            .insert(new_id, decl.as_ref().item_meta().name.clone());
473        if let (ItemId::TraitDecl(orig_trait_id), ItemId::TraitDecl(new_trait_id)) =
474            (*orig_id, new_id)
475        {
476            let names = self.ctx.translated.assoc_item_names[orig_trait_id].clone();
477            self.ctx
478                .translated
479                .assoc_item_names
480                .insert(new_trait_id, names);
481        }
482
483        decl
484    }
485}
486
487impl VisitorWithSpan for PartialMonomorphizer<'_> {
488    fn current_span(&mut self) -> &mut Span {
489        &mut self.span
490    }
491}
492impl VisitAstMut for PartialMonomorphizer<'_> {
493    fn visit<T: AstVisitable>(&mut self, x: &mut T) -> ControlFlow<Self::Break> {
494        // Track a useful enclosing span, for error messages.
495        VisitWithSpan::new(self).visit(x)
496    }
497
498    fn exit_type_decl_ref(&mut self, x: &mut TypeDeclRef) {
499        if self.specialize_adts
500            && let Some(id) = x.as_adt()
501            && let Some(new_decl_ref) = self.process_generics(id.into(), &x.generics)
502        {
503            *x = new_decl_ref.try_into().unwrap()
504        }
505    }
506    fn exit_fn_ptr(&mut self, x: &mut FnPtr) {
507        // TODO: methods. any `Trait::method<&mut A>` requires monomorphizing all the instances of
508        // that method just in case :>>>
509        if let FnPtrKind::Fun(FunId::Regular(id)) = *x.kind
510            && let Some(new_decl_ref) = self.process_generics(id.into(), &x.generics)
511        {
512            *x = new_decl_ref.try_into().unwrap()
513        }
514    }
515    fn exit_fun_decl_ref(&mut self, x: &mut FunDeclRef) {
516        if let Some(new_decl_ref) = self.process_generics(x.id.into(), &x.generics) {
517            *x = new_decl_ref.try_into().unwrap()
518        }
519    }
520    fn exit_global_decl_ref(&mut self, x: &mut GlobalDeclRef) {
521        if let Some(new_decl_ref) = self.process_generics(x.id.into(), &x.generics) {
522            *x = new_decl_ref.try_into().unwrap()
523        }
524    }
525    fn exit_trait_decl_ref(&mut self, x: &mut TraitDeclRef) {
526        if let Some(new_decl_ref) = self.process_generics(x.id.into(), &x.generics) {
527            *x = new_decl_ref.try_into().unwrap()
528        }
529    }
530    fn exit_trait_impl_ref(&mut self, x: &mut TraitImplRef) {
531        if let Some(new_decl_ref) = self.process_generics(x.id.into(), &x.generics) {
532            *x = new_decl_ref.try_into().unwrap()
533        }
534    }
535}
536
537pub struct Transform;
538impl TransformPass for Transform {
539    fn transform_ctx(&self, ctx: &mut TransformCtx) {
540        let Some(include_types) = ctx.options.monomorphize_mut else {
541            return;
542        };
543        // TODO: test name matcher, also with methods
544        let mut visitor =
545            PartialMonomorphizer::new(ctx, matches!(include_types, MonomorphizeMut::All));
546        while let Some(id) = visitor.to_process.pop_front() {
547            // Get the item corresponding to this id, either by creating it or by getting an
548            // existing one.
549            let mut decl = if visitor.reverse_shape_map.contains_key(&id) {
550                // Create the required item by instantiating the item it's based on.
551                visitor.create_pending_instantiation(id)
552            } else {
553                // Take the item out so we can modify it. Warning: don't look up other items in the
554                // meantime as this would break in recursive cases.
555                match visitor.ctx.translated.remove_item_temporarily(id) {
556                    Some(decl) => decl,
557                    None => continue,
558                }
559            };
560            // Visit the item, replacing type instantiations with references to soon-to-be-created
561            // partially-monomorphized types.
562            visitor.process_item(&mut decl.as_mut());
563            // Put the item back.
564            visitor.ctx.translated.put_item_back(id, decl);
565        }
566    }
567}
charon_lib/transform/normalize/partial_monomorphization.rs

charon_lib/transform/normalize/
partial_monomorphization.rs
