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

//! # Micro-pass: remove the overflow checks for arithmetic operations we couldn't remove in
//! [`remove_dynamic_checks`]. See comments there for more details.
use crate::transform::TransformCtx;
use crate::ullbc_ast::*;

use super::ctx::UllbcPass;

pub struct Transform;

impl Transform {
    /// After `remove_dynamic_checks`, the only remaining dynamic checks are overflow checks. We
    /// couldn't remove them in ullbc because the generated code spans two basic blocks. They are
    /// inserted only in constants since we otherwise compile in release mode. These assertions
    /// look like:
    /// ```text
    /// r := x checked.+ y;
    /// assert(move r.1 == false);
    /// z := move r.0;
    /// ```
    /// We replace that with:
    /// ```text
    /// z := x + y;
    /// ```
    fn update_statements(seq: &mut [Statement]) {
        if let [Statement {
            content:
                RawStatement::Assign(
                    binop,
                    Rvalue::BinaryOp(
                        op @ (BinOp::CheckedAdd | BinOp::CheckedSub | BinOp::CheckedMul),
                        _,
                        _,
                    ),
                ),
            ..
        }, Statement {
            content:
                RawStatement::Assert(Assert {
                    cond: Operand::Move(assert_cond),
                    expected: false,
                }),
            ..
        }, Statement {
            content: RawStatement::Assign(final_value, Rvalue::Use(Operand::Move(assigned))),
            ..
        }, ..] = seq
        {
            // assigned should be: binop.0
            // assert_cond should be: binop.1
            if let (
                [ProjectionElem::Field(FieldProjKind::Tuple(..), fid0)],
                [ProjectionElem::Field(FieldProjKind::Tuple(..), fid1)],
            ) = (
                assigned.projection.as_slice(),
                assert_cond.projection.as_slice(),
            ) {
                if assert_cond.var_id == binop.var_id
                    && assigned.var_id == binop.var_id
                    && binop.projection.len() == 0
                    && fid0.index() == 0
                    && fid1.index() == 1
                {
                    // Switch to the unchecked operation.
                    *op = match op {
                        BinOp::CheckedAdd => BinOp::Add,
                        BinOp::CheckedSub => BinOp::Sub,
                        BinOp::CheckedMul => BinOp::Mul,
                        _ => unreachable!(),
                    };
                    // Assign to the correct value in the first statement.
                    std::mem::swap(binop, final_value);
                    // Remove the other two statements.
                    seq[1].content = RawStatement::Nop;
                    seq[2].content = RawStatement::Nop;
                }
            }
        }
    }
}

impl UllbcPass for Transform {
    fn transform_body(&self, _ctx: &mut TransformCtx<'_>, b: &mut ExprBody) {
        b.transform_sequences(&mut |_, seq| {
            Transform::update_statements(seq);
            Vec::new()
        })
    }
}