charon_lib/transform/
remove_read_discriminant.rs

1//! The MIR code often contains variables with type `Never`, and we want to get
2//! rid of those. We proceed in two steps. First, we remove the instructions
3//! `drop(v)` where `v` has type `Never` (it can happen - this module does the
4//! filtering). Then, we filter the unused variables ([crate::remove_unused_locals]).
5
6use crate::errors::register_error;
7use crate::formatter::IntoFormatter;
8use crate::llbc_ast::*;
9use crate::name_matcher::NamePattern;
10use crate::pretty::FmtWithCtx;
11use crate::transform::TransformCtx;
12use itertools::Itertools;
13use std::collections::{HashMap, HashSet};
14
15use super::ctx::LlbcPass;
16
17/// Generate `match _y { 0 => { _x = 0 }, 1 => { _x = 1; }, .. }`.
18fn generate_discr_assignment(
19    span: Span,
20    variants: &Vector<VariantId, Variant>,
21    scrutinee: &Place,
22    dest: &Place,
23) -> RawStatement {
24    let targets = variants
25        .iter_indexed_values()
26        .map(|(id, variant)| {
27            let discr_value = Rvalue::Use(Operand::Const(variant.discriminant.to_constant()));
28            let statement = Statement::new(span, RawStatement::Assign(dest.clone(), discr_value));
29            (vec![id], statement.into_block())
30        })
31        .collect();
32    RawStatement::Switch(Switch::Match(scrutinee.clone(), targets, None))
33}
34
35pub struct Transform;
36impl Transform {
37    fn update_block(
38        ctx: &mut TransformCtx,
39        block: &mut Block,
40        discriminant_intrinsic: Option<FunDeclId>,
41    ) {
42        // Iterate through the statements.
43        for i in 0..block.statements.len() {
44            let suffix = &mut block.statements[i..];
45            match suffix {
46                [Statement {
47                    content: RawStatement::Assign(dest, Rvalue::Discriminant(p, adt_id)),
48                    span: span1,
49                    ..
50                }, rest @ ..] => {
51                    // The destination should be a variable
52                    assert!(dest.is_local());
53
54                    // Lookup the type of the scrutinee
55                    let tkind = ctx.translated.type_decls.get(*adt_id).map(|x| &x.kind);
56                    let Some(TypeDeclKind::Enum(variants)) = tkind else {
57                        match tkind {
58                            // This can happen if the type was declared as invisible or opaque.
59                            None | Some(TypeDeclKind::Opaque) => {
60                                let name = ctx.translated.item_name(*adt_id).unwrap();
61                                register_error!(
62                                    ctx,
63                                    block.span,
64                                    "reading the discriminant of an opaque enum. \
65                                    Add `--include {}` to the `charon` arguments \
66                                    to translate this enum.",
67                                    name.fmt_with_ctx(&ctx.into_fmt())
68                                );
69                            }
70                            // Don't double-error
71                            Some(TypeDeclKind::Error(..)) => {}
72                            Some(_) => {
73                                register_error!(
74                                    ctx,
75                                    block.span,
76                                    "reading the discriminant of a non-enum type"
77                                );
78                            }
79                        }
80                        block.statements[i].content = RawStatement::Error(
81                            "error reading the discriminant of this type".to_owned(),
82                        );
83                        return;
84                    };
85
86                    // We look for a `SwitchInt` just after the discriminant read.
87                    match rest {
88                        [Statement {
89                            content:
90                                RawStatement::Switch(switch @ Switch::SwitchInt(Operand::Move(_), ..)),
91                            ..
92                        }, ..] => {
93                            // Convert between discriminants and variant indices. Remark: the discriminant can
94                            // be of any *signed* integer type (`isize`, `i8`, etc.).
95                            let discr_to_id: HashMap<ScalarValue, VariantId> = variants
96                                .iter_indexed_values()
97                                .map(|(id, variant)| (variant.discriminant, id))
98                                .collect();
99
100                            take_mut::take(switch, |switch| {
101                                let (Operand::Move(op_p), _, targets, otherwise) =
102                                    switch.to_switch_int().unwrap()
103                                else {
104                                    unreachable!()
105                                };
106                                assert!(op_p.is_local() && op_p.local_id() == dest.local_id());
107
108                                let mut covered_discriminants: HashSet<ScalarValue> =
109                                    HashSet::default();
110                                let targets = targets
111                                    .into_iter()
112                                    .map(|(v, e)| {
113                                        let targets = v
114                                            .into_iter()
115                                            .filter_map(|discr| {
116                                                covered_discriminants.insert(discr);
117                                                discr_to_id.get(&discr).or_else(|| {
118                                                    register_error!(
119                                                        ctx,
120                                                        block.span,
121                                                        "Found incorrect discriminant \
122                                                        {discr} for enum {adt_id}"
123                                                    );
124                                                    None
125                                                })
126                                            })
127                                            .copied()
128                                            .collect_vec();
129                                        (targets, e)
130                                    })
131                                    .collect_vec();
132                                // Filter the otherwise branch if it is not necessary.
133                                let covers_all = covered_discriminants.len() == discr_to_id.len();
134                                let otherwise = if covers_all { None } else { Some(otherwise) };
135
136                                // Replace the old switch with a match.
137                                Switch::Match(p.clone(), targets, otherwise)
138                            });
139                            // `Nop` the discriminant read.
140                            block.statements[i].content = RawStatement::Nop;
141                        }
142                        _ => {
143                            // The discriminant read is not followed by a `SwitchInt`. This can happen
144                            // in optimized MIR. We replace `_x = Discr(_y)` with `match _y { 0 => { _x
145                            // = 0 }, 1 => { _x = 1; }, .. }`.
146                            block.statements[i].content =
147                                generate_discr_assignment(*span1, variants, p, dest)
148                        }
149                    }
150                }
151                // Replace calls of `core::intrinsics::discriminant_value` on a known enum with the
152                // appropriate MIR.
153                [Statement {
154                    content: RawStatement::Call(call),
155                    span: span1,
156                    ..
157                }, ..]
158                    if let Some(discriminant_intrinsic) = discriminant_intrinsic
159                        // Detect a call to the intrinsic...
160                        && let FnOperand::Regular(fn_ptr) = &call.func
161                        && let FunIdOrTraitMethodRef::Fun(FunId::Regular(fun_id)) = fn_ptr.func
162                        && fun_id == discriminant_intrinsic
163                        // on a known enum...
164                        && let ty = &fn_ptr.generics.types[0]
165                        && let TyKind::Adt(TypeId::Adt(type_id), _) = *ty.kind()
166                        && let Some(TypeDeclKind::Enum(variants)) =
167                            ctx.translated.type_decls.get(type_id).map(|x| &x.kind)
168                        // passing it a reference.
169                        && let Operand::Move(p) = &call.args[0]
170                        && let TyKind::Ref(_, sub_ty, _) = p.ty().kind() =>
171                {
172                    let p = p.clone().project(ProjectionElem::Deref, sub_ty.clone());
173                    block.statements[i].content =
174                        generate_discr_assignment(*span1, variants, &p, &call.dest)
175                }
176                _ => {}
177            }
178        }
179    }
180}
181
182const DISCRIMINANT_INTRINSIC: &str = "core::intrinsics::discriminant_value";
183
184impl LlbcPass for Transform {
185    fn transform_ctx(&self, ctx: &mut TransformCtx) {
186        let pat = NamePattern::parse(DISCRIMINANT_INTRINSIC).unwrap();
187        let discriminant_intrinsic: Option<FunDeclId> = ctx
188            .translated
189            .item_names
190            .iter()
191            .find(|(_, name)| pat.matches(&ctx.translated, name))
192            .and_then(|(id, _)| id.as_fun())
193            .copied();
194
195        ctx.for_each_fun_decl(|ctx, decl| {
196            if let Ok(body) = &mut decl.body {
197                let body = body.as_structured_mut().unwrap();
198                self.log_before_body(ctx, &decl.item_meta.name, Ok(&*body));
199                body.body.visit_blocks_bwd(|block: &mut Block| {
200                    Transform::update_block(ctx, block, discriminant_intrinsic);
201                });
202            };
203        });
204    }
205}