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<TypeId>,
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<TypeId> = 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            .map(TypeId::Adt)
47            .chain([TypeId::Builtin(BuiltinTy::Box)])
48            .collect();
49
50        for &id in &types {
51            if let Some(&id) = id.as_adt()
52                && let Some(tdecl) = ctx.translated.type_decls.get_mut(id)
53            {
54                if tdecl.generics.types.is_empty() {
55                    // We monomorpohized this type.
56                    let args = tdecl.item_meta.name.mono_args_mut().unwrap();
57                    args.types.pop().unwrap();
58                } else {
59                    struct SubstWithErrorVisitor(TypeVarId);
60                    impl VarsVisitor for SubstWithErrorVisitor {
61                        fn visit_type_var(&mut self, v: TypeDbVar) -> Option<Ty> {
62                            if let DeBruijnVar::Bound(DeBruijnId::ZERO, var_id) = v
63                                && var_id == self.0
64                            {
65                                Some(
66                                    TyKind::Error("removed allocator parameter".to_owned())
67                                        .into_ty(),
68                                )
69                            } else {
70                                None
71                            }
72                        }
73                    }
74                    let tvar = tdecl.generics.types.pop().unwrap();
75                    tdecl.visit_vars(&mut SubstWithErrorVisitor(tvar.index));
76                }
77            }
78        }
79
80        let _ = ctx
81            .translated
82            .drive_mut(&mut RemoveLastParamVisitor { types });
83    }
84}