1
use conjure_cp::ast::Expression as Expr;
2
use conjure_cp::ast::Moo;
3
use conjure_cp::ast::SymbolTable;
4
use conjure_cp::ast::{AbstractLiteral, GroundDomain};
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
fn index_tuple_to_atom(expr: &Expr, symbols: &SymbolTable) -> ApplicationResult {
21
    // i assume the MkOpIndexing is the same as matrix indexing
22
    let Expr::SafeIndex(_, subject, indices) = expr else {
23
        return Err(RuleNotApplicable);
24
    };
25

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
238
    let Expr::AbstractLiteral(_, AbstractLiteral::Tuple(elems2)) = &**right else {
239
        return Err(RuleNotApplicable);
240
    };
241

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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