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
72216
fn index_tuple_to_atom(expr: &Expr, symbols: &SymbolTable) -> ApplicationResult {
21
    // i assume the MkOpIndexing is the same as matrix indexing
22
72216
    let Expr::SafeIndex(_, subject, indices) = expr else {
23
67368
        return Err(RuleNotApplicable);
24
    };
25

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

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

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

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

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

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

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

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

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

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

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

            
71
84
    Ok(Reduction::pure(subject))
72
72216
}
73

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
220
6
    Ok(Reduction::pure(new_expr))
221
72216
}
222

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

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

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

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

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

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

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

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

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

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

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

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

            
291
12
    Ok(Reduction::pure(new_expr))
292
72216
}
293

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
382
6
    Ok(Reduction::pure(new_expr))
383
72216
}