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