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

            
11
use conjure_cp::ast::Atom;
12
use conjure_cp::ast::Expression;
13
use conjure_cp::ast::Literal;
14
use conjure_cp::ast::Metadata;
15
use conjure_cp::ast::Name;
16
use conjure_cp::rule_engine::ApplicationError;
17

            
18
//TODO: largely copied from the matrix rules, This should be possible to simplify
19
#[register_rule(("Base", 2000))]
20
105669
fn index_tuple_to_atom(expr: &Expr, symbols: &SymbolTable) -> ApplicationResult {
21
    // i assume the MkOpIndexing is the same as matrix indexing
22
105669
    let Expr::SafeIndex(_, subject, indices) = expr else {
23
98397
        return Err(RuleNotApplicable);
24
    };
25

            
26
7272
    let Expr::Atomic(_, Atom::Reference(decl)) = &**subject else {
27
513
        return Err(RuleNotApplicable);
28
    };
29

            
30
6759
    let Name::WithRepresentation(name, reprs) = &decl.name() as &Name else {
31
6561
        return Err(RuleNotApplicable);
32
    };
33

            
34
198
    if reprs.first().is_none_or(|x| x.as_str() != "tuple_to_atom") {
35
72
        return Err(RuleNotApplicable);
36
126
    }
37

            
38
    // tuples are always one dimensional
39
126
    if indices.len() != 1 {
40
        return Err(RuleNotApplicable);
41
126
    }
42

            
43
126
    let repr = symbols
44
126
        .get_representation(name, &["tuple_to_atom"])
45
126
        .unwrap()[0]
46
126
        .clone();
47

            
48
    // let decl = symbols.lookup(name).unwrap();
49

            
50
126
    let Some(GroundDomain::Tuple(_)) = decl.resolved_domain().as_deref() else {
51
        return Err(RuleNotApplicable);
52
    };
53

            
54
126
    let mut indices_as_lit: Literal = Literal::Bool(false);
55

            
56
126
    for index in indices {
57
126
        let Some(index) = index.clone().into_literal() else {
58
            return Err(RuleNotApplicable); // we don't support non-literal indices
59
        };
60
126
        indices_as_lit = index;
61
    }
62

            
63
126
    let indices_as_name = Name::Represented(Box::new((
64
126
        name.as_ref().clone(),
65
126
        "tuple_to_atom".into(),
66
126
        indices_as_lit.into(),
67
126
    )));
68

            
69
126
    let subject = repr.expression_down(symbols)?[&indices_as_name].clone();
70

            
71
126
    Ok(Reduction::pure(subject))
72
105669
}
73

            
74
#[register_rule(("Bubble", 8000))]
75
872982
fn tuple_index_to_bubble(expr: &Expr, _: &SymbolTable) -> ApplicationResult {
76
872982
    let Expr::UnsafeIndex(_, subject, indices) = expr else {
77
856863
        return Err(RuleNotApplicable);
78
    };
79

            
80
16119
    let Expr::Atomic(_, Atom::Reference(decl)) = &**subject else {
81
36
        return Err(RuleNotApplicable);
82
    };
83

            
84
16083
    let Name::WithRepresentation(_, reprs) = &decl.name() as &Name else {
85
2988
        return Err(RuleNotApplicable);
86
    };
87

            
88
13095
    if reprs.first().is_none_or(|x| x.as_str() != "tuple_to_atom") {
89
13077
        return Err(RuleNotApplicable);
90
18
    }
91

            
92
18
    let domain = subject.domain_of().ok_or(ApplicationError::DomainError)?;
93

            
94
18
    let Some(elems) = domain.as_tuple() else {
95
        return Err(RuleNotApplicable);
96
    };
97

            
98
18
    assert_eq!(indices.len(), 1, "tuple indexing is always one dimensional");
99
18
    let index = indices[0].clone();
100

            
101
18
    let bubble_constraint = Moo::new(Expression::And(
102
18
        Metadata::new(),
103
18
        Moo::new(matrix_expr![
104
18
            Expression::Leq(
105
18
                Metadata::new(),
106
18
                Moo::new(index.clone()),
107
18
                Moo::new(Expression::Atomic(
108
18
                    Metadata::new(),
109
18
                    Atom::Literal(Literal::Int(elems.len() as i32))
110
18
                ))
111
18
            ),
112
18
            Expression::Geq(
113
18
                Metadata::new(),
114
18
                Moo::new(index),
115
18
                Moo::new(Expression::Atomic(
116
18
                    Metadata::new(),
117
18
                    Atom::Literal(Literal::Int(1))
118
18
                ))
119
18
            )
120
18
        ]),
121
18
    ));
122

            
123
18
    let new_expr = Moo::new(Expression::SafeIndex(
124
18
        Metadata::new(),
125
18
        subject.clone(),
126
18
        indices.clone(),
127
18
    ));
128

            
129
18
    Ok(Reduction::pure(Expression::Bubble(
130
18
        Metadata::new(),
131
18
        new_expr,
132
18
        bubble_constraint,
133
18
    )))
134
872982
}
135

            
136
// convert equality to tuple equality
137
#[register_rule(("Base", 2000))]
138
105669
fn tuple_equality(expr: &Expr, _: &SymbolTable) -> ApplicationResult {
139
105669
    let Expr::Eq(_, left, right) = expr else {
140
100200
        return Err(RuleNotApplicable);
141
    };
142

            
143
5469
    let Expr::Atomic(_, Atom::Reference(decl)) = &**left else {
144
2721
        return Err(RuleNotApplicable);
145
    };
146

            
147
2748
    let Name::WithRepresentation(_, reprs) = &decl.name() as &Name else {
148
2703
        return Err(RuleNotApplicable);
149
    };
150

            
151
45
    let Expr::Atomic(_, Atom::Reference(decl2)) = &**right else {
152
27
        return Err(RuleNotApplicable);
153
    };
154

            
155
18
    let Name::WithRepresentation(_, reprs2) = &decl2.name() as &Name else {
156
        return Err(RuleNotApplicable);
157
    };
158

            
159
18
    if reprs.first().is_none_or(|x| x.as_str() != "tuple_to_atom") {
160
9
        return Err(RuleNotApplicable);
161
9
    }
162

            
163
9
    if reprs2.first().is_none_or(|x| x.as_str() != "tuple_to_atom") {
164
        return Err(RuleNotApplicable);
165
9
    }
166

            
167
    // let decl = symbols.lookup(name).unwrap();
168
    // let decl2 = symbols.lookup(name2).unwrap();
169

            
170
9
    let domain = decl
171
9
        .resolved_domain()
172
9
        .ok_or(ApplicationError::DomainError)?;
173
9
    let domain2 = decl2
174
9
        .resolved_domain()
175
9
        .ok_or(ApplicationError::DomainError)?;
176

            
177
9
    let GroundDomain::Tuple(elems) = domain.as_ref() else {
178
        return Err(RuleNotApplicable);
179
    };
180

            
181
9
    let GroundDomain::Tuple(elems2) = domain2.as_ref() else {
182
        return Err(RuleNotApplicable);
183
    };
184

            
185
9
    if elems.len() != elems2.len() {
186
        return Err(RuleNotApplicable);
187
9
    }
188

            
189
9
    let mut equality_constraints = vec![];
190
18
    for i in 0..elems.len() {
191
18
        let left_elem = Expression::SafeIndex(
192
18
            Metadata::new(),
193
18
            Moo::clone(left),
194
18
            vec![Expression::Atomic(
195
18
                Metadata::new(),
196
18
                Atom::Literal(Literal::Int((i + 1) as i32)),
197
18
            )],
198
18
        );
199
18
        let right_elem = Expression::SafeIndex(
200
18
            Metadata::new(),
201
18
            Moo::clone(right),
202
18
            vec![Expression::Atomic(
203
18
                Metadata::new(),
204
18
                Atom::Literal(Literal::Int((i + 1) as i32)),
205
18
            )],
206
18
        );
207
18

            
208
18
        equality_constraints.push(Expression::Eq(
209
18
            Metadata::new(),
210
18
            Moo::new(left_elem),
211
18
            Moo::new(right_elem),
212
18
        ));
213
18
    }
214

            
215
9
    let new_expr = Expression::And(
216
9
        Metadata::new(),
217
9
        Moo::new(into_matrix_expr!(equality_constraints)),
218
9
    );
219

            
220
9
    Ok(Reduction::pure(new_expr))
221
105669
}
222

            
223
//tuple equality where the left is a variable and the right is a tuple literal
224
#[register_rule(("Base", 2000))]
225
105669
fn tuple_to_constant(expr: &Expr, symbols: &SymbolTable) -> ApplicationResult {
226
105669
    let Expr::Eq(_, left, right) = expr else {
227
100200
        return Err(RuleNotApplicable);
228
    };
229

            
230
5469
    let Expr::Atomic(_, Atom::Reference(decl)) = &**left else {
231
2721
        return Err(RuleNotApplicable);
232
    };
233

            
234
2748
    let Name::WithRepresentation(name, reprs) = &decl.name() as &Name else {
235
2703
        return Err(RuleNotApplicable);
236
    };
237

            
238
45
    let Some(rhs_tuple_len) = crate::utils::constant_tuple_len(right.as_ref()) else {
239
27
        return Err(RuleNotApplicable);
240
    };
241

            
242
18
    if reprs.first().is_none_or(|x| x.as_str() != "tuple_to_atom") {
243
        return Err(RuleNotApplicable);
244
18
    }
245

            
246
18
    let decl = symbols.lookup(name).unwrap();
247

            
248
18
    let domain = decl
249
18
        .resolved_domain()
250
18
        .ok_or(ApplicationError::DomainError)?;
251

            
252
18
    let GroundDomain::Tuple(elems) = domain.as_ref() else {
253
        return Err(RuleNotApplicable);
254
    };
255

            
256
18
    if elems.len() != rhs_tuple_len {
257
        return Err(RuleNotApplicable);
258
18
    }
259

            
260
18
    let mut equality_constraints = vec![];
261
36
    for i in 0..elems.len() {
262
36
        let left_elem = Expression::SafeIndex(
263
36
            Metadata::new(),
264
36
            Moo::clone(left),
265
36
            vec![Expression::Atomic(
266
36
                Metadata::new(),
267
36
                Atom::Literal(Literal::Int((i + 1) as i32)),
268
36
            )],
269
36
        );
270
36
        let right_elem = Expression::SafeIndex(
271
36
            Metadata::new(),
272
36
            Moo::clone(right),
273
36
            vec![Expression::Atomic(
274
36
                Metadata::new(),
275
36
                Atom::Literal(Literal::Int((i + 1) as i32)),
276
36
            )],
277
36
        );
278
36

            
279
36
        equality_constraints.push(Expression::Eq(
280
36
            Metadata::new(),
281
36
            Moo::new(left_elem),
282
36
            Moo::new(right_elem),
283
36
        ));
284
36
    }
285

            
286
18
    let new_expr = Expression::And(
287
18
        Metadata::new(),
288
18
        Moo::new(into_matrix_expr!(equality_constraints)),
289
18
    );
290

            
291
18
    Ok(Reduction::pure(new_expr))
292
105669
}
293

            
294
// convert equality to tuple inequality
295
#[register_rule(("Base", 2000))]
296
105669
fn tuple_inequality(expr: &Expr, _: &SymbolTable) -> ApplicationResult {
297
105669
    let Expr::Neq(_, left, right) = expr else {
298
103194
        return Err(RuleNotApplicable);
299
    };
300

            
301
2475
    let Expr::Atomic(_, Atom::Reference(decl)) = &**left else {
302
1449
        return Err(RuleNotApplicable);
303
    };
304

            
305
1026
    let Name::WithRepresentation(_, reprs) = &decl.name() as &Name else {
306
1017
        return Err(RuleNotApplicable);
307
    };
308

            
309
9
    let Expr::Atomic(_, Atom::Reference(decl2)) = &**right else {
310
        return Err(RuleNotApplicable);
311
    };
312

            
313
9
    let Name::WithRepresentation(_, reprs2) = &decl2.name() as &Name else {
314
        return Err(RuleNotApplicable);
315
    };
316

            
317
9
    if reprs.first().is_none_or(|x| x.as_str() != "tuple_to_atom") {
318
        return Err(RuleNotApplicable);
319
9
    }
320

            
321
9
    if reprs2.first().is_none_or(|x| x.as_str() != "tuple_to_atom") {
322
        return Err(RuleNotApplicable);
323
9
    }
324

            
325
9
    let domain = decl
326
9
        .resolved_domain()
327
9
        .ok_or(ApplicationError::DomainError)?;
328

            
329
9
    let domain2 = decl2
330
9
        .resolved_domain()
331
9
        .ok_or(ApplicationError::DomainError)?;
332

            
333
9
    let GroundDomain::Tuple(elems) = domain.as_ref() else {
334
        return Err(RuleNotApplicable);
335
    };
336

            
337
9
    let GroundDomain::Tuple(elems2) = domain2.as_ref() else {
338
        return Err(RuleNotApplicable);
339
    };
340

            
341
9
    assert_eq!(
342
9
        elems.len(),
343
9
        elems2.len(),
344
        "tuple inequality requires same length domains"
345
    );
346

            
347
9
    let mut equality_constraints = vec![];
348
18
    for i in 0..elems.len() {
349
18
        let left_elem = Expression::SafeIndex(
350
18
            Metadata::new(),
351
18
            Moo::clone(left),
352
18
            vec![Expression::Atomic(
353
18
                Metadata::new(),
354
18
                Atom::Literal(Literal::Int((i + 1) as i32)),
355
18
            )],
356
18
        );
357
18
        let right_elem = Expression::SafeIndex(
358
18
            Metadata::new(),
359
18
            Moo::clone(right),
360
18
            vec![Expression::Atomic(
361
18
                Metadata::new(),
362
18
                Atom::Literal(Literal::Int((i + 1) as i32)),
363
18
            )],
364
18
        );
365
18

            
366
18
        equality_constraints.push(Expression::Eq(
367
18
            Metadata::new(),
368
18
            Moo::new(left_elem),
369
18
            Moo::new(right_elem),
370
18
        ));
371
18
    }
372

            
373
    // Just copied from Conjure, would it be better to DeMorgan this?
374
9
    let new_expr = Expression::Not(
375
9
        Metadata::new(),
376
9
        Moo::new(Expression::And(
377
9
            Metadata::new(),
378
9
            Moo::new(into_matrix_expr!(equality_constraints)),
379
9
        )),
380
9
    );
381

            
382
9
    Ok(Reduction::pure(new_expr))
383
105669
}