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}