Skip to main content

charon_lib/transform/simplify_output/
hide_allocator_param.rs

1use itertools::Itertools;
2use std::collections::HashSet;
3
4use crate::{ast::*, name_matcher::NamePattern};
5
6use crate::transform::{TransformCtx, ctx::TransformPass};
7
8#[derive(Visitor)]
9struct RemoveLastParamVisitor {
10    types: HashSet<TypeDeclId>,
11}
12
13impl VisitAstMut for RemoveLastParamVisitor {
14    fn enter_type_decl_ref(&mut self, x: &mut TypeDeclRef) {
15        if self.types.contains(&x.id) {
16            // Remove the last param.
17            x.generics.types.pop();
18        }
19    }
20}
21
22pub struct Transform;
23impl TransformPass for Transform {
24    fn transform_ctx(&self, ctx: &mut TransformCtx) {
25        if !ctx.options.hide_allocator {
26            return;
27        }
28        let types = &[
29            "alloc::boxed::Box",
30            "alloc::vec::Vec",
31            "alloc::rc::Rc",
32            "alloc::sync::Arc",
33        ];
34
35        let types: Vec<NamePattern> = types
36            .iter()
37            .map(|s| NamePattern::parse(s).unwrap())
38            .collect_vec();
39        let types: HashSet<TypeDeclId> = ctx
40            .translated
41            .item_names
42            .iter()
43            .filter(|(_, name)| types.iter().any(|p| p.matches(&ctx.translated, name)))
44            .filter_map(|(id, _)| id.as_type())
45            .copied()
46            .collect();
47
48        for &id in &types {
49            if let Some(tdecl) = ctx.translated.type_decls.get_mut(id) {
50                if tdecl.generics.types.is_empty() {
51                    // We monomorpohized this type.
52                    let args = tdecl.item_meta.name.mono_args_mut().unwrap();
53                    args.types.pop().unwrap();
54                } else {
55                    struct SubstWithErrorVisitor(TypeVarId);
56                    impl VarsVisitor for SubstWithErrorVisitor {
57                        fn visit_type_var(&mut self, v: TypeDbVar) -> Option<Ty> {
58                            if let DeBruijnVar::Bound(DeBruijnId::ZERO, var_id) = v
59                                && var_id == self.0
60                            {
61                                Some(
62                                    TyKind::Error("removed allocator parameter".to_owned())
63                                        .into_ty(),
64                                )
65                            } else {
66                                None
67                            }
68                        }
69                    }
70                    let tvar = tdecl.generics.types.pop().unwrap();
71                    tdecl.visit_vars(&mut SubstWithErrorVisitor(tvar.index));
72                }
73            }
74        }
75
76        let _ = ctx
77            .translated
78            .drive_mut(&mut RemoveLastParamVisitor { types });
79    }
80}