1
use conjure_cp::ast::comprehension::ComprehensionQualifier;
2
use conjure_cp::ast::{Atom, Metadata};
3
use conjure_cp::ast::{Expression as Expr, Moo, SymbolTable};
4
use conjure_cp::into_matrix_expr;
5
use conjure_cp::rule_engine::Reduction;
6
use conjure_cp::rule_engine::{
7
    ApplicationError::RuleNotApplicable, ApplicationResult, register_rule,
8
};
9

            
10
use crate::utils::replace_expression_generator_source;
11

            
12
// [ return_expr | i <- A union B, qualifiers...] -> flatten([[ return_expr | i <- A, qualifiers...], [ return_expr | i <- B, !(i in A), qualifiers...]; int(1..2)])
13
#[register_rule("Base", 8700, [Comprehension])]
14
1965557
fn union_set(expr: &Expr, _: &SymbolTable) -> ApplicationResult {
15
1965557
    match expr {
16
16176
        Expr::Comprehension(_, comp) => {
17
            // find if any of the generators are generating from expressions
18
26951
            for qualifier in &comp.qualifiers {
19
26951
                if let ComprehensionQualifier::ExpressionGenerator { ptr } = qualifier {
20
74
                    let gen_decl = ptr.clone();
21

            
22
                    // match on expression being of form A union B
23
74
                    let Some((a, b)) = (match ptr.as_quantified_expr() {
24
74
                        Some(expr_guard) => match &*expr_guard {
25
10
                            Expr::Union(_, a, b) => Some((a.clone(), b.clone())),
26
64
                            _ => None,
27
                        },
28
                        None => None,
29
                    }) else {
30
64
                        continue;
31
                    };
32

            
33
                    // [ return_expr | i <- A, guards...] part
34
10
                    let (comprehension1, _) = replace_expression_generator_source(
35
10
                        comp.as_ref(),
36
10
                        &gen_decl,
37
10
                        a.clone().into(),
38
10
                    );
39

            
40
                    // [ return_expr | i <- B, !(i in A), guards...] part
41
10
                    let (mut comprehension2, b_ptr) =
42
10
                        replace_expression_generator_source(comp.as_ref(), &gen_decl, b.into());
43

            
44
                    // add the condition !(i in A)
45
10
                    comprehension2
46
10
                        .qualifiers
47
10
                        .push(ComprehensionQualifier::Condition(Expr::Not(
48
10
                            Metadata::new(),
49
10
                            Moo::new(Expr::In(
50
10
                                Metadata::new(),
51
10
                                Moo::new(Expr::Atomic(Metadata::new(), Atom::new_ref(b_ptr))),
52
10
                                a,
53
10
                            )),
54
10
                        )));
55

            
56
10
                    return Ok(Reduction::pure(Expr::Flatten(
57
10
                        Metadata::new(),
58
10
                        None,
59
10
                        Moo::new(into_matrix_expr!(vec![
60
10
                            Expr::Comprehension(Metadata::new(), comprehension1.into()),
61
10
                            Expr::Comprehension(Metadata::new(), comprehension2.into())
62
10
                        ])),
63
10
                    )));
64
26877
                }
65
            }
66

            
67
16166
            Err(RuleNotApplicable)
68
        }
69
1949381
        _ => Err(RuleNotApplicable),
70
    }
71
1965557
}