1
#![allow(dead_code)]
2
use crate::ast::{AbstractLiteral, Atom, Expression as Expr, Literal as Lit, Metadata, matrix};
3
use crate::into_matrix;
4
use itertools::{Itertools as _, izip};
5
use std::cmp::Ordering as CmpOrdering;
6
use std::collections::HashSet;
7

            
8
/// Simplify an expression to a constant if possible
9
/// Returns:
10
/// `None` if the expression cannot be simplified to a constant (e.g. if it contains a variable)
11
/// `Some(Const)` if the expression can be simplified to a constant
12
1911080
pub fn eval_constant(expr: &Expr) -> Option<Lit> {
13
857980
    match expr {
14
160
        Expr::Supset(_, a, b) => match (a.as_ref(), b.as_ref()) {
15
            (
16
160
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
17
160
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
18
            ) => {
19
160
                let a_set: HashSet<Lit> = a.iter().cloned().collect();
20
160
                let b_set: HashSet<Lit> = b.iter().cloned().collect();
21

            
22
160
                if a_set.difference(&b_set).count() > 0 {
23
120
                    Some(Lit::Bool(a_set.is_superset(&b_set)))
24
                } else {
25
40
                    Some(Lit::Bool(false))
26
                }
27
            }
28
            _ => None,
29
        },
30
160
        Expr::SupsetEq(_, a, b) => match (a.as_ref(), b.as_ref()) {
31
            (
32
160
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
33
160
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
34
160
            ) => Some(Lit::Bool(
35
160
                a.iter()
36
160
                    .cloned()
37
160
                    .collect::<HashSet<Lit>>()
38
160
                    .is_superset(&b.iter().cloned().collect::<HashSet<Lit>>()),
39
160
            )),
40
            _ => None,
41
        },
42
200
        Expr::Subset(_, a, b) => match (a.as_ref(), b.as_ref()) {
43
            (
44
200
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
45
200
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
46
            ) => {
47
200
                let a_set: HashSet<Lit> = a.iter().cloned().collect();
48
200
                let b_set: HashSet<Lit> = b.iter().cloned().collect();
49

            
50
200
                if b_set.difference(&a_set).count() > 0 {
51
160
                    Some(Lit::Bool(a_set.is_subset(&b_set)))
52
                } else {
53
40
                    Some(Lit::Bool(false))
54
                }
55
            }
56
            _ => None,
57
        },
58
400
        Expr::SubsetEq(_, a, b) => match (a.as_ref(), b.as_ref()) {
59
            (
60
400
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
61
400
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
62
400
            ) => Some(Lit::Bool(
63
400
                a.iter()
64
400
                    .cloned()
65
400
                    .collect::<HashSet<Lit>>()
66
400
                    .is_subset(&b.iter().cloned().collect::<HashSet<Lit>>()),
67
400
            )),
68
            _ => None,
69
        },
70
120
        Expr::Intersect(_, a, b) => match (a.as_ref(), b.as_ref()) {
71
            (
72
120
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
73
120
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
74
            ) => {
75
120
                let mut res: Vec<Lit> = Vec::new();
76
280
                for lit in a.iter() {
77
280
                    if b.contains(lit) && !res.contains(lit) {
78
200
                        res.push(lit.clone());
79
200
                    }
80
                }
81
120
                Some(Lit::AbstractLiteral(AbstractLiteral::Set(res)))
82
            }
83
            _ => None,
84
        },
85
120
        Expr::Union(_, a, b) => match (a.as_ref(), b.as_ref()) {
86
            (
87
120
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
88
120
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
89
            ) => {
90
120
                let mut res: Vec<Lit> = Vec::new();
91
320
                for lit in a.iter() {
92
320
                    res.push(lit.clone());
93
320
                }
94
320
                for lit in b.iter() {
95
320
                    if !res.contains(lit) {
96
240
                        res.push(lit.clone());
97
240
                    }
98
                }
99
120
                Some(Lit::AbstractLiteral(AbstractLiteral::Set(res)))
100
            }
101
            _ => None,
102
        },
103
1140
        Expr::In(_, a, b) => {
104
            if let (
105
40
                Expr::Atomic(_, Atom::Literal(Lit::Int(c))),
106
40
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(d)))),
107
1140
            ) = (a.as_ref(), b.as_ref())
108
            {
109
120
                for lit in d.iter() {
110
120
                    if let Lit::Int(x) = lit
111
120
                        && c == x
112
                    {
113
40
                        return Some(Lit::Bool(true));
114
80
                    }
115
                }
116
                Some(Lit::Bool(false))
117
            } else {
118
1100
                None
119
            }
120
        }
121
        Expr::FromSolution(_, _) => None,
122
        Expr::DominanceRelation(_, _) => None,
123
9560
        Expr::InDomain(_, e, domain) => {
124
9560
            let Expr::Atomic(_, Atom::Literal(lit)) = e.as_ref() else {
125
6400
                return None;
126
            };
127

            
128
3160
            domain.contains(lit).ok().map(Into::into)
129
        }
130
56780
        Expr::Atomic(_, Atom::Literal(c)) => Some(c.clone()),
131
801200
        Expr::Atomic(_, Atom::Reference(reference)) => reference.resolve_constant(),
132
111040
        Expr::AbstractLiteral(_, a) => Some(Lit::AbstractLiteral(a.clone().into_literals()?)),
133
8940
        Expr::Comprehension(_, _) => None,
134
40
        Expr::AbstractComprehension(_, _) => None,
135
225200
        Expr::UnsafeIndex(_, subject, indices) | Expr::SafeIndex(_, subject, indices) => {
136
282160
            let subject: Lit = eval_constant(subject.as_ref())?;
137
2040
            let indices: Vec<Lit> = indices
138
2040
                .iter()
139
2040
                .map(eval_constant)
140
2040
                .collect::<Option<Vec<Lit>>>()?;
141

            
142
380
            match subject {
143
240
                Lit::AbstractLiteral(subject @ AbstractLiteral::Matrix(_, _)) => {
144
240
                    matrix::flatten_enumerate(subject)
145
840
                        .find(|(i, _)| i == &indices)
146
240
                        .map(|(_, x)| x)
147
                }
148
100
                Lit::AbstractLiteral(subject @ AbstractLiteral::Tuple(_)) => {
149
100
                    let AbstractLiteral::Tuple(elems) = subject else {
150
                        return None;
151
                    };
152

            
153
100
                    assert!(indices.len() == 1, "nested tuples not supported yet");
154

            
155
100
                    let Lit::Int(index) = indices[0].clone() else {
156
                        return None;
157
                    };
158

            
159
100
                    if elems.len() < index as usize || index < 1 {
160
                        return None;
161
100
                    }
162

            
163
                    // -1 for 0-indexing vs 1-indexing
164
100
                    let item = elems[index as usize - 1].clone();
165

            
166
100
                    Some(item)
167
                }
168
40
                Lit::AbstractLiteral(subject @ AbstractLiteral::Record(_)) => {
169
40
                    let AbstractLiteral::Record(elems) = subject else {
170
                        return None;
171
                    };
172

            
173
40
                    assert!(indices.len() == 1, "nested record not supported yet");
174

            
175
40
                    let Lit::Int(index) = indices[0].clone() else {
176
                        return None;
177
                    };
178

            
179
40
                    if elems.len() < index as usize || index < 1 {
180
                        return None;
181
40
                    }
182

            
183
                    // -1 for 0-indexing vs 1-indexing
184
40
                    let item = elems[index as usize - 1].clone();
185
40
                    Some(item.value)
186
                }
187
                _ => None,
188
            }
189
        }
190
10020
        Expr::UnsafeSlice(_, subject, indices) | Expr::SafeSlice(_, subject, indices) => {
191
13420
            let subject: Lit = eval_constant(subject.as_ref())?;
192
20
            let Lit::AbstractLiteral(subject @ AbstractLiteral::Matrix(_, _)) = subject else {
193
                return None;
194
            };
195

            
196
20
            let hole_dim = indices
197
20
                .iter()
198
20
                .cloned()
199
40
                .position(|x| x.is_none())
200
20
                .expect("slice expression should have a hole dimension");
201

            
202
20
            let missing_domain = matrix::index_domains(subject.clone())[hole_dim].clone();
203

            
204
20
            let indices: Vec<Option<Lit>> = indices
205
20
                .iter()
206
20
                .cloned()
207
40
                .map(|x| {
208
                    // the outer option represents success of this iterator, the inner the index
209
                    // slice.
210
40
                    match x {
211
20
                        Some(x) => eval_constant(&x).map(Some),
212
20
                        None => Some(None),
213
                    }
214
40
                })
215
20
                .collect::<Option<Vec<Option<Lit>>>>()?;
216

            
217
20
            let indices_in_slice: Vec<Vec<Lit>> = missing_domain
218
20
                .values()
219
20
                .ok()?
220
60
                .map(|i| {
221
60
                    let mut indices = indices.clone();
222
60
                    indices[hole_dim] = Some(i);
223
                    // These unwraps will only fail if we have multiple holes.
224
                    // As this is invalid, panicking is fine.
225
120
                    indices.into_iter().map(|x| x.unwrap()).collect_vec()
226
60
                })
227
20
                .collect_vec();
228

            
229
            // Note: indices_in_slice is not necessarily sorted, so this is the best way.
230
20
            let elems = matrix::flatten_enumerate(subject)
231
180
                .filter(|(i, _)| indices_in_slice.contains(i))
232
20
                .map(|(_, elem)| elem)
233
20
                .collect();
234

            
235
20
            Some(Lit::AbstractLiteral(into_matrix![elems]))
236
        }
237
3400
        Expr::Abs(_, e) => un_op::<i32, i32>(|a| a.abs(), e).map(Lit::Int),
238
110980
        Expr::Eq(_, a, b) => bin_op::<i32, bool>(|a, b| a == b, a, b)
239
110980
            .or_else(|| bin_op::<bool, bool>(|a, b| a == b, a, b))
240
110980
            .map(Lit::Bool),
241
23680
        Expr::Neq(_, a, b) => bin_op::<i32, bool>(|a, b| a != b, a, b).map(Lit::Bool),
242
31080
        Expr::Lt(_, a, b) => bin_op::<i32, bool>(|a, b| a < b, a, b).map(Lit::Bool),
243
760
        Expr::Gt(_, a, b) => bin_op::<i32, bool>(|a, b| a > b, a, b).map(Lit::Bool),
244
10660
        Expr::Leq(_, a, b) => bin_op::<i32, bool>(|a, b| a <= b, a, b).map(Lit::Bool),
245
4600
        Expr::Geq(_, a, b) => bin_op::<i32, bool>(|a, b| a >= b, a, b).map(Lit::Bool),
246
4700
        Expr::Not(_, expr) => un_op::<bool, bool>(|e| !e, expr).map(Lit::Bool),
247
29240
        Expr::And(_, e) => {
248
29240
            vec_lit_op::<bool, bool>(|e| e.iter().all(|&e| e), e.as_ref()).map(Lit::Bool)
249
        }
250
69420
        Expr::Root(_, _) => None,
251
25560
        Expr::Or(_, es) => {
252
            // possibly cheating; definitely should be in partial eval instead
253
60760
            for e in (**es).clone().unwrap_list()? {
254
620
                if let Expr::Atomic(_, Atom::Literal(Lit::Bool(true))) = e {
255
420
                    return Some(Lit::Bool(true));
256
60340
                };
257
            }
258

            
259
21500
            vec_lit_op::<bool, bool>(|e| e.iter().any(|&e| e), es.as_ref()).map(Lit::Bool)
260
        }
261
6500
        Expr::Imply(_, box1, box2) => {
262
6500
            let a: &Atom = (&**box1).try_into().ok()?;
263
1940
            let b: &Atom = (&**box2).try_into().ok()?;
264

            
265
1060
            let a: bool = a.try_into().ok()?;
266
            let b: bool = b.try_into().ok()?;
267

            
268
            if a {
269
                // true -> b ~> b
270
                Some(Lit::Bool(b))
271
            } else {
272
                // false -> b ~> true
273
                Some(Lit::Bool(true))
274
            }
275
        }
276
980
        Expr::Iff(_, box1, box2) => {
277
980
            let a: &Atom = (&**box1).try_into().ok()?;
278
820
            let b: &Atom = (&**box2).try_into().ok()?;
279

            
280
60
            let a: bool = a.try_into().ok()?;
281
20
            let b: bool = b.try_into().ok()?;
282

            
283
            Some(Lit::Bool(a == b))
284
        }
285
71380
        Expr::Sum(_, exprs) => vec_lit_op::<i32, i32>(|e| e.iter().sum(), exprs).map(Lit::Int),
286
18180
        Expr::Product(_, exprs) => {
287
18180
            vec_lit_op::<i32, i32>(|e| e.iter().product(), exprs).map(Lit::Int)
288
        }
289
21100
        Expr::FlatIneq(_, a, b, c) => {
290
21100
            let a: i32 = a.try_into().ok()?;
291
3580
            let b: i32 = b.try_into().ok()?;
292
180
            let c: i32 = c.try_into().ok()?;
293

            
294
180
            Some(Lit::Bool(a <= b + c))
295
        }
296
21920
        Expr::FlatSumGeq(_, exprs, a) => {
297
38980
            let sum = exprs.iter().try_fold(0, |acc, atom: &Atom| {
298
38980
                let n: i32 = atom.try_into().ok()?;
299
18340
                let acc = acc + n;
300
18340
                Some(acc)
301
38980
            })?;
302

            
303
1280
            Some(Lit::Bool(sum >= a.try_into().ok()?))
304
        }
305
24900
        Expr::FlatSumLeq(_, exprs, a) => {
306
42700
            let sum = exprs.iter().try_fold(0, |acc, atom: &Atom| {
307
42700
                let n: i32 = atom.try_into().ok()?;
308
17820
                let acc = acc + n;
309
17820
                Some(acc)
310
42700
            })?;
311

            
312
20
            Some(Lit::Bool(sum >= a.try_into().ok()?))
313
        }
314
1740
        Expr::Min(_, e) => {
315
1740
            opt_vec_lit_op::<i32, i32>(|e| e.iter().min().copied(), e.as_ref()).map(Lit::Int)
316
        }
317
880
        Expr::Max(_, e) => {
318
880
            opt_vec_lit_op::<i32, i32>(|e| e.iter().max().copied(), e.as_ref()).map(Lit::Int)
319
        }
320
15620
        Expr::UnsafeDiv(_, a, b) | Expr::SafeDiv(_, a, b) => {
321
19860
            if unwrap_expr::<i32>(b)? == 0 {
322
40
                return None;
323
3500
            }
324
3500
            bin_op::<i32, i32>(|a, b| ((a as f32) / (b as f32)).floor() as i32, a, b).map(Lit::Int)
325
        }
326
8940
        Expr::UnsafeMod(_, a, b) | Expr::SafeMod(_, a, b) => {
327
11420
            if unwrap_expr::<i32>(b)? == 0 {
328
                return None;
329
660
            }
330
660
            bin_op::<i32, i32>(|a, b| a - b * (a as f32 / b as f32).floor() as i32, a, b)
331
660
                .map(Lit::Int)
332
        }
333
2240
        Expr::MinionDivEqUndefZero(_, a, b, c) => {
334
            // div always rounds down
335
2240
            let a: i32 = a.try_into().ok()?;
336
40
            let b: i32 = b.try_into().ok()?;
337
            let c: i32 = c.try_into().ok()?;
338

            
339
            if b == 0 {
340
                return None;
341
            }
342

            
343
            let a = a as f32;
344
            let b = b as f32;
345
            let div: i32 = (a / b).floor() as i32;
346
            Some(Lit::Bool(div == c))
347
        }
348
10600
        Expr::Bubble(_, a, b) => bin_op::<bool, bool>(|a, b| a && b, a, b).map(Lit::Bool),
349
9180
        Expr::MinionReify(_, a, b) => {
350
9180
            let result = eval_constant(a)?;
351

            
352
3200
            let result: bool = result.try_into().ok()?;
353
3200
            let b: bool = b.try_into().ok()?;
354

            
355
            Some(Lit::Bool(b == result))
356
        }
357
3140
        Expr::MinionReifyImply(_, a, b) => {
358
3140
            let result = eval_constant(a)?;
359

            
360
            let result: bool = result.try_into().ok()?;
361
            let b: bool = b.try_into().ok()?;
362

            
363
            if b {
364
                Some(Lit::Bool(result))
365
            } else {
366
                Some(Lit::Bool(true))
367
            }
368
        }
369
720
        Expr::MinionModuloEqUndefZero(_, a, b, c) => {
370
            // From Savile Row. Same semantics as division.
371
            //
372
            //   a - (b * floor(a/b))
373
            //
374
            // We don't use % as it has the same semantics as /. We don't use / as we want to round
375
            // down instead, not towards zero.
376

            
377
720
            let a: i32 = a.try_into().ok()?;
378
40
            let b: i32 = b.try_into().ok()?;
379
            let c: i32 = c.try_into().ok()?;
380

            
381
            if b == 0 {
382
                return None;
383
            }
384

            
385
            let modulo = a - b * (a as f32 / b as f32).floor() as i32;
386
            Some(Lit::Bool(modulo == c))
387
        }
388
1620
        Expr::MinionPow(_, a, b, c) => {
389
            // only available for positive a b c
390

            
391
1620
            let a: i32 = a.try_into().ok()?;
392
            let b: i32 = b.try_into().ok()?;
393
            let c: i32 = c.try_into().ok()?;
394

            
395
            if a <= 0 {
396
                return None;
397
            }
398

            
399
            if b <= 0 {
400
                return None;
401
            }
402

            
403
            if c <= 0 {
404
                return None;
405
            }
406

            
407
            Some(Lit::Bool(a ^ b == c))
408
        }
409
160
        Expr::MinionWInSet(_, _, _) => None,
410
380
        Expr::MinionWInIntervalSet(_, x, intervals) => {
411
380
            let x_lit: &Lit = x.try_into().ok()?;
412

            
413
            let x_lit = match x_lit.clone() {
414
                Lit::Int(i) => Some(i),
415
                Lit::Bool(true) => Some(1),
416
                Lit::Bool(false) => Some(0),
417
                _ => None,
418
            }?;
419

            
420
            let mut intervals = intervals.iter();
421
            loop {
422
                let Some(lower) = intervals.next() else {
423
                    break;
424
                };
425

            
426
                let Some(upper) = intervals.next() else {
427
                    break;
428
                };
429
                if &x_lit >= lower && &x_lit <= upper {
430
                    return Some(Lit::Bool(true));
431
                }
432
            }
433

            
434
            Some(Lit::Bool(false))
435
        }
436
        Expr::Flatten(_, _, _) => {
437
            // TODO
438
2860
            None
439
        }
440
19700
        Expr::AllDiff(_, e) => {
441
19700
            let es = (**e).clone().unwrap_list()?;
442
1660
            let mut lits: HashSet<Lit> = HashSet::new();
443
1780
            for expr in es {
444
820
                let Expr::Atomic(_, Atom::Literal(x)) = expr else {
445
1620
                    return None;
446
                };
447
160
                match x {
448
                    Lit::Int(_) | Lit::Bool(_) => {
449
160
                        if lits.contains(&x) {
450
                            return Some(Lit::Bool(false));
451
160
                        } else {
452
160
                            lits.insert(x.clone());
453
160
                        }
454
                    }
455
                    Lit::AbstractLiteral(_) => return None, // Reject AbstractLiteral cases
456
                }
457
            }
458
40
            Some(Lit::Bool(true))
459
        }
460
4060
        Expr::FlatAllDiff(_, es) => {
461
4060
            let mut lits: HashSet<Lit> = HashSet::new();
462
4060
            for atom in es {
463
4060
                let Atom::Literal(x) = atom else {
464
4060
                    return None;
465
                };
466

            
467
                match x {
468
                    Lit::Int(_) | Lit::Bool(_) => {
469
                        if lits.contains(x) {
470
                            return Some(Lit::Bool(false));
471
                        } else {
472
                            lits.insert(x.clone());
473
                        }
474
                    }
475
                    Lit::AbstractLiteral(_) => return None, // Reject AbstractLiteral cases
476
                }
477
            }
478
            Some(Lit::Bool(true))
479
        }
480
1540
        Expr::FlatWatchedLiteral(_, _, _) => None,
481
12240
        Expr::AuxDeclaration(_, _, _) => None,
482
13540
        Expr::Neg(_, a) => {
483
13540
            let a: &Atom = a.try_into().ok()?;
484
7360
            let a: i32 = a.try_into().ok()?;
485
2660
            Some(Lit::Int(-a))
486
        }
487
2200
        Expr::Minus(_, a, b) => {
488
2200
            let a: &Atom = a.try_into().ok()?;
489
1800
            let a: i32 = a.try_into().ok()?;
490

            
491
260
            let b: &Atom = b.try_into().ok()?;
492
260
            let b: i32 = b.try_into().ok()?;
493

            
494
140
            Some(Lit::Int(a - b))
495
        }
496
1200
        Expr::FlatMinusEq(_, a, b) => {
497
1200
            let a: i32 = a.try_into().ok()?;
498
960
            let b: i32 = b.try_into().ok()?;
499
            Some(Lit::Bool(a == -b))
500
        }
501
400
        Expr::FlatProductEq(_, a, b, c) => {
502
400
            let a: i32 = a.try_into().ok()?;
503
60
            let b: i32 = b.try_into().ok()?;
504
60
            let c: i32 = c.try_into().ok()?;
505
60
            Some(Lit::Bool(a * b == c))
506
        }
507
2500
        Expr::FlatWeightedSumLeq(_, cs, vs, total) => {
508
2500
            let cs: Vec<i32> = cs
509
2500
                .iter()
510
6600
                .map(|x| TryInto::<i32>::try_into(x).ok())
511
2500
                .collect::<Option<Vec<i32>>>()?;
512
2500
            let vs: Vec<i32> = vs
513
2500
                .iter()
514
4060
                .map(|x| TryInto::<i32>::try_into(x).ok())
515
2500
                .collect::<Option<Vec<i32>>>()?;
516
            let total: i32 = total.try_into().ok()?;
517

            
518
            let sum: i32 = izip!(cs, vs).fold(0, |acc, (c, v)| acc + (c * v));
519

            
520
            Some(Lit::Bool(sum <= total))
521
        }
522
2060
        Expr::FlatWeightedSumGeq(_, cs, vs, total) => {
523
2060
            let cs: Vec<i32> = cs
524
2060
                .iter()
525
5260
                .map(|x| TryInto::<i32>::try_into(x).ok())
526
2060
                .collect::<Option<Vec<i32>>>()?;
527
2060
            let vs: Vec<i32> = vs
528
2060
                .iter()
529
3440
                .map(|x| TryInto::<i32>::try_into(x).ok())
530
2060
                .collect::<Option<Vec<i32>>>()?;
531
            let total: i32 = total.try_into().ok()?;
532

            
533
            let sum: i32 = izip!(cs, vs).fold(0, |acc, (c, v)| acc + (c * v));
534

            
535
            Some(Lit::Bool(sum >= total))
536
        }
537
420
        Expr::FlatAbsEq(_, x, y) => {
538
420
            let x: i32 = x.try_into().ok()?;
539
20
            let y: i32 = y.try_into().ok()?;
540

            
541
            Some(Lit::Bool(x == y.abs()))
542
        }
543
11540
        Expr::UnsafePow(_, a, b) | Expr::SafePow(_, a, b) => {
544
14840
            let a: &Atom = a.try_into().ok()?;
545
14180
            let a: i32 = a.try_into().ok()?;
546

            
547
360
            let b: &Atom = b.try_into().ok()?;
548
360
            let b: i32 = b.try_into().ok()?;
549

            
550
360
            if (a != 0 || b != 0) && b >= 0 {
551
360
                Some(Lit::Int(a.pow(b as u32)))
552
            } else {
553
                None
554
            }
555
        }
556
        Expr::Scope(_, _) => None,
557
        Expr::Metavar(_, _) => None,
558
3180
        Expr::MinionElementOne(_, _, _, _) => None,
559
520
        Expr::ToInt(_, expression) => {
560
520
            let lit = eval_constant(expression.as_ref())?;
561
            match lit {
562
                Lit::Int(_) => Some(lit),
563
                Lit::Bool(true) => Some(Lit::Int(1)),
564
                Lit::Bool(false) => Some(Lit::Int(0)),
565
                _ => None,
566
            }
567
        }
568
        Expr::SATInt(..) => None,
569
        Expr::PairwiseSum(_, a, b) => {
570
            match (eval_constant(a.as_ref())?, eval_constant(b.as_ref())?) {
571
                (Lit::Int(a_int), Lit::Int(b_int)) => Some(Lit::Int(a_int + b_int)),
572
                _ => None,
573
            }
574
        }
575
        Expr::PairwiseProduct(_, a, b) => {
576
            match (eval_constant(a.as_ref())?, eval_constant(b.as_ref())?) {
577
                (Lit::Int(a_int), Lit::Int(b_int)) => Some(Lit::Int(a_int * b_int)),
578
                _ => None,
579
            }
580
        }
581
        Expr::Defined(_, _) => todo!(),
582
        Expr::Range(_, _) => todo!(),
583
        Expr::Image(_, _, _) => todo!(),
584
        Expr::ImageSet(_, _, _) => todo!(),
585
        Expr::PreImage(_, _, _) => todo!(),
586
        Expr::Inverse(_, _, _) => todo!(),
587
        Expr::Restrict(_, _, _) => todo!(),
588
280
        Expr::LexLt(_, a, b) => {
589
280
            let lt = vec_expr_pairs_op::<i32, _>(a, b, |pairs, (a_len, b_len)| {
590
80
                pairs
591
80
                    .iter()
592
160
                    .find_map(|(a, b)| match a.cmp(b) {
593
40
                        CmpOrdering::Less => Some(true),     // First difference is <
594
                        CmpOrdering::Greater => Some(false), // First difference is >
595
120
                        CmpOrdering::Equal => None,          // No difference
596
160
                    })
597
80
                    .unwrap_or(a_len < b_len) // [1,1] <lex [1,1,x]
598
200
            })?;
599
80
            Some(lt.into())
600
        }
601
3040
        Expr::LexLeq(_, a, b) => {
602
3040
            let lt = vec_expr_pairs_op::<i32, _>(a, b, |pairs, (a_len, b_len)| {
603
40
                pairs
604
40
                    .iter()
605
80
                    .find_map(|(a, b)| match a.cmp(b) {
606
40
                        CmpOrdering::Less => Some(true),
607
                        CmpOrdering::Greater => Some(false),
608
40
                        CmpOrdering::Equal => None,
609
80
                    })
610
40
                    .unwrap_or(a_len <= b_len) // [1,1] <=lex [1,1,x]
611
3000
            })?;
612
40
            Some(lt.into())
613
        }
614
        Expr::LexGt(_, a, b) => eval_constant(&Expr::LexLt(Metadata::new(), b.clone(), a.clone())),
615
        Expr::LexGeq(_, a, b) => {
616
            eval_constant(&Expr::LexLeq(Metadata::new(), b.clone(), a.clone()))
617
        }
618
40
        Expr::FlatLexLt(_, a, b) => {
619
40
            let lt = atoms_pairs_op::<i32, _>(a, b, |pairs, (a_len, b_len)| {
620
                pairs
621
                    .iter()
622
                    .find_map(|(a, b)| match a.cmp(b) {
623
                        CmpOrdering::Less => Some(true),
624
                        CmpOrdering::Greater => Some(false),
625
                        CmpOrdering::Equal => None,
626
                    })
627
                    .unwrap_or(a_len < b_len)
628
40
            })?;
629
            Some(lt.into())
630
        }
631
40
        Expr::FlatLexLeq(_, a, b) => {
632
40
            let lt = atoms_pairs_op::<i32, _>(a, b, |pairs, (a_len, b_len)| {
633
                pairs
634
                    .iter()
635
                    .find_map(|(a, b)| match a.cmp(b) {
636
                        CmpOrdering::Less => Some(true),
637
                        CmpOrdering::Greater => Some(false),
638
                        CmpOrdering::Equal => None,
639
                    })
640
                    .unwrap_or(a_len <= b_len)
641
40
            })?;
642
            Some(lt.into())
643
        }
644
    }
645
1911080
}
646

            
647
8100
pub fn un_op<T, A>(f: fn(T) -> A, a: &Expr) -> Option<A>
648
8100
where
649
8100
    T: TryFrom<Lit>,
650
{
651
8100
    let a = unwrap_expr::<T>(a)?;
652
100
    Some(f(a))
653
8100
}
654

            
655
306600
pub fn bin_op<T, A>(f: fn(T, T) -> A, a: &Expr, b: &Expr) -> Option<A>
656
306600
where
657
306600
    T: TryFrom<Lit>,
658
{
659
306600
    let a = unwrap_expr::<T>(a)?;
660
4160
    let b = unwrap_expr::<T>(b)?;
661
2480
    Some(f(a, b))
662
306600
}
663

            
664
#[allow(dead_code)]
665
pub fn tern_op<T, A>(f: fn(T, T, T) -> A, a: &Expr, b: &Expr, c: &Expr) -> Option<A>
666
where
667
    T: TryFrom<Lit>,
668
{
669
    let a = unwrap_expr::<T>(a)?;
670
    let b = unwrap_expr::<T>(b)?;
671
    let c = unwrap_expr::<T>(c)?;
672
    Some(f(a, b, c))
673
}
674

            
675
4590
pub fn vec_op<T, A>(f: fn(Vec<T>) -> A, a: &[Expr]) -> Option<A>
676
4590
where
677
4590
    T: TryFrom<Lit>,
678
{
679
4590
    let a = a.iter().map(unwrap_expr).collect::<Option<Vec<T>>>()?;
680
6
    Some(f(a))
681
4590
}
682

            
683
140300
pub fn vec_lit_op<T, A>(f: fn(Vec<T>) -> A, a: &Expr) -> Option<A>
684
140300
where
685
140300
    T: TryFrom<Lit>,
686
{
687
    // we don't care about preserving indices here, as we will be getting rid of the vector
688
    // anyways!
689
140300
    let a = a.clone().unwrap_matrix_unchecked()?.0;
690
123560
    let a = a.iter().map(unwrap_expr).collect::<Option<Vec<T>>>()?;
691
4880
    Some(f(a))
692
140300
}
693

            
694
type PairsCallback<T, A> = fn(Vec<(T, T)>, (usize, usize)) -> A;
695

            
696
/// Calls the given function on each consecutive pair of elements in the list expressions.
697
/// Also passes the length of the two lists.
698
3320
fn vec_expr_pairs_op<T, A>(a: &Expr, b: &Expr, f: PairsCallback<T, A>) -> Option<A>
699
3320
where
700
3320
    T: TryFrom<Lit>,
701
{
702
3320
    let a_exprs = a.clone().unwrap_matrix_unchecked()?.0;
703
280
    let b_exprs = b.clone().unwrap_matrix_unchecked()?.0;
704
200
    let lens = (a_exprs.len(), b_exprs.len());
705

            
706
200
    let lit_pairs = std::iter::zip(a_exprs, b_exprs)
707
320
        .map(|(a, b)| Some((unwrap_expr(&a)?, unwrap_expr(&b)?)))
708
200
        .collect::<Option<Vec<(T, T)>>>()?;
709
120
    Some(f(lit_pairs, lens))
710
3320
}
711

            
712
/// Same as [`vec_expr_pairs_op`], but over slices of atoms.
713
80
fn atoms_pairs_op<T, A>(a: &[Atom], b: &[Atom], f: PairsCallback<T, A>) -> Option<A>
714
80
where
715
80
    T: TryFrom<Atom>,
716
{
717
80
    let lit_pairs = Iterator::zip(a.iter(), b.iter())
718
80
        .map(|(a, b)| Some((a.clone().try_into().ok()?, b.clone().try_into().ok()?)))
719
80
        .collect::<Option<Vec<(T, T)>>>()?;
720
    Some(f(lit_pairs, (a.len(), b.len())))
721
80
}
722

            
723
pub fn opt_vec_op<T, A>(f: fn(Vec<T>) -> Option<A>, a: &[Expr]) -> Option<A>
724
where
725
    T: TryFrom<Lit>,
726
{
727
    let a = a.iter().map(unwrap_expr).collect::<Option<Vec<T>>>()?;
728
    f(a)
729
}
730

            
731
2620
pub fn opt_vec_lit_op<T, A>(f: fn(Vec<T>) -> Option<A>, a: &Expr) -> Option<A>
732
2620
where
733
2620
    T: TryFrom<Lit>,
734
{
735
2620
    let a = a.clone().unwrap_list()?;
736
    // FIXME: deal with explicit matrix domains
737
440
    let a = a.iter().map(unwrap_expr).collect::<Option<Vec<T>>>()?;
738
20
    f(a)
739
2620
}
740

            
741
#[allow(dead_code)]
742
pub fn flat_op<T, A>(f: fn(Vec<T>, T) -> A, a: &[Expr], b: &Expr) -> Option<A>
743
where
744
    T: TryFrom<Lit>,
745
{
746
    let a = a.iter().map(unwrap_expr).collect::<Option<Vec<T>>>()?;
747
    let b = unwrap_expr::<T>(b)?;
748
    Some(f(a, b))
749
}
750

            
751
521400
pub fn unwrap_expr<T: TryFrom<Lit>>(expr: &Expr) -> Option<T> {
752
521400
    let c = eval_constant(expr)?;
753
37800
    TryInto::<T>::try_into(c).ok()
754
521400
}