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

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

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

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

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

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

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

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

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

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

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

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

            
71
42
    Ok(Reduction::pure(subject))
72
24944
}
73

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
220
3
    Ok(Reduction::pure(new_expr))
221
24944
}
222

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

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

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

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

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

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

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

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

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

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

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

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

            
291
6
    Ok(Reduction::pure(new_expr))
292
24944
}
293

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
382
3
    Ok(Reduction::pure(new_expr))
383
24944
}