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, Uniplate};
13

            
14
use super::utils::is_all_constant;
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
2436378
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
2435610
        _ => Err(ApplicationError::RuleNotApplicable),
37
    }
38
2436378
}
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
2177189
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
2177189
    if matches!(expr, Expression::Root(_, _) | Expression::Bubble(_, _, _)) {
52
60312
        return Err(RuleNotApplicable);
53
2116877
    }
54

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

            
67
2106635
    let mut sub = expr.children();
68
2106635
    let mut bubbled_conditions = vec![];
69
2106635
    for e in sub.iter_mut() {
70
1789709
        if let Expression::Bubble(_, a, b) = e
71
882
            && a.return_type() != ReturnType::Bool
72
882
        {
73
882
            let a = Moo::unwrap_or_clone(Moo::clone(a));
74
882
            let b = Moo::unwrap_or_clone(Moo::clone(b));
75
882
            bubbled_conditions.push(b);
76
882
            *e = a;
77
1788827
        }
78
    }
79
2106635
    if bubbled_conditions.is_empty() {
80
2105753
        Err(ApplicationError::RuleNotApplicable)
81
882
    } else if bubbled_conditions.len() == 1 {
82
882
        let new_expr = Expression::Bubble(
83
882
            Metadata::new(),
84
882
            Moo::new(expr.with_children(sub)),
85
882
            Moo::new(bubbled_conditions[0].clone()),
86
882
        );
87

            
88
882
        Ok(Reduction::pure(new_expr))
89
    } else {
90
        Ok(Reduction::pure(Expression::Bubble(
91
            Metadata::new(),
92
            Moo::new(expr.with_children(sub)),
93
            Moo::new(Expression::And(
94
                Metadata::new(),
95
                Moo::new(into_matrix_expr![bubbled_conditions]),
96
            )),
97
        )))
98
    }
99
2177189
}
100

            
101
// Bubble applications
102

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

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

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

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

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

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

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