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

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

            
26
320
                if a_set.difference(&b_set).count() > 0 {
27
240
                    Some(Lit::Bool(a_set.is_superset(&b_set)))
28
                } else {
29
80
                    Some(Lit::Bool(false))
30
                }
31
            }
32
320
            _ => None,
33
        },
34
640
        Expr::SupsetEq(_, a, b) => match (a.as_ref(), b.as_ref()) {
35
            (
36
320
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
37
320
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
38
320
            ) => Some(Lit::Bool(
39
320
                a.iter()
40
320
                    .cloned()
41
320
                    .collect::<HashSet<Lit>>()
42
320
                    .is_superset(&b.iter().cloned().collect::<HashSet<Lit>>()),
43
320
            )),
44
320
            _ => None,
45
        },
46
1200
        Expr::Subset(_, a, b) => match (a.as_ref(), b.as_ref()) {
47
            (
48
400
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
49
400
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
50
            ) => {
51
400
                let a_set: HashSet<Lit> = a.iter().cloned().collect();
52
400
                let b_set: HashSet<Lit> = b.iter().cloned().collect();
53

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

            
132
288
            domain.contains(lit).ok().map(Into::into)
133
        }
134
11370620
        Expr::Atomic(_, Atom::Literal(c)) => Some(c.clone()),
135
16461572
        Expr::Atomic(_, Atom::Reference(reference)) => reference.resolve_constant(),
136
4889808
        Expr::AbstractLiteral(_, a) => Some(Lit::AbstractLiteral(a.clone().into_literals()?)),
137
555392
        Expr::Comprehension(_, comprehension) => {
138
555392
            eval_constant_comprehension(comprehension.as_ref())
139
        }
140
        Expr::AbstractComprehension(_, _) => None,
141
16400
        Expr::RecordField(_, rec, fld_name) => {
142
            if let Expr::Atomic(
143
                _,
144
160
                Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Record(ents))),
145
16400
            ) = rec.as_ref()
146
            {
147
240
                for Field { name, value } in ents {
148
240
                    if name.eq(fld_name) {
149
160
                        return Some(value.clone());
150
80
                    }
151
                }
152
16240
            }
153
16240
            None
154
        }
155
4565444
        Expr::UnsafeIndex(_, subject, indices) | Expr::SafeIndex(_, subject, indices) => {
156
6604880
            let subject: Lit = eval_constant(subject.as_ref())?;
157
166240
            let indices: Vec<Lit> = indices
158
166240
                .iter()
159
166240
                .map(eval_constant)
160
166240
                .collect::<Option<Vec<Lit>>>()?;
161

            
162
30520
            match subject {
163
30120
                Lit::AbstractLiteral(subject @ AbstractLiteral::Matrix(_, _)) => {
164
30120
                    matrix::flatten_enumerate(subject)
165
82348
                        .find(|(i, _)| i == &indices)
166
30120
                        .map(|(_, x)| x)
167
                }
168
400
                Lit::AbstractLiteral(subject @ AbstractLiteral::Tuple(_)) => {
169
400
                    let AbstractLiteral::Tuple(elems) = subject else {
170
                        return None;
171
                    };
172

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

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

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

            
183
                    // -1 for 0-indexing vs 1-indexing
184
400
                    let item = elems[index as usize - 1].clone();
185

            
186
400
                    Some(item)
187
                }
188
                Lit::AbstractLiteral(subject @ AbstractLiteral::Record(_)) => {
189
                    let AbstractLiteral::Record(elems) = subject else {
190
                        return None;
191
                    };
192

            
193
                    assert!(indices.len() == 1, "nested record not supported yet");
194

            
195
                    let Lit::Int(index) = indices[0].clone() else {
196
                        return None;
197
                    };
198

            
199
                    if elems.len() < index as usize || index < 1 {
200
                        return None;
201
                    }
202

            
203
                    // -1 for 0-indexing vs 1-indexing
204
                    let item = elems[index as usize - 1].clone();
205
                    Some(item.value)
206
                }
207
                _ => None,
208
            }
209
        }
210
54080
        Expr::UnsafeSlice(_, subject, indices) | Expr::SafeSlice(_, subject, indices) => {
211
102960
            let subject: Lit = eval_constant(subject.as_ref())?;
212
80
            let Lit::AbstractLiteral(subject @ AbstractLiteral::Matrix(_, _)) = subject else {
213
                return None;
214
            };
215

            
216
80
            let hole_dim = indices
217
80
                .iter()
218
80
                .cloned()
219
160
                .position(|x| x.is_none())
220
80
                .expect("slice expression should have a hole dimension");
221

            
222
80
            let missing_domain = matrix::index_domains(&subject)[hole_dim].clone();
223

            
224
80
            let indices: Vec<Option<Lit>> = indices
225
80
                .iter()
226
80
                .cloned()
227
160
                .map(|x| {
228
                    // the outer option represents success of this iterator, the inner the index
229
                    // slice.
230
160
                    match x {
231
80
                        Some(x) => eval_constant(&x).map(Some),
232
80
                        None => Some(None),
233
                    }
234
160
                })
235
80
                .collect::<Option<Vec<Option<Lit>>>>()?;
236

            
237
80
            let indices_in_slice: Vec<Vec<Lit>> = missing_domain
238
80
                .values()
239
80
                .ok()?
240
240
                .map(|i| {
241
240
                    let mut indices = indices.clone();
242
240
                    indices[hole_dim] = Some(i);
243
                    // These unwraps will only fail if we have multiple holes.
244
                    // As this is invalid, panicking is fine.
245
480
                    indices.into_iter().map(|x| x.unwrap()).collect_vec()
246
240
                })
247
80
                .collect_vec();
248

            
249
            // Note: indices_in_slice is not necessarily sorted, so this is the best way.
250
80
            let elems = matrix::flatten_enumerate(subject)
251
720
                .filter(|(i, _)| indices_in_slice.contains(i))
252
80
                .map(|(_, elem)| elem)
253
80
                .collect();
254

            
255
80
            Some(Lit::AbstractLiteral(into_matrix![elems]))
256
        }
257
25920
        Expr::Abs(_, e) => un_op::<i32, i32>(|a| a.abs(), e).map(Lit::Int),
258
1586676
        Expr::Eq(_, a, b) => bin_op::<i32, bool>(|a, b| a == b, a, b)
259
1586676
            .or_else(|| bin_op::<bool, bool>(|a, b| a == b, a, b))
260
1586676
            .map(Lit::Bool),
261
905592
        Expr::Neq(_, a, b) => bin_op::<i32, bool>(|a, b| a != b, a, b).map(Lit::Bool),
262
273944
        Expr::Lt(_, a, b) => bin_op::<i32, bool>(|a, b| a < b, a, b).map(Lit::Bool),
263
14496
        Expr::Gt(_, a, b) => bin_op::<i32, bool>(|a, b| a > b, a, b).map(Lit::Bool),
264
1160772
        Expr::Leq(_, a, b) => bin_op::<i32, bool>(|a, b| a <= b, a, b).map(Lit::Bool),
265
346152
        Expr::Geq(_, a, b) => bin_op::<i32, bool>(|a, b| a >= b, a, b).map(Lit::Bool),
266
450996
        Expr::Not(_, expr) => un_op::<bool, bool>(|e| !e, expr).map(Lit::Bool),
267
1635888
        Expr::And(_, e) => {
268
1635888
            vec_lit_op::<bool, bool>(|e| e.iter().all(|&e| e), e.as_ref()).map(Lit::Bool)
269
        }
270
1440
        Expr::Table(_, _, _) => None,
271
240
        Expr::NegativeTable(_, _, _) => None,
272
505988
        Expr::Root(_, _) => None,
273
1024104
        Expr::Or(_, es) => {
274
            // possibly cheating; definitely should be in partial eval instead
275
1651824
            for e in (**es).clone().unwrap_list()? {
276
5736
                if let Expr::Atomic(_, Atom::Literal(Lit::Bool(true))) = e {
277
3904
                    return Some(Lit::Bool(true));
278
1647920
                };
279
            }
280

            
281
793328
            vec_lit_op::<bool, bool>(|e| e.iter().any(|&e| e), es.as_ref()).map(Lit::Bool)
282
        }
283
538008
        Expr::Imply(_, box1, box2) => {
284
538008
            let a: &Atom = (&**box1).try_into().ok()?;
285
306160
            let b: &Atom = (&**box2).try_into().ok()?;
286

            
287
134736
            let a: bool = a.try_into().ok()?;
288
123400
            let b: bool = b.try_into().ok()?;
289

            
290
123400
            if a {
291
                // true -> b ~> b
292
60500
                Some(Lit::Bool(b))
293
            } else {
294
                // false -> b ~> true
295
62900
                Some(Lit::Bool(true))
296
            }
297
        }
298
25784
        Expr::Iff(_, box1, box2) => {
299
25784
            let a: &Atom = (&**box1).try_into().ok()?;
300
4256
            let b: &Atom = (&**box2).try_into().ok()?;
301

            
302
336
            let a: bool = a.try_into().ok()?;
303
96
            let b: bool = b.try_into().ok()?;
304

            
305
16
            Some(Lit::Bool(a == b))
306
        }
307
2713680
        Expr::Sum(_, exprs) => vec_lit_op::<i32, i32>(|e| e.iter().sum(), exprs).map(Lit::Int),
308
959844
        Expr::Product(_, exprs) => {
309
959844
            vec_lit_op::<i32, i32>(|e| e.iter().product(), exprs).map(Lit::Int)
310
        }
311
351076
        Expr::FlatIneq(_, a, b, c) => {
312
351076
            let a: i32 = a.try_into().ok()?;
313
216468
            let b: i32 = b.try_into().ok()?;
314
162800
            let c: i32 = c.try_into().ok()?;
315

            
316
162800
            Some(Lit::Bool(a <= b + c))
317
        }
318
414068
        Expr::FlatSumGeq(_, exprs, a) => {
319
778468
            let sum = exprs.iter().try_fold(0, |acc, atom: &Atom| {
320
778468
                let n: i32 = atom.try_into().ok()?;
321
364400
                let acc = acc + n;
322
364400
                Some(acc)
323
778468
            })?;
324

            
325
            Some(Lit::Bool(sum >= a.try_into().ok()?))
326
        }
327
462180
        Expr::FlatSumLeq(_, exprs, a) => {
328
838524
            let sum = exprs.iter().try_fold(0, |acc, atom: &Atom| {
329
838524
                let n: i32 = atom.try_into().ok()?;
330
376504
                let acc = acc + n;
331
376504
                Some(acc)
332
838524
            })?;
333

            
334
160
            Some(Lit::Bool(sum >= a.try_into().ok()?))
335
        }
336
66120
        Expr::Min(_, e) => {
337
66120
            opt_vec_lit_op::<i32, i32>(|e| e.iter().min().copied(), e.as_ref()).map(Lit::Int)
338
        }
339
68052
        Expr::Max(_, e) => {
340
68052
            opt_vec_lit_op::<i32, i32>(|e| e.iter().max().copied(), e.as_ref()).map(Lit::Int)
341
        }
342
150120
        Expr::UnsafeDiv(_, a, b) | Expr::SafeDiv(_, a, b) => {
343
251520
            if unwrap_expr::<i32>(b)? == 0 {
344
                return None;
345
21680
            }
346
21680
            bin_op::<i32, i32>(|a, b| ((a as f32) / (b as f32)).floor() as i32, a, b).map(Lit::Int)
347
        }
348
40800
        Expr::UnsafeMod(_, a, b) | Expr::SafeMod(_, a, b) => {
349
60080
            if unwrap_expr::<i32>(b)? == 0 {
350
                return None;
351
4480
            }
352
4480
            bin_op::<i32, i32>(|a, b| a - b * (a as f32 / b as f32).floor() as i32, a, b)
353
4480
                .map(Lit::Int)
354
        }
355
        Expr::Substring(_, s, t) => match (s.as_ref(), t.as_ref()) {
356
            (
357
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Sequence(s)))),
358
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Sequence(t)))),
359
            ) => {
360
                if s.len() > t.len() {
361
                    return Some(Lit::Bool(false));
362
                }
363

            
364
                let found = t.windows(s.len()).any(|window| window == s.as_slice());
365
                Some(Lit::Bool(found))
366
            }
367
            _ => None,
368
        },
369
        Expr::Subsequence(_, s, t) => match (s.as_ref(), t.as_ref()) {
370
            (
371
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Sequence(s)))),
372
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Sequence(t)))),
373
            ) => {
374
                let mut i = 0;
375
                let mut j = 0;
376

            
377
                while i < s.len() && j < t.len() {
378
                    if s[i] == t[j] {
379
                        i += 1;
380
                    }
381
                    j += 1;
382
                }
383

            
384
                Some(Lit::Bool(i == s.len()))
385
            }
386
            _ => None,
387
        },
388
12040
        Expr::MinionDivEqUndefZero(_, a, b, c) => {
389
            // div always rounds down
390
12040
            let a: i32 = a.try_into().ok()?;
391
240
            let b: i32 = b.try_into().ok()?;
392
            let c: i32 = c.try_into().ok()?;
393

            
394
            if b == 0 {
395
                return None;
396
            }
397

            
398
            let a = a as f32;
399
            let b = b as f32;
400
            let div: i32 = (a / b).floor() as i32;
401
            Some(Lit::Bool(div == c))
402
        }
403
53472
        Expr::Bubble(_, a, b) => bin_op::<bool, bool>(|a, b| a && b, a, b).map(Lit::Bool),
404
180820
        Expr::MinionReify(_, a, b) => {
405
180820
            let result = eval_constant(a)?;
406

            
407
10840
            let result: bool = result.try_into().ok()?;
408
10840
            let b: bool = b.try_into().ok()?;
409

            
410
            Some(Lit::Bool(b == result))
411
        }
412
83952
        Expr::MinionReifyImply(_, a, b) => {
413
83952
            let result = eval_constant(a)?;
414

            
415
            let result: bool = result.try_into().ok()?;
416
            let b: bool = b.try_into().ok()?;
417

            
418
            if b {
419
                Some(Lit::Bool(result))
420
            } else {
421
                Some(Lit::Bool(true))
422
            }
423
        }
424
2960
        Expr::MinionModuloEqUndefZero(_, a, b, c) => {
425
            // From Savile Row. Same semantics as division.
426
            //
427
            //   a - (b * floor(a/b))
428
            //
429
            // We don't use % as it has the same semantics as /. We don't use / as we want to round
430
            // down instead, not towards zero.
431

            
432
2960
            let a: i32 = a.try_into().ok()?;
433
240
            let b: i32 = b.try_into().ok()?;
434
            let c: i32 = c.try_into().ok()?;
435

            
436
            if b == 0 {
437
                return None;
438
            }
439

            
440
            let modulo = a - b * (a as f32 / b as f32).floor() as i32;
441
            Some(Lit::Bool(modulo == c))
442
        }
443
5100
        Expr::MinionPow(_, a, b, c) => {
444
            // only available for positive a b c
445

            
446
5100
            let a: i32 = a.try_into().ok()?;
447
            let b: i32 = b.try_into().ok()?;
448
            let c: i32 = c.try_into().ok()?;
449

            
450
            if a <= 0 {
451
                return None;
452
            }
453

            
454
            if b <= 0 {
455
                return None;
456
            }
457

            
458
            if c <= 0 {
459
                return None;
460
            }
461

            
462
            Some(Lit::Bool(a ^ b == c))
463
        }
464
640
        Expr::MinionWInSet(_, _, _) => None,
465
2040
        Expr::MinionWInIntervalSet(_, x, intervals) => {
466
2040
            let x_lit: &Lit = x.try_into().ok()?;
467

            
468
            let x_lit = match x_lit.clone() {
469
                Lit::Int(i) => Some(i),
470
                Lit::Bool(true) => Some(1),
471
                Lit::Bool(false) => Some(0),
472
                _ => None,
473
            }?;
474

            
475
            let mut intervals = intervals.iter();
476
            while let Some(lower) = intervals.next() {
477
                let Some(upper) = intervals.next() else {
478
                    break;
479
                };
480
                if &x_lit >= lower && &x_lit <= upper {
481
                    return Some(Lit::Bool(true));
482
                }
483
            }
484

            
485
            Some(Lit::Bool(false))
486
        }
487
        Expr::Flatten(_, _, _) => {
488
            // TODO
489
22312
            None
490
        }
491
106832
        Expr::AllDiff(_, e) => {
492
106832
            let es = (**e).clone().unwrap_list()?;
493
12308
            let mut lits: HashSet<Lit> = HashSet::new();
494
13428
            for expr in es {
495
6028
                let Expr::Atomic(_, Atom::Literal(x)) = expr else {
496
11988
                    return None;
497
                };
498
1440
                match x {
499
                    Lit::Int(_) | Lit::Bool(_) => {
500
1440
                        if lits.contains(&x) {
501
                            return Some(Lit::Bool(false));
502
1440
                        } else {
503
1440
                            lits.insert(x.clone());
504
1440
                        }
505
                    }
506
                    Lit::AbstractLiteral(_) => return None, // Reject AbstractLiteral cases
507
                }
508
            }
509
320
            Some(Lit::Bool(true))
510
        }
511
62948
        Expr::FlatAllDiff(_, es) => {
512
62948
            let mut lits: HashSet<Lit> = HashSet::new();
513
62948
            for atom in es {
514
62948
                let Atom::Literal(x) = atom else {
515
62948
                    return None;
516
                };
517

            
518
                match x {
519
                    Lit::Int(_) | Lit::Bool(_) => {
520
                        if lits.contains(x) {
521
                            return Some(Lit::Bool(false));
522
                        } else {
523
                            lits.insert(x.clone());
524
                        }
525
                    }
526
                    Lit::AbstractLiteral(_) => return None, // Reject AbstractLiteral cases
527
                }
528
            }
529
            Some(Lit::Bool(true))
530
        }
531
80972
        Expr::FlatWatchedLiteral(_, _, _) => None,
532
202172
        Expr::AuxDeclaration(_, _, _) => None,
533
125580
        Expr::Neg(_, a) => match eval_constant(a.as_ref())? {
534
19912
            Lit::Int(a) => Some(Lit::Int(-a)),
535
            _ => None,
536
        },
537
        Expr::Factorial(_, _) => None,
538
697044
        Expr::Minus(_, a, b) => bin_op::<i32, i32>(|a, b| a - b, a, b).map(Lit::Int),
539
1680
        Expr::FlatMinusEq(_, a, b) => {
540
1680
            let a: i32 = a.try_into().ok()?;
541
            let b: i32 = b.try_into().ok()?;
542
            Some(Lit::Bool(a == -b))
543
        }
544
6560
        Expr::FlatProductEq(_, a, b, c) => {
545
6560
            let a: i32 = a.try_into().ok()?;
546
            let b: i32 = b.try_into().ok()?;
547
            let c: i32 = c.try_into().ok()?;
548
            Some(Lit::Bool(a * b == c))
549
        }
550
155800
        Expr::FlatWeightedSumLeq(_, cs, vs, total) => {
551
155800
            let cs: Vec<i32> = cs
552
155800
                .iter()
553
321160
                .map(|x| TryInto::<i32>::try_into(x).ok())
554
155800
                .collect::<Option<Vec<i32>>>()?;
555
155800
            let vs: Vec<i32> = vs
556
155800
                .iter()
557
259560
                .map(|x| TryInto::<i32>::try_into(x).ok())
558
155800
                .collect::<Option<Vec<i32>>>()?;
559
2640
            let total: i32 = total.try_into().ok()?;
560

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

            
563
            Some(Lit::Bool(sum <= total))
564
        }
565
152280
        Expr::FlatWeightedSumGeq(_, cs, vs, total) => {
566
152280
            let cs: Vec<i32> = cs
567
152280
                .iter()
568
310160
                .map(|x| TryInto::<i32>::try_into(x).ok())
569
152280
                .collect::<Option<Vec<i32>>>()?;
570
152280
            let vs: Vec<i32> = vs
571
152280
                .iter()
572
253560
                .map(|x| TryInto::<i32>::try_into(x).ok())
573
152280
                .collect::<Option<Vec<i32>>>()?;
574
2120
            let total: i32 = total.try_into().ok()?;
575

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

            
578
            Some(Lit::Bool(sum >= total))
579
        }
580
3760
        Expr::FlatAbsEq(_, x, y) => {
581
3760
            let x: i32 = x.try_into().ok()?;
582
160
            let y: i32 = y.try_into().ok()?;
583

            
584
            Some(Lit::Bool(x == y.abs()))
585
        }
586
293056
        Expr::UnsafePow(_, a, b) | Expr::SafePow(_, a, b) => {
587
561200
            let a: &Atom = a.try_into().ok()?;
588
556800
            let a: i32 = a.try_into().ok()?;
589

            
590
481652
            let b: &Atom = b.try_into().ok()?;
591
481652
            let b: i32 = b.try_into().ok()?;
592

            
593
481652
            if (a != 0 || b != 0) && b >= 0 {
594
481652
                Some(Lit::Int(a.pow(b as u32)))
595
            } else {
596
                None
597
            }
598
        }
599
        Expr::Metavar(_, _) => None,
600
218448
        Expr::MinionElementOne(_, _, _, _) => None,
601
98072
        Expr::ToInt(_, expression) => {
602
98072
            let lit = eval_constant(expression.as_ref())?;
603
3224
            match lit {
604
                Lit::Int(_) => Some(lit),
605
3220
                Lit::Bool(true) => Some(Lit::Int(1)),
606
4
                Lit::Bool(false) => Some(Lit::Int(0)),
607
                _ => None,
608
            }
609
        }
610
        Expr::SATInt(_, _, _, _) => {
611
            // TODO: If this SATInt is composed of literals, we should evaluate it back to an
612
            // integer literal.
613
            //
614
            // This is important because `is_all_constant` currently returns true for SATInts
615
            // containing no references. If we don't evaluate them here, bubble rules will skip
616
            // them (thinking they'll be constant-folded later), but they'll actually reach
617
            // the solver adaptors as un-encoded unsafe operations, causing panics.
618
2177940
            None
619
        }
620
        Expr::PairwiseSum(_, a, b) => {
621
            match (eval_constant(a.as_ref())?, eval_constant(b.as_ref())?) {
622
                (Lit::Int(a_int), Lit::Int(b_int)) => Some(Lit::Int(a_int + b_int)),
623
                _ => None,
624
            }
625
        }
626
        Expr::PairwiseProduct(_, a, b) => {
627
            match (eval_constant(a.as_ref())?, eval_constant(b.as_ref())?) {
628
                (Lit::Int(a_int), Lit::Int(b_int)) => Some(Lit::Int(a_int * b_int)),
629
                _ => None,
630
            }
631
        }
632
        Expr::Defined(_, _) => todo!(),
633
        Expr::Range(_, _) => todo!(),
634
        Expr::Image(_, _, _) => todo!(),
635
        Expr::ImageSet(_, _, _) => todo!(),
636
        Expr::PreImage(_, _, _) => todo!(),
637
        Expr::Inverse(_, _, _) => todo!(),
638
        Expr::Restrict(_, _, _) => todo!(),
639
        Expr::Active(_, _, _) => todo!(),
640
        Expr::ToSet(_, _) => todo!(),
641
        Expr::ToMSet(_, _) => todo!(),
642
        Expr::ToRelation(_, _) => todo!(),
643
        Expr::RelationProj(_, _, _) => todo!(),
644
        Expr::Apart(_, _, _) => todo!(),
645
        Expr::Together(_, _, _) => todo!(),
646
        Expr::Participants(_, _) => todo!(),
647
        Expr::Party(_, _, _) => todo!(),
648
        Expr::Parts(_, _) => todo!(),
649
        Expr::Card(_, _) => todo!(),
650
2004
        Expr::LexLt(_, a, b) => {
651
2004
            let lt = vec_expr_pairs_op::<i32, _>(a, b, |pairs, (a_len, b_len)| {
652
320
                pairs
653
320
                    .iter()
654
640
                    .find_map(|(a, b)| match a.cmp(b) {
655
160
                        CmpOrdering::Less => Some(true),     // First difference is <
656
                        CmpOrdering::Greater => Some(false), // First difference is >
657
480
                        CmpOrdering::Equal => None,          // No difference
658
640
                    })
659
320
                    .unwrap_or(a_len < b_len) // [1,1] <lex [1,1,x]
660
1684
            })?;
661
320
            Some(lt.into())
662
        }
663
82760
        Expr::LexLeq(_, a, b) => {
664
82760
            let lt = vec_expr_pairs_op::<i32, _>(a, b, |pairs, (a_len, b_len)| {
665
160
                pairs
666
160
                    .iter()
667
320
                    .find_map(|(a, b)| match a.cmp(b) {
668
160
                        CmpOrdering::Less => Some(true),
669
                        CmpOrdering::Greater => Some(false),
670
160
                        CmpOrdering::Equal => None,
671
320
                    })
672
160
                    .unwrap_or(a_len <= b_len) // [1,1] <=lex [1,1,x]
673
82600
            })?;
674
160
            Some(lt.into())
675
        }
676
120
        Expr::LexGt(_, a, b) => eval_constant(&Expr::LexLt(Metadata::new(), b.clone(), a.clone())),
677
240
        Expr::LexGeq(_, a, b) => {
678
240
            eval_constant(&Expr::LexLeq(Metadata::new(), b.clone(), a.clone()))
679
        }
680
320
        Expr::FlatLexLt(_, a, b) => {
681
320
            let lt = atoms_pairs_op::<i32, _>(a, b, |pairs, (a_len, b_len)| {
682
                pairs
683
                    .iter()
684
                    .find_map(|(a, b)| match a.cmp(b) {
685
                        CmpOrdering::Less => Some(true),
686
                        CmpOrdering::Greater => Some(false),
687
                        CmpOrdering::Equal => None,
688
                    })
689
                    .unwrap_or(a_len < b_len)
690
320
            })?;
691
            Some(lt.into())
692
        }
693
480
        Expr::FlatLexLeq(_, a, b) => {
694
480
            let lt = atoms_pairs_op::<i32, _>(a, b, |pairs, (a_len, b_len)| {
695
                pairs
696
                    .iter()
697
                    .find_map(|(a, b)| match a.cmp(b) {
698
                        CmpOrdering::Less => Some(true),
699
                        CmpOrdering::Greater => Some(false),
700
                        CmpOrdering::Equal => None,
701
                    })
702
                    .unwrap_or(a_len <= b_len)
703
480
            })?;
704
            Some(lt.into())
705
        }
706
    }
707
58989964
}
708

            
709
476916
pub fn un_op<T, A>(f: fn(T) -> A, a: &Expr) -> Option<A>
710
476916
where
711
476916
    T: TryFrom<Lit>,
712
{
713
476916
    let a = unwrap_expr::<T>(a)?;
714
124076
    Some(f(a))
715
476916
}
716

            
717
6482416
pub fn bin_op<T, A>(f: fn(T, T) -> A, a: &Expr, b: &Expr) -> Option<A>
718
6482416
where
719
6482416
    T: TryFrom<Lit>,
720
{
721
6482416
    let a = unwrap_expr::<T>(a)?;
722
779520
    let b = unwrap_expr::<T>(b)?;
723
638844
    Some(f(a, b))
724
6482416
}
725

            
726
#[allow(dead_code)]
727
pub fn tern_op<T, A>(f: fn(T, T, T) -> A, a: &Expr, b: &Expr, c: &Expr) -> Option<A>
728
where
729
    T: TryFrom<Lit>,
730
{
731
    let a = unwrap_expr::<T>(a)?;
732
    let b = unwrap_expr::<T>(b)?;
733
    let c = unwrap_expr::<T>(c)?;
734
    Some(f(a, b, c))
735
}
736

            
737
52746
pub fn vec_op<T, A>(f: fn(Vec<T>) -> A, a: &[Expr]) -> Option<A>
738
52746
where
739
52746
    T: TryFrom<Lit>,
740
{
741
52746
    let a = a.iter().map(unwrap_expr).collect::<Option<Vec<T>>>()?;
742
24
    Some(f(a))
743
52746
}
744

            
745
6102740
pub fn vec_lit_op<T, A>(f: fn(Vec<T>) -> A, a: &Expr) -> Option<A>
746
6102740
where
747
6102740
    T: TryFrom<Lit>,
748
{
749
6102740
    Some(f(eval_list_items(a)?))
750
6102740
}
751

            
752
type PairsCallback<T, A> = fn(Vec<(T, T)>, (usize, usize)) -> A;
753

            
754
/// Calls the given function on each consecutive pair of elements in the list expressions.
755
/// Also passes the length of the two lists.
756
84764
fn vec_expr_pairs_op<T, A>(a: &Expr, b: &Expr, f: PairsCallback<T, A>) -> Option<A>
757
84764
where
758
84764
    T: TryFrom<Lit>,
759
{
760
84764
    let a_exprs = a.clone().unwrap_matrix_unchecked()?.0;
761
2080
    let b_exprs = b.clone().unwrap_matrix_unchecked()?.0;
762
1280
    let lens = (a_exprs.len(), b_exprs.len());
763

            
764
1280
    let lit_pairs = std::iter::zip(a_exprs, b_exprs)
765
1760
        .map(|(a, b)| Some((unwrap_expr(&a)?, unwrap_expr(&b)?)))
766
1280
        .collect::<Option<Vec<(T, T)>>>()?;
767
480
    Some(f(lit_pairs, lens))
768
84764
}
769

            
770
/// Same as [`vec_expr_pairs_op`], but over slices of atoms.
771
800
fn atoms_pairs_op<T, A>(a: &[Atom], b: &[Atom], f: PairsCallback<T, A>) -> Option<A>
772
800
where
773
800
    T: TryFrom<Atom>,
774
{
775
800
    let lit_pairs = Iterator::zip(a.iter(), b.iter())
776
800
        .map(|(a, b)| Some((a.clone().try_into().ok()?, b.clone().try_into().ok()?)))
777
800
        .collect::<Option<Vec<(T, T)>>>()?;
778
    Some(f(lit_pairs, (a.len(), b.len())))
779
800
}
780

            
781
pub fn opt_vec_op<T, A>(f: fn(Vec<T>) -> Option<A>, a: &[Expr]) -> Option<A>
782
where
783
    T: TryFrom<Lit>,
784
{
785
    let a = a.iter().map(unwrap_expr).collect::<Option<Vec<T>>>()?;
786
    f(a)
787
}
788

            
789
134172
pub fn opt_vec_lit_op<T, A>(f: fn(Vec<T>) -> Option<A>, a: &Expr) -> Option<A>
790
134172
where
791
134172
    T: TryFrom<Lit>,
792
{
793
134172
    f(eval_list_items(a)?)
794
134172
}
795

            
796
#[allow(dead_code)]
797
pub fn flat_op<T, A>(f: fn(Vec<T>, T) -> A, a: &[Expr], b: &Expr) -> Option<A>
798
where
799
    T: TryFrom<Lit>,
800
{
801
    let a = a.iter().map(unwrap_expr).collect::<Option<Vec<T>>>()?;
802
    let b = unwrap_expr::<T>(b)?;
803
    Some(f(a, b))
804
}
805

            
806
16070280
pub fn unwrap_expr<T: TryFrom<Lit>>(expr: &Expr) -> Option<T> {
807
16070280
    let c = eval_constant(expr)?;
808
4491180
    TryInto::<T>::try_into(c).ok()
809
16070280
}
810

            
811
6236912
fn eval_list_items<T>(expr: &Expr) -> Option<Vec<T>>
812
6236912
where
813
6236912
    T: TryFrom<Lit>,
814
{
815
6236912
    if let Some(items) = expr
816
6236912
        .clone()
817
6236912
        .unwrap_matrix_unchecked()
818
6236912
        .map(|(items, _)| items)
819
    {
820
5896056
        return items.iter().map(unwrap_expr).collect();
821
340856
    }
822

            
823
340856
    let Lit::AbstractLiteral(list) = eval_constant(expr)? else {
824
        return None;
825
    };
826

            
827
6456
    let items = list.unwrap_list()?;
828
6456
    items
829
6456
        .iter()
830
6456
        .cloned()
831
6456
        .map(TryInto::try_into)
832
6456
        .collect::<Result<Vec<_>, _>>()
833
6456
        .ok()
834
6236912
}
835

            
836
555392
fn eval_constant_comprehension(comprehension: &Comprehension) -> Option<Lit> {
837
555392
    let mut values = Vec::new();
838
555392
    eval_comprehension_qualifiers(comprehension, 0, &mut values)?;
839
7656
    Some(Lit::AbstractLiteral(
840
7656
        AbstractLiteral::matrix_implied_indices(values),
841
7656
    ))
842
555392
}
843

            
844
1329412
fn eval_comprehension_qualifiers(
845
1329412
    comprehension: &Comprehension,
846
1329412
    qualifier_index: usize,
847
1329412
    values: &mut Vec<Lit>,
848
1329412
) -> Option<()> {
849
1329412
    if qualifier_index == comprehension.qualifiers.len() {
850
534452
        values.push(eval_constant(&comprehension.return_expression)?);
851
14880
        return Some(());
852
794960
    }
853

            
854
794960
    match &comprehension.qualifiers[qualifier_index] {
855
667664
        ComprehensionQualifier::Generator { ptr } => {
856
667664
            let domain = ptr.domain()?;
857
667664
            let generator_values = domain
858
667664
                .resolve()
859
667664
                .and_then(|x| x.values())
860
667664
                .ok()?
861
666212
                .collect_vec();
862

            
863
683076
            for value in generator_values {
864
683076
                with_temporary_quantified_binding(ptr, &value, || {
865
683076
                    eval_comprehension_qualifiers(comprehension, qualifier_index + 1, values)
866
683076
                })?;
867
            }
868
        }
869
820
        ComprehensionQualifier::ExpressionGenerator { ptr } => {
870
            // clone immediately so the read lock guard is dropped
871
820
            let expr = ptr.as_quantified_expr()?.clone();
872
820
            let generator_values = generator_values_from_expr(&expr)?;
873

            
874
16
            for value in generator_values {
875
16
                with_temporary_quantified_binding(ptr, &value, || {
876
16
                    eval_comprehension_qualifiers(comprehension, qualifier_index + 1, values)
877
16
                })?;
878
            }
879
        }
880
126476
        ComprehensionQualifier::Condition(condition) => match eval_constant(condition)? {
881
            Lit::Bool(true) => {
882
90928
                eval_comprehension_qualifiers(comprehension, qualifier_index + 1, values)?
883
            }
884
9640
            Lit::Bool(false) => {}
885
            _ => return None,
886
        },
887
    }
888

            
889
21520
    Some(())
890
1329412
}
891

            
892
820
fn generator_values_from_expr(expr: &Expr) -> Option<Vec<Lit>> {
893
820
    match eval_constant(expr)? {
894
16
        Lit::AbstractLiteral(AbstractLiteral::Set(values))
895
        | Lit::AbstractLiteral(AbstractLiteral::MSet(values))
896
16
        | Lit::AbstractLiteral(AbstractLiteral::Tuple(values)) => Some(values),
897
        Lit::AbstractLiteral(list) => list.unwrap_list().cloned(),
898
        _ => None,
899
    }
900
820
}
901

            
902
683092
fn with_temporary_quantified_binding<T>(
903
683092
    quantified: &crate::ast::DeclarationPtr,
904
683092
    value: &Lit,
905
683092
    f: impl FnOnce() -> Option<T>,
906
683092
) -> Option<T> {
907
683092
    let mut targets = vec![quantified.clone()];
908
683092
    if let DeclarationKind::Quantified(inner) = &*quantified.kind()
909
683076
        && let Some(generator) = inner.generator()
910
    {
911
        targets.push(generator.clone());
912
683092
    }
913

            
914
683092
    let mut originals = Vec::with_capacity(targets.len());
915
683092
    for mut target in targets {
916
683092
        let old_kind = target.replace_kind(DeclarationKind::TemporaryValueLetting(Expr::Atomic(
917
683092
            Metadata::new(),
918
683092
            Atom::Literal(value.clone()),
919
683092
        )));
920
683092
        originals.push((target, old_kind));
921
683092
    }
922

            
923
683092
    let result = f();
924

            
925
683092
    for (mut target, old_kind) in originals.into_iter().rev() {
926
683092
        let _ = target.replace_kind(old_kind);
927
683092
    }
928

            
929
683092
    result
930
683092
}