1
use conjure_cp::{
2
    ast::Metadata,
3
    ast::{
4
        Atom, DeclarationKind, Expression, Literal, Moo, Name, ReturnType, SymbolTable, Typeable,
5
    },
6
    into_matrix_expr, matrix_expr,
7
    rule_engine::{
8
        ApplicationError::{self, RuleNotApplicable},
9
        ApplicationResult, Reduction, register_rule, register_rule_set,
10
    },
11
};
12
use uniplate::Biplate;
13

            
14
use super::utils::{is_all_constant, rewrite_children};
15

            
16
register_rule_set!("Bubble", ("Base"));
17

            
18
// Bubble reduction rules
19

            
20
/*
21
    Reduce bubbles with a boolean expression to a conjunction with their condition.
22

            
23
    e.g. (a / b = c) @ (b != 0) => (a / b = c) & (b != 0)
24
*/
25
#[register_rule("Bubble", 8900, [Bubble])]
26
2430875
fn expand_bubble(expr: &Expression, _: &SymbolTable) -> ApplicationResult {
27
1650
    match expr {
28
1650
        Expression::Bubble(_, a, b) if a.return_type() == ReturnType::Bool => {
29
768
            let a = Moo::unwrap_or_clone(Moo::clone(a));
30
768
            let b = Moo::unwrap_or_clone(Moo::clone(b));
31
768
            Ok(Reduction::pure(Expression::And(
32
768
                Metadata::new(),
33
768
                Moo::new(matrix_expr![a, b]),
34
768
            )))
35
        }
36
2430107
        _ => Err(ApplicationError::RuleNotApplicable),
37
    }
38
2430875
}
39

            
40
/*
41
    Bring bubbles with a non-boolean expression higher up the tree.
42

            
43
    E.g. ((a / b) @ (b != 0)) = c => (a / b = c) @ (b != 0)
44
*/
45
#[register_rule("Bubble", 8800)]
46
2171028
fn bubble_up(expr: &Expression, syms: &SymbolTable) -> ApplicationResult {
47
    // do not put root inside a bubble
48
    //
49
    // also do not bubble bubbles inside bubbles, as this does nothing productive it just shuffles
50
    // the conditions around, shuffles them back, then gets stuck in a loop doing this ad infinitum
51
2171028
    if matches!(expr, Expression::Root(_, _) | Expression::Bubble(_, _, _)) {
52
60632
        return Err(RuleNotApplicable);
53
2110396
    }
54

            
55
    // do not bubble things containing lettings
56
5325757
    if expr.universe_bi().iter().any(|x: &Name| {
57
5325757
        syms.lookup(x).is_some_and(|x| {
58
5054994
            matches!(
59
5065236
                &x.kind() as &DeclarationKind,
60
                DeclarationKind::ValueLetting(_, _)
61
            )
62
5065236
        })
63
5325757
    }) {
64
10242
        return Err(RuleNotApplicable);
65
2100154
    };
66

            
67
2100154
    let mut bubbled_conditions = vec![];
68
2100154
    let (new_inner, num_changed) = rewrite_children(expr, |child| match child {
69
882
        Expression::Bubble(_, a, b) if a.return_type() != ReturnType::Bool => {
70
882
            let a = Moo::unwrap_or_clone(a);
71
882
            let b = Moo::unwrap_or_clone(b);
72
882
            bubbled_conditions.push(b);
73
882
            (a, true)
74
        }
75
1782232
        child => (child, false),
76
1783114
    });
77
2100154
    if num_changed == 0 {
78
2099272
        Err(ApplicationError::RuleNotApplicable)
79
882
    } else if bubbled_conditions.len() == 1 {
80
882
        let new_expr = Expression::Bubble(
81
882
            Metadata::new(),
82
882
            Moo::new(new_inner),
83
882
            Moo::new(bubbled_conditions[0].clone()),
84
882
        );
85

            
86
882
        Ok(Reduction::pure(new_expr))
87
    } else {
88
        Ok(Reduction::pure(Expression::Bubble(
89
            Metadata::new(),
90
            Moo::new(new_inner),
91
            Moo::new(Expression::And(
92
                Metadata::new(),
93
                Moo::new(into_matrix_expr![bubbled_conditions]),
94
            )),
95
        )))
96
    }
97
2171028
}
98

            
99
// Bubble applications
100

            
101
/// Converts an unsafe division to a safe division with a bubble condition.
102
///
103
/// ```text
104
///     a / b => (a / b) @ (b != 0)
105
/// ```
106
///
107
/// Division by zero is undefined and therefore not allowed, so we add a condition to check for it.
108
/// This condition is brought up the tree and expanded into a conjunction with the first
109
/// boolean-type expression it is paired with.
110

            
111
#[register_rule("Bubble", 6000, [UnsafeDiv])]
112
1272697
fn div_to_bubble(expr: &Expression, _: &SymbolTable) -> ApplicationResult {
113
1272697
    if is_all_constant(expr) {
114
258646
        return Err(RuleNotApplicable);
115
1014051
    }
116
1014051
    if let Expression::UnsafeDiv(_, a, b) = expr {
117
        // bubble bottom up
118
648
        if !a.is_safe() || !b.is_safe() {
119
84
            return Err(RuleNotApplicable);
120
564
        }
121

            
122
564
        return Ok(Reduction::pure(Expression::Bubble(
123
564
            Metadata::new(),
124
564
            Moo::new(Expression::SafeDiv(Metadata::new(), a.clone(), b.clone())),
125
564
            Moo::new(Expression::Neq(
126
564
                Metadata::new(),
127
564
                b.clone(),
128
564
                Moo::new(Expression::from(0)),
129
564
            )),
130
564
        )));
131
1013403
    }
132
1013403
    Err(ApplicationError::RuleNotApplicable)
133
1272697
}
134

            
135
/// Converts an unsafe mod to a safe mod with a bubble condition.
136
///
137
/// ```text
138
/// a % b => (a % b) @ (b != 0)
139
/// ```
140
///
141
/// Mod zero is undefined and therefore not allowed, so we add a condition to check for it.
142
/// This condition is brought up the tree and expanded into a conjunction with the first
143
/// boolean-type expression it is paired with.
144
///
145
#[register_rule("Bubble", 6000, [UnsafeMod])]
146
1272697
fn mod_to_bubble(expr: &Expression, _: &SymbolTable) -> ApplicationResult {
147
1272697
    if is_all_constant(expr) {
148
258646
        return Err(RuleNotApplicable);
149
1014051
    }
150
1014051
    if let Expression::UnsafeMod(_, a, b) = expr {
151
        // bubble bottom up
152
192
        if !a.is_safe() || !b.is_safe() {
153
36
            return Err(RuleNotApplicable);
154
156
        }
155

            
156
156
        return Ok(Reduction::pure(Expression::Bubble(
157
156
            Metadata::new(),
158
156
            Moo::new(Expression::SafeMod(Metadata::new(), a.clone(), b.clone())),
159
156
            Moo::new(Expression::Neq(
160
156
                Metadata::new(),
161
156
                b.clone(),
162
156
                Moo::new(Expression::from(0)),
163
156
            )),
164
156
        )));
165
1013859
    }
166
1013859
    Err(ApplicationError::RuleNotApplicable)
167
1272697
}
168

            
169
/// Converts an unsafe pow to a safe pow with a bubble condition.
170
///
171
/// ```text
172
/// a**b => (a ** b) @ ((a!=0 \/ b!=0) /\ b>=0
173
/// ```
174
///
175
/// Pow is only defined when `(a!=0 \/ b!=0) /\ b>=0`, so we add a condition to check for it.
176
/// This condition is brought up the tree and expanded into a conjunction with the first
177
/// boolean-type expression it is paired with.
178
///
179
#[register_rule("Bubble", 6000, [UnsafePow])]
180
1272697
fn pow_to_bubble(expr: &Expression, _: &SymbolTable) -> ApplicationResult {
181
1272697
    if is_all_constant(expr) {
182
258646
        return Err(RuleNotApplicable);
183
1014051
    }
184
1014051
    if let Expression::UnsafePow(_, a, b) = expr.clone() {
185
        // bubble bottom up
186
222
        if !a.is_safe() || !b.is_safe() {
187
12
            return Err(RuleNotApplicable);
188
210
        }
189

            
190
210
        return Ok(Reduction::pure(Expression::Bubble(
191
210
            Metadata::new(),
192
210
            Moo::new(Expression::SafePow(Metadata::new(), a.clone(), b.clone())),
193
210
            Moo::new(Expression::And(
194
210
                Metadata::new(),
195
210
                Moo::new(matrix_expr![
196
210
                    Expression::Or(
197
210
                        Metadata::new(),
198
210
                        Moo::new(matrix_expr![
199
210
                            Expression::Neq(
200
210
                                Metadata::new(),
201
210
                                a,
202
210
                                Moo::new(Atom::Literal(Literal::Int(0)).into()),
203
210
                            ),
204
210
                            Expression::Neq(
205
210
                                Metadata::new(),
206
210
                                b.clone(),
207
210
                                Moo::new(Atom::Literal(Literal::Int(0)).into()),
208
210
                            ),
209
210
                        ]),
210
210
                    ),
211
210
                    Expression::Geq(
212
210
                        Metadata::new(),
213
210
                        b,
214
210
                        Moo::new(Atom::Literal(Literal::Int(0)).into()),
215
210
                    ),
216
210
                ]),
217
210
            )),
218
210
        )));
219
1013829
    }
220
1013829
    Err(ApplicationError::RuleNotApplicable)
221
1272697
}