Skip to main content

charon_lib/transform/simplify_output/
lift_associated_item_clauses.rs

1//! Move clauses on non-generic associated types to be implied clauses of the trait. The
2//! distinction is not semantically meaningful.
3use std::mem;
4
5use crate::{
6    ast::*,
7    ids::{IndexMap, IndexVec},
8};
9
10use crate::transform::{TransformCtx, ctx::TransformPass};
11
12pub struct Transform;
13impl TransformPass for Transform {
14    fn should_run(&self, options: &crate::options::TranslateOptions) -> bool {
15        !options.no_normalize
16    }
17    fn transform_ctx(&self, ctx: &mut TransformCtx) {
18        // For each trait, we move the item-local clauses to be top-level parent clauses, and
19        // record the mapping from the old to the new ids.
20        let trait_item_clause_ids: IndexMap<
21            TraitDeclId,
22            IndexMap<AssocTypeId, IndexVec<TraitClauseId, TraitClauseId>>,
23        > = ctx.translated.trait_decls.map_ref_mut(|decl| {
24            let mut map = IndexMap::default();
25            for (type_id, assoc_ty) in decl
26                .types
27                .iter_mut_enumerated()
28                .filter(|(_, assoc_ty)| !assoc_ty.params.has_explicits())
29            {
30                // Try to move them all out of the binder, or bail.
31                if let Some(clauses) = assoc_ty
32                    .skip_binder
33                    .clone()
34                    .implied_clauses
35                    .move_from_under_binder()
36                {
37                    assoc_ty.skip_binder.implied_clauses.clear();
38                    let id_map = clauses.map(|mut clause| {
39                        decl.implied_clauses.push_with(|id| {
40                            clause.clause_id = id;
41                            clause
42                        })
43                    });
44                    if assoc_ty.params.trait_clauses.is_empty() {
45                        // Move non-trait-clause-predicates of non-GAT types to be predicates on
46                        // the trait itself.
47                        decl.generics.take_predicates_from(
48                            mem::take(&mut assoc_ty.params)
49                                .move_from_under_binder()
50                                .unwrap(),
51                        );
52                    }
53                    map.set_slot_extend(type_id, id_map);
54                }
55            }
56            map
57        });
58
59        // Move the item-local trait refs to match what we did in the trait declarations.
60        for timpl in ctx.translated.trait_impls.iter_mut() {
61            for (type_id, assoc_ty) in timpl.types.iter_mut_indexed() {
62                if let Some(m) = trait_item_clause_ids.get(timpl.impl_trait.id)
63                    && m.get(type_id).is_some()
64                {
65                    for trait_ref in mem::take(&mut assoc_ty.skip_binder.implied_trait_refs) {
66                        let trait_ref = trait_ref.move_from_under_binder().unwrap();
67                        // Note: this assumes that we listed the types in the same order as in the
68                        // trait decl, which we do.
69                        timpl.implied_trait_refs.push(trait_ref);
70                    }
71                }
72            }
73        }
74
75        // Update trait refs.
76        ctx.translated.dyn_visit_mut(|trkind: &mut TraitRefKind| {
77            use TraitRefKind::*;
78            match trkind {
79                ItemClause { .. } => take_mut::take(trkind, |trkind| {
80                    let ItemClause {
81                        trait_ref,
82                        type_id,
83                        generics,
84                        clause_id,
85                    } = trkind
86                    else {
87                        unreachable!()
88                    };
89                    let new_id = (|| {
90                        let new_id = *trait_item_clause_ids
91                            .get(trait_ref.trait_decl_ref.skip_binder.id)?
92                            .get(type_id)?
93                            .get(clause_id)?;
94                        Some(new_id)
95                    })();
96                    match new_id {
97                        Some(new_id) => ParentClause(trait_ref, new_id),
98                        None => ItemClause {
99                            trait_ref,
100                            type_id,
101                            generics,
102                            clause_id,
103                        },
104                    }
105                }),
106                BuiltinOrAuto {
107                    parent_trait_refs,
108                    types,
109                    ..
110                } => {
111                    for assoc_ty in types.iter_mut() {
112                        for tref in std::mem::take(&mut assoc_ty.implied_trait_refs) {
113                            // Note: this assumes that we listed the types in the same order as in
114                            // the trait decl, which we do. And that builtin assoc types are never
115                            // GATs, which they so far aren't.
116                            parent_trait_refs.push(tref);
117                        }
118                    }
119                }
120                _ => {}
121            }
122        });
123    }
124}