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::collections::HashMap;
4use std::mem;
5
6use crate::{
7    ast::*,
8    ids::{IndexMap, IndexVec},
9};
10
11use crate::transform::{TransformCtx, ctx::TransformPass};
12
13pub struct Transform;
14impl TransformPass for Transform {
15    fn should_run(&self, options: &crate::options::TranslateOptions) -> bool {
16        !options.no_normalize
17    }
18    fn transform_ctx(&self, ctx: &mut TransformCtx) {
19        // For each trait, we move the item-local clauses to be top-level parent clauses, and
20        // record the mapping from the old to the new ids.
21        let trait_item_clause_ids: IndexMap<
22            TraitDeclId,
23            HashMap<TraitItemName, IndexVec<TraitClauseId, TraitClauseId>>,
24        > = ctx.translated.trait_decls.map_ref_mut(|decl| {
25            decl.types
26                .iter_mut()
27                .filter(|assoc_ty| !assoc_ty.params.has_explicits())
28                .map(|assoc_ty| {
29                    let id_map =
30                        mem::take(&mut assoc_ty.skip_binder.implied_clauses).map(|clause| {
31                            let mut clause = clause.move_from_under_binder().unwrap();
32                            decl.implied_clauses.push_with(|id| {
33                                clause.clause_id = id;
34                                clause
35                            })
36                        });
37                    if assoc_ty.params.trait_clauses.is_empty() {
38                        // Move non-trait-clause-predicates of non-GAT types to be predicates on
39                        // the trait itself.
40                        decl.generics.take_predicates_from(
41                            mem::take(&mut assoc_ty.params)
42                                .move_from_under_binder()
43                                .unwrap(),
44                        );
45                    }
46                    (*assoc_ty.name(), id_map)
47                })
48                .collect()
49        });
50
51        // Move the item-local trait refs to match what we did in the trait declarations.
52        for timpl in ctx.translated.trait_impls.iter_mut() {
53            for (_, assoc_ty) in &mut timpl.types {
54                if !assoc_ty.params.has_explicits() {
55                    for trait_ref in mem::take(&mut assoc_ty.skip_binder.implied_trait_refs) {
56                        let trait_ref = trait_ref.move_from_under_binder().unwrap();
57                        // Note: this assumes that we listed the types in the same order as in the
58                        // trait decl, which we do.
59                        timpl.implied_trait_refs.push(trait_ref);
60                    }
61                }
62            }
63        }
64
65        // Update trait refs.
66        ctx.translated.dyn_visit_mut(|trkind: &mut TraitRefKind| {
67            use TraitRefKind::*;
68            match trkind {
69                ItemClause(..) => take_mut::take(trkind, |trkind| {
70                    let ItemClause(tref, item_name, item_clause_id) = trkind else {
71                        unreachable!()
72                    };
73                    let new_id = (|| {
74                        let new_id = *trait_item_clause_ids
75                            .get(tref.trait_decl_ref.skip_binder.id)?
76                            .get(&item_name)?
77                            .get(item_clause_id)?;
78                        Some(new_id)
79                    })();
80                    match new_id {
81                        Some(new_id) => ParentClause(tref, new_id),
82                        None => ItemClause(tref, item_name, item_clause_id),
83                    }
84                }),
85                BuiltinOrAuto {
86                    parent_trait_refs,
87                    types,
88                    ..
89                } => {
90                    for (_, assoc_ty) in types {
91                        for tref in std::mem::take(&mut assoc_ty.implied_trait_refs) {
92                            // Note: this assumes that we listed the types in the same order as in
93                            // the trait decl, which we do.
94                            parent_trait_refs.push(tref);
95                        }
96                    }
97                }
98                _ => {}
99            }
100        });
101    }
102}