1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

//! The MIR code often contains variables with type `Never`, and we want to get
//! rid of those. We proceed in two steps. First, we remove the instructions
//! `drop(v)` where `v` has type `Never` (it can happen - this module does the
//! filtering). Then, we filter the unused variables ([crate::remove_unused_locals]).

use crate::errors::register_error_or_panic;
use crate::llbc_ast::*;
use crate::transform::TransformCtx;
use derive_visitor::visitor_enter_fn_mut;
use derive_visitor::DriveMut;
use itertools::Itertools;
use std::collections::{HashMap, HashSet};

use super::ctx::LlbcPass;

pub struct Transform;
impl Transform {
    fn update_block(ctx: &mut TransformCtx<'_>, block: &mut Block) {
        // Iterate through the statements.
        for i in 0..block.statements.len() {
            let suffix = &mut block.statements[i..];
            if let [Statement {
                content: RawStatement::Assign(dest, Rvalue::Discriminant(p, adt_id)),
                span: span1,
                ..
            }, rest @ ..] = suffix
            {
                // The destination should be a variable
                assert!(dest.projection.is_empty());

                // Lookup the type of the scrutinee
                let variants = match ctx.translated.type_decls.get(*adt_id) {
                    // This can happen if there was an error while extracting the definitions
                    None => None,
                    Some(d) => {
                        match &d.kind {
                            TypeDeclKind::Struct(..)
                            | TypeDeclKind::Union(..)
                            | TypeDeclKind::Opaque
                            | TypeDeclKind::Alias(..) => {
                                // We shouldn't get there
                                register_error_or_panic!(ctx, block.span, "Unreachable case");
                                None
                            }
                            TypeDeclKind::Error(_) => None,
                            TypeDeclKind::Enum(variants) => Some(variants),
                        }
                    }
                };
                let Some(variants) = variants else {
                    // An error occurred. We can't keep the `Rvalue::Discriminant` around so we
                    // `Nop` the discriminant read.
                    block.statements[i].content = RawStatement::Nop;
                    return;
                };

                // We look for a `SwitchInt` just after the discriminant read.
                match rest {
                    [Statement {
                        content:
                            RawStatement::Switch(switch @ Switch::SwitchInt(Operand::Move(_), ..)),
                        ..
                    }, ..] => {
                        // Convert between discriminants and variant indices. Remark: the discriminant can
                        // be of any *signed* integer type (`isize`, `i8`, etc.).
                        let discr_to_id: HashMap<ScalarValue, VariantId> = variants
                            .iter_indexed_values()
                            .map(|(id, variant)| (variant.discriminant, id))
                            .collect();

                        take_mut::take(switch, |switch| {
                            let (Operand::Move(op_p), _, targets, otherwise) =
                                switch.to_switch_int().unwrap()
                            else {
                                unreachable!()
                            };
                            assert!(op_p.projection.is_empty() && op_p.var_id == dest.var_id);

                            let mut covered_discriminants: HashSet<ScalarValue> =
                                HashSet::default();
                            let targets = targets
                                .into_iter()
                                .map(|(v, e)| {
                                    (
                                        v.into_iter()
                                            .filter_map(|discr| {
                                                covered_discriminants.insert(discr);
                                                discr_to_id.get(&discr).or_else(|| {
                                                    register_error_or_panic!(
                                                        ctx,
                                                        block.span,
                                                        "Found incorrect discriminant {discr} for enum {adt_id}"
                                                    );
                                                    None
                                                })
                                            })
                                            .copied()
                                            .collect_vec(),
                                        e,
                                    )
                                })
                                .collect_vec();
                            // Filter the otherwise branch if it is not necessary.
                            let covers_all = covered_discriminants.len() == discr_to_id.len();
                            let otherwise = if covers_all { None } else { Some(otherwise) };

                            // Replace the old switch with a match.
                            Switch::Match(p.clone(), targets, otherwise)
                        });
                        // `Nop` the discriminant read.
                        block.statements[i].content = RawStatement::Nop;
                    }
                    _ => {
                        // The discriminant read is not followed by a `SwitchInt`. This can happen
                        // in optimized MIR. We replace `_x = Discr(_y)` with `match _y { 0 => { _x
                        // = 0 }, 1 => { _x = 1; }, .. }`.
                        let targets = variants
                            .iter_indexed_values()
                            .map(|(id, variant)| {
                                let discr_value =
                                    Rvalue::Use(Operand::Const(variant.discriminant.to_constant()));
                                let statement = Statement::new(
                                    *span1,
                                    RawStatement::Assign(dest.clone(), discr_value),
                                );
                                (vec![id], statement.into_block())
                            })
                            .collect();
                        block.statements[i].content =
                            RawStatement::Switch(Switch::Match(p.clone(), targets, None))
                    }
                }
            }
        }
    }
}

impl LlbcPass for Transform {
    fn transform_body(&self, ctx: &mut TransformCtx<'_>, b: &mut ExprBody) {
        b.body
            .drive_mut(&mut visitor_enter_fn_mut(|block: &mut Block| {
                Transform::update_block(ctx, block);
            }));
    }
}