charon_lib/transform/resugar/
reconstruct_matches.rs

1//! The way to match on enums in MIR is in two steps: first read the discriminant, then switch on
2//! the resulting integer. This pass records the enum place as the switch scrutinee and replaces
3//! the integer cases with symbolic enum discriminants.
4use std::{
5    collections::{HashMap, HashSet},
6    mem,
7};
8
9use crate::formatter::IntoFormatter;
10use crate::name_matcher::NamePattern;
11use crate::pretty::FmtWithCtx;
12use crate::transform::TransformCtx;
13use crate::transform::ctx::UllbcPass;
14use crate::ullbc_ast::*;
15use crate::{errors::register_error, transform::CowBox};
16
17pub struct Transform {
18    discriminant_intrinsics: HashSet<FunDeclId>,
19}
20
21impl Transform {
22    fn replace_mem_discriminant_call(&self, block: &mut BlockData) {
23        if let TerminatorKind::Call { call, target, .. } = &block.terminator.kind
24            && let FnOperand::Regular(fn_ptr) = &call.func
25            && let FnPtrKind::Fun(fun_id) = fn_ptr.kind.as_ref()
26            && self.discriminant_intrinsics.contains(fun_id)
27            && let [Operand::Move(p)] = call.args.as_slice()
28            && let TyKind::Ref(_, sub_ty, _) = p.ty().kind()
29        {
30            let dest = call.dest.clone();
31            let p = p.clone().project(ProjectionElem::Deref, sub_ty.clone());
32            let target = *target;
33            let mut statement = Statement::new(
34                block.terminator.span,
35                StatementKind::Assign(dest, Rvalue::Discriminant(p)),
36            );
37            statement.comments_before = mem::take(&mut block.terminator.comments_before);
38            block.statements.push(statement);
39            block.terminator.kind = TerminatorKind::Goto { target };
40        }
41    }
42
43    fn reconstruct_match(&self, ctx: &mut TransformCtx, block: &mut BlockData) {
44        let span = block.terminator.span;
45        let TerminatorKind::Switch { data, .. } = &mut block.terminator.kind else {
46            return;
47        };
48        let SwitchScrutinee::Value(Operand::Move(op_p)) = &data.scrutinee else {
49            return;
50        };
51        // If the last statement is a discriminant read.
52        let Some(last_st) = block.statements.last_mut() else {
53            return;
54        };
55        let StatementKind::Assign(dest, Rvalue::Discriminant(p)) = &last_st.kind else {
56            return;
57        };
58        assert!(dest.is_local()); // The destination should be a variable.
59        if op_p != dest {
60            return;
61        }
62        let scrut_ty = p.ty().clone();
63
64        // Merge the two statements.
65        data.scrutinee = SwitchScrutinee::Discriminant(p.clone());
66        last_st.kind = StatementKind::Nop;
67
68        // Lookup the type of the scrutinee.
69        let TyKind::Adt(tdecl_ref) = scrut_ty.kind() else {
70            return;
71        };
72        let tdecl_ref = tdecl_ref.clone();
73        let adt_id = tdecl_ref.id;
74        let tkind = ctx.translated.type_decls.get(adt_id).map(|x| &x.kind);
75        let Some(TypeDeclKind::Enum(variants)) = tkind else {
76            match tkind {
77                // This can happen if the type was declared as invisible or opaque.
78                None | Some(TypeDeclKind::Opaque) => {
79                    let name = ctx.translated.item_name(adt_id);
80                    register_error!(
81                        ctx,
82                        span,
83                        "reading the discriminant of an opaque enum. \
84                        Add `--include {}` to the `charon` arguments \
85                        to translate this enum.",
86                        name.with_ctx(&ctx.into_fmt())
87                    );
88                }
89                // Don't double-error.
90                Some(TypeDeclKind::Error(..)) => {}
91                Some(_) => {
92                    register_error!(ctx, span, "reading the discriminant of a non-enum type");
93                }
94            }
95            return;
96        };
97
98        // Map from discriminants to variant indices. The discriminant can be of any integer type.
99        let discr_to_id: HashMap<IntegerValue, VariantId> = variants
100            .iter_enumerated()
101            .map(|(id, variant)| (variant.discriminant, id))
102            .collect();
103
104        // Replace the branch values with discriminant constants.
105        let mut covered_discriminants: HashSet<IntegerValue> = HashSet::default();
106        for (value, _) in &mut data.branches {
107            if let ConstantExprKind::Integer(discr) = value.kind()
108                && let Some(variant_id) = discr_to_id.get(discr).copied()
109            {
110                covered_discriminants.insert(*discr);
111                *value = ConstantExpr::new(
112                    ConstantExprKind::Discriminant(tdecl_ref.clone(), variant_id),
113                    value.ty().clone(),
114                );
115            } else {
116                register_error!(
117                    ctx,
118                    block.terminator.span,
119                    "Found incorrect discriminant {value} for enum {adt_id}"
120                );
121            }
122        }
123
124        // Remove the fallback if the explicit cases cover every variant.
125        if covered_discriminants.len() == discr_to_id.len() {
126            data.fallback.take();
127        }
128    }
129}
130
131const DISCRIMINANT_INTRINSIC: &str = "core::intrinsics::discriminant_value";
132
133impl Transform {
134    pub fn new(ctx: &mut TransformCtx) -> CowBox<dyn UllbcPass> {
135        let pat = NamePattern::parse(DISCRIMINANT_INTRINSIC).unwrap();
136        // There can be many if we're in mono mode.
137        let discriminant_intrinsics = ctx
138            .translated
139            .item_names
140            .iter()
141            .filter(|(_, name)| pat.matches(&ctx.translated, name))
142            .filter_map(|(id, _)| id.as_fun())
143            .copied()
144            .collect();
145        CowBox::Owned(Box::new(Transform {
146            discriminant_intrinsics,
147        }))
148    }
149}
150
151impl UllbcPass for Transform {
152    fn should_run(&self, options: &crate::options::TranslateOptions) -> bool {
153        options.reconstruct_matches
154    }
155
156    fn transform_body(&self, ctx: &mut TransformCtx, body: &mut ExprBody) {
157        for block in &mut body.body {
158            self.reconstruct_match(ctx, block);
159            self.replace_mem_discriminant_call(block);
160        }
161    }
162}
charon_lib/transform/resugar/reconstruct_matches.rs

charon_lib/transform/resugar/
reconstruct_matches.rs
