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(tref, type_id, item_clause_id) = trkind else {
81                        unreachable!()
82                    };
83                    let new_id = (|| {
84                        let new_id = *trait_item_clause_ids
85                            .get(tref.trait_decl_ref.skip_binder.id)?
86                            .get(type_id)?
87                            .get(item_clause_id)?;
88                        Some(new_id)
89                    })();
90                    match new_id {
91                        Some(new_id) => ParentClause(tref, new_id),
92                        None => ItemClause(tref, type_id, item_clause_id),
93                    }
94                }),
95                BuiltinOrAuto {
96                    parent_trait_refs,
97                    types,
98                    ..
99                } => {
100                    for assoc_ty in types.iter_mut() {
101                        for tref in std::mem::take(&mut assoc_ty.implied_trait_refs) {
102                            // Note: this assumes that we listed the types in the same order as in
103                            // the trait decl, which we do.
104                            parent_trait_refs.push(tref);
105                        }
106                    }
107                }
108                _ => {}
109            }
110        });
111    }
112}