1
use crate::ast::{
2
    DomainOpError,
3
    domains::{Int, OXIDE_INT_MAX, OXIDE_INT_MIN},
4
};
5
use funcmap::{FuncMap, TryFuncMap};
6
use num_traits::Num;
7
use polyquine::Quine;
8
use serde::{Deserialize, Serialize};
9
use std::fmt::Display;
10

            
11
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, FuncMap, TryFuncMap, Quine)]
12
#[path_prefix(conjure_cp::ast)]
13
pub enum Range<A = Int> {
14
    Single(A),
15
    Bounded(A, A),
16
    UnboundedL(A),
17
    UnboundedR(A),
18
    Unbounded,
19
}
20

            
21
impl Range<Int> {
22
    /// Effective endpoints for a finite integer solver. Missing endpoints use the default
23
    /// machine limits without replacing the unbounded range stored in the model.
24
5250202
    pub fn int_bounds_with_defaults(&self) -> (Int, Int) {
25
5250202
        (
26
5250202
            *self.low().unwrap_or(&OXIDE_INT_MIN),
27
5250202
            *self.high().unwrap_or(&OXIDE_INT_MAX),
28
5250202
        )
29
5250202
    }
30
}
31

            
32
impl<A> Range<A> {
33
    /// Whether the range is **bounded** on either side. A bounded range may still be infinite.
34
    /// See also: [Range::is_finite].
35
35730295
    pub fn is_lower_or_upper_bounded(&self) -> bool {
36
35730295
        match &self {
37
            Range::Single(_)
38
            | Range::Bounded(_, _)
39
            | Range::UnboundedL(_)
40
35489652
            | Range::UnboundedR(_) => true,
41
240643
            Range::Unbounded => false,
42
        }
43
35730295
    }
44

            
45
    /// Whether the range is **unbounded** on both sides.
46
    pub fn is_unbounded(&self) -> bool {
47
        !self.is_lower_or_upper_bounded()
48
    }
49

            
50
    /// Whether the range is **finite**. See also: [Range::is_lower_or_upper_bounded].
51
17621783
    pub fn is_finite(&self) -> bool {
52
17621783
        match &self {
53
17621730
            Range::Single(_) | Range::Bounded(_, _) => true,
54
53
            Range::Unbounded | Range::UnboundedL(_) | Range::UnboundedR(_) => false,
55
        }
56
17621783
    }
57
}
58

            
59
impl<A: Ord> Range<A> {
60
49528550
    pub fn contains(&self, val: &A) -> bool {
61
49528550
        match self {
62
9804875
            Range::Single(x) => x == val,
63
38095389
            Range::Bounded(x, y) => x <= val && val <= y,
64
87549
            Range::UnboundedR(x) => x <= val,
65
1006292
            Range::UnboundedL(x) => val <= x,
66
534445
            Range::Unbounded => true,
67
        }
68
49528550
    }
69

            
70
    /// Returns the lower bound of the range, if it has one
71
72416700
    pub fn low(&self) -> Option<&A> {
72
72416700
        match self {
73
24835711
            Range::Single(a) => Some(a),
74
47350111
            Range::Bounded(a, _) => Some(a),
75
136506
            Range::UnboundedR(a) => Some(a),
76
94372
            Range::UnboundedL(_) | Range::Unbounded => None,
77
        }
78
72416700
    }
79

            
80
    /// Returns the upper bound of the range, if it has one
81
71961252
    pub fn high(&self) -> Option<&A> {
82
71961252
        match self {
83
22204452
            Range::Single(a) => Some(a),
84
49525914
            Range::Bounded(_, a) => Some(a),
85
165
            Range::UnboundedL(a) => Some(a),
86
230721
            Range::UnboundedR(_) | Range::Unbounded => None,
87
        }
88
71961252
    }
89
}
90

            
91
impl<A: Ord + Clone> Range<A> {
92
    /// Create a new range with a lower and upper bound
93
24199130
    pub fn new(lo: Option<A>, hi: Option<A>) -> Range<A> {
94
24199130
        match (lo, hi) {
95
20
            (None, None) => Range::Unbounded,
96
192
            (Some(l), None) => Range::UnboundedR(l),
97
3
            (None, Some(r)) => Range::UnboundedL(r),
98
            (Some(l), Some(r)) => {
99
24198915
                if l == r {
100
1150304
                    Range::Single(l)
101
                } else {
102
23048611
                    let min = Ord::min(&l, &r).clone();
103
23048611
                    let max = Ord::max(l, r);
104
23048611
                    Range::Bounded(min, max)
105
                }
106
            }
107
        }
108
24199130
    }
109

            
110
    /// Given a slice of ranges, create a single range that spans from the start
111
    /// of the leftmost range to the end of the rightmost range.
112
    /// An empty slice is considered equivalent to `Range::unbounded`.
113
376967
    pub fn spanning(rngs: &[Range<A>]) -> Range<A> {
114
376967
        if rngs.is_empty() {
115
1
            return Range::Unbounded;
116
376966
        }
117

            
118
376966
        let mut lo = rngs[0].low();
119
376966
        let mut hi = rngs[0].high();
120
376966
        for rng in rngs {
121
381379
            lo = match (lo, rng.low()) {
122
381371
                (Some(curr), Some(new)) => Some(curr.min(new)),
123
8
                _ => None,
124
            };
125
381379
            hi = match (hi, rng.high()) {
126
381373
                (Some(curr), Some(new)) => Some(curr.max(new)),
127
6
                _ => None,
128
            };
129
        }
130
376966
        Range::new(lo.cloned(), hi.cloned())
131
376967
    }
132
    /// Find the range such that:
133
    /// - the lower bound is the maximum of the lower bounds
134
    /// - the upper bound is the minimum of the upper bounds
135
    /// - **ranges must not be disjoint**
136
    ///
137
    /// * `DomainopError::ConflictingArgs`: if given disjoint ranges; e.g. (2..4) (6..8)
138
14014
    pub fn minimal(rngs: &[Range<A>]) -> Result<Range<A>, DomainOpError> {
139
14014
        if rngs.is_empty() {
140
            return Ok(Range::Unbounded);
141
14014
        }
142
14014
        let mut lo = rngs[0].low();
143
14014
        let mut hi = rngs[0].high();
144
14014
        for rng in rngs {
145
28028
            lo = match (lo, rng.low()) {
146
28028
                (Some(curr), Some(new)) => Some(curr.max(new)),
147
                (None, Some(new)) => Some(new),
148
                (Some(curr), None) => Some(curr),
149
                _ => None,
150
            };
151
28028
            hi = match (hi, rng.high()) {
152
28028
                (Some(curr), Some(new)) => Some(curr.min(new)),
153
                (None, Some(new)) => Some(new),
154
                (Some(curr), None) => Some(curr),
155
                _ => None,
156
            };
157
28028
            if let (Some(l), Some(h)) = (lo, hi)
158
28028
                && l > h
159
            {
160
3954
                return Err(DomainOpError::ConflictingAttrs);
161
24074
            }
162
        }
163
10060
        Ok(Range::new(lo.cloned(), hi.cloned()))
164
14014
    }
165
}
166

            
167
impl<A: Num + Ord + Clone> Range<A> {
168
7648330
    pub fn length(&self) -> Option<A> {
169
7648330
        match self {
170
1105211
            Range::Single(_) => Some(A::one()),
171
6543106
            Range::Bounded(i, j) if i > j => Some(A::zero()),
172
6543106
            Range::Bounded(i, j) => Some(j.clone() - i.clone() + A::one()),
173
13
            Range::UnboundedR(_) | Range::UnboundedL(_) | Range::Unbounded => None,
174
        }
175
7648330
    }
176

            
177
    /// Returns true if this interval overlaps another one, i.e. at least one
178
    /// number is part of both `self` and `other`
179
    /// E.g:
180
    /// - [0, 2] overlaps [2, 4]
181
    /// - [1, 3] overlaps [2, 4]
182
    /// - [4, 6] overlaps [2, 4]
183
5716971
    pub fn overlaps(&self, other: &Range<A>) -> bool {
184
5716971
        self.low()
185
5716969
            .is_none_or(|la| other.high().is_none_or(|rb| la <= rb))
186
5243035
            && self
187
5243035
                .high()
188
5243034
                .is_none_or(|ra| other.low().is_none_or(|lb| ra >= lb))
189
5716971
    }
190

            
191
    /// Returns true if this interval touches another one on the left
192
    /// E.g: [1, 2] touches_left  [3, 4]
193
6608013
    pub fn touches_left(&self, other: &Range<A>) -> bool {
194
6608013
        self.high().zip(other.low()).is_some_and(|(ra, lb)| {
195
            // The strict ordering rules out overflowing at the largest endpoint.
196
6608011
            ra < lb && ra.clone() + A::one() == *lb
197
6608011
        })
198
6608013
    }
199

            
200
    /// Returns true if this interval touches another one on the right
201
    /// E.g: [3, 4] touches_right  [1, 2]
202
2207007
    pub fn touches_right(&self, other: &Range<A>) -> bool {
203
2207007
        self.low().zip(other.high()).is_some_and(|(la, rb)| {
204
            // The strict ordering rules out overflowing at the smallest endpoint.
205
2206976
            la > rb && la.clone() - A::one() == *rb
206
2206976
        })
207
2207007
    }
208

            
209
    /// Returns true if this interval overlaps or touches another one
210
    /// E.g:
211
    /// - [1, 3] joins [4, 6]
212
    /// - [2, 4] joins [4, 6]
213
    /// - [3, 5] joins [4, 6]
214
    /// - [6, 8] joins [4, 6]
215
    /// - [7, 8] joins [4, 6]
216
6608007
    pub fn joins(&self, other: &Range<A>) -> bool {
217
6608007
        self.touches_left(other) || self.overlaps(other) || self.touches_right(other)
218
6608007
    }
219

            
220
    /// Returns true if this interval is strictly before another one
221
7
    pub fn is_before(&self, other: &Range<A>) -> bool {
222
7
        self.high()
223
7
            .is_some_and(|ra| other.low().is_some_and(|lb| ra < &(lb.clone() - A::one())))
224
7
    }
225

            
226
    /// Returns true if this interval is strictly after another one
227
6
    pub fn is_after(&self, other: &Range<A>) -> bool {
228
6
        self.low()
229
6
            .is_some_and(|la| other.high().is_some_and(|rb| la > &(rb.clone() + A::one())))
230
6
    }
231

            
232
    /// If the two ranges join, return a new range which spans both
233
6608007
    pub fn join(&self, other: &Range<A>) -> Option<Range<A>> {
234
6608007
        if self.joins(other) {
235
4587798
            let lo = Ord::min(self.low(), other.low());
236
4587798
            let hi = match (self.high(), other.high()) {
237
4587795
                (Some(a), Some(b)) => Some(Ord::max(a, b)),
238
3
                _ => None,
239
            };
240
4587798
            return Some(Range::new(lo.cloned(), hi.cloned()));
241
2020209
        }
242
2020209
        None
243
6608007
    }
244

            
245
    /// Merge all joining ranges in the list, and return a new vec of disjoint ranges.
246
    /// E.g:
247
    /// ```ignore
248
    /// [(2..3), (4), (..1), (6..8)] -> [(..4), (6..8)]
249
    /// ```
250
    ///
251
    /// # Performance
252
    /// Currently uses a naive O(n^2) algorithm.
253
    /// A more optimal approach based on interval trees is planned.
254
189919707
    pub fn squeeze(rngs: &[Range<A>]) -> Vec<Range<A>> {
255
189919707
        let mut ans = Vec::from(rngs);
256

            
257
189919707
        if ans.is_empty() {
258
20405
            return ans;
259
189899302
        }
260

            
261
        loop {
262
194487092
            let mut merged = false;
263

            
264
            // Check every pair of ranges and join them if possible
265
194487092
            'outer: for i in 0..ans.len() {
266
195825377
                for j in (i + 1)..ans.len() {
267
6607997
                    if let Some(joined) = ans[i].join(&ans[j]) {
268
4587790
                        ans[i] = joined;
269
                        // Safe to delete here because we restart the outer loop immediately
270
4587790
                        ans.remove(j);
271
4587790
                        merged = true;
272
4587790
                        break 'outer;
273
2020207
                    }
274
                }
275
            }
276

            
277
            // If no merges occurred, we're done
278
194487092
            if !merged {
279
189899302
                break;
280
4587790
            }
281
        }
282

            
283
189899302
        ans
284
189919707
    }
285

            
286
    /// If this range is bounded, returns a lazy iterator over all values within the range.
287
    /// Otherwise, returns None.
288
24103539
    pub fn iter(&self) -> Option<RangeIterator<A>> {
289
24103539
        match self {
290
952637
            Range::Single(val) => Some(RangeIterator::Single(Some(val.clone()))),
291
23099713
            Range::Bounded(start, end) => Some(RangeIterator::Bounded {
292
23099713
                current: start.clone(),
293
23099713
                end: end.clone(),
294
23099713
            }),
295
51189
            Range::UnboundedL(_) | Range::UnboundedR(_) | Range::Unbounded => None,
296
        }
297
24103539
    }
298

            
299
    /// Lazily iterate all values in a list of ranges
300
12122
    pub fn values(rngs: &[Range<A>]) -> Option<impl Iterator<Item = A>> {
301
12122
        let itrs = rngs
302
12122
            .iter()
303
12122
            .map(Range::iter)
304
12122
            .collect::<Option<Vec<RangeIterator<A>>>>()?;
305
12122
        Some(itrs.into_iter().flatten())
306
12122
    }
307

            
308
    /// True if the list of ranges is contiguous
309
44095
    pub fn is_contiguous(rngs: &[Range<A>]) -> bool {
310
44095
        Self::squeeze(rngs).len() <= 1
311
44095
    }
312

            
313
    /// Lowest value from a sequence of ranges
314
37977
    pub fn low_of(rngs: &[Range<A>]) -> Option<&A> {
315
37977
        let mut low = rngs.first()?.low()?;
316
37977
        for rng in rngs {
317
37977
            low = Ord::min(low, rng.low()?);
318
        }
319
37977
        Some(low)
320
37977
    }
321

            
322
    /// Lowest value from a sequence of ranges
323
    pub fn high_of(rngs: &[Range<A>]) -> Option<&A> {
324
        let mut hi = rngs.first()?.high()?;
325
        for rng in rngs {
326
            hi = Ord::min(hi, rng.low()?);
327
        }
328
        Some(hi)
329
    }
330

            
331
    /// Total number of values across a slice of ranges.
332
    /// Returns `None` if any range is unbounded.
333
    pub fn total_length(rngs: &[Range<A>]) -> Option<A> {
334
        rngs.iter()
335
            .try_fold(A::zero(), |acc, r| Some(acc + r.length()?))
336
    }
337
}
338

            
339
/// Iterator for Range<A> that yields values lazily
340
pub enum RangeIterator<A> {
341
    Single(Option<A>),
342
    Bounded { current: A, end: A },
343
}
344

            
345
impl<A: Num + Ord + Clone> Iterator for RangeIterator<A> {
346
    type Item = A;
347

            
348
172601863
    fn next(&mut self) -> Option<Self::Item> {
349
172601863
        match self {
350
1905013
            RangeIterator::Single(val) => val.take(),
351
            RangeIterator::Bounded { current, end } => {
352
170696850
                if current > end {
353
23093216
                    return None;
354
147603634
                }
355

            
356
147603634
                let result = current.clone();
357
147603634
                *current = current.clone() + A::one();
358

            
359
147603634
                Some(result)
360
            }
361
        }
362
172601863
    }
363
}
364

            
365
impl<A: Display> Display for Range<A> {
366
35489648
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
367
35489648
        match self {
368
7276515
            Range::Single(i) => write!(f, "{i}"),
369
26239449
            Range::Bounded(i, j) => write!(f, "{i}..{j}"),
370
1973589
            Range::UnboundedR(i) => write!(f, "{i}.."),
371
95
            Range::UnboundedL(i) => write!(f, "..{i}"),
372
            Range::Unbounded => write!(f, ""),
373
        }
374
35489648
    }
375
}
376

            
377
#[allow(unused_imports)]
378
mod test {
379
    use super::*;
380
    use crate::range;
381

            
382
    #[test]
383
1
    pub fn test_range_macros() {
384
1
        assert_eq!(range!(1..3), Range::Bounded(1, 3));
385
1
        assert_eq!(range!(1..), Range::UnboundedR(1));
386
1
        assert_eq!(range!(..3), Range::UnboundedL(3));
387
1
        assert_eq!(range!(1), Range::Single(1));
388
1
    }
389

            
390
    #[test]
391
1
    pub fn test_range_low() {
392
1
        assert_eq!(range!(1..3).low(), Some(&1));
393
1
        assert_eq!(range!(1..).low(), Some(&1));
394
1
        assert_eq!(range!(1).low(), Some(&1));
395
1
        assert_eq!(range!(..3).low(), None);
396
1
        assert_eq!(Range::<Int>::Unbounded.low(), None);
397
1
    }
398

            
399
    #[test]
400
1
    pub fn test_range_high() {
401
1
        assert_eq!(range!(1..3).high(), Some(&3));
402
1
        assert_eq!(range!(1..).high(), None);
403
1
        assert_eq!(range!(1).high(), Some(&1));
404
1
        assert_eq!(range!(..3).high(), Some(&3));
405
1
        assert_eq!(Range::<Int>::Unbounded.high(), None);
406
1
    }
407

            
408
    #[test]
409
1
    pub fn test_range_is_finite() {
410
1
        assert!(range!(1..3).is_finite());
411
1
        assert!(range!(1).is_finite());
412
1
        assert!(!range!(1..).is_finite());
413
1
        assert!(!range!(..3).is_finite());
414
1
        assert!(!Range::<Int>::Unbounded.is_finite());
415
1
    }
416

            
417
    #[test]
418
1
    pub fn test_range_bounded() {
419
1
        assert!(range!(1..3).is_lower_or_upper_bounded());
420
1
        assert!(range!(1).is_lower_or_upper_bounded());
421
1
        assert!(range!(1..).is_lower_or_upper_bounded());
422
1
        assert!(range!(..3).is_lower_or_upper_bounded());
423
1
        assert!(!Range::<Int>::Unbounded.is_lower_or_upper_bounded());
424
1
    }
425

            
426
    #[test]
427
1
    pub fn test_range_length() {
428
1
        assert_eq!(range!(1..3).length(), Some(3));
429
1
        assert_eq!(range!(1).length(), Some(1));
430
1
        assert_eq!(range!(1..).length(), None);
431
1
        assert_eq!(range!(..3).length(), None);
432
1
        assert_eq!(Range::<Int>::Unbounded.length(), None);
433
1
    }
434

            
435
    #[test]
436
1
    pub fn test_range_contains_value() {
437
1
        assert!(range!(1..3).contains(&2));
438
1
        assert!(!range!(1..3).contains(&4));
439
1
        assert!(range!(1).contains(&1));
440
1
        assert!(!range!(1).contains(&2));
441
1
        assert!(Range::Unbounded.contains(&42));
442
1
    }
443

            
444
    #[test]
445
1
    pub fn test_range_overlaps() {
446
1
        assert!(range!(1..3).overlaps(&range!(2..4)));
447
1
        assert!(range!(1..3).overlaps(&range!(3..5)));
448
1
        assert!(!range!(1..3).overlaps(&range!(4..6)));
449
1
        assert!(Range::Unbounded.overlaps(&range!(1..3)));
450
1
    }
451

            
452
    #[test]
453
1
    pub fn test_range_touches_left() {
454
1
        assert!(range!(1..2).touches_left(&range!(3..4)));
455
1
        assert!(range!(1..2).touches_left(&range!(3)));
456
1
        assert!(range!(-5..-4).touches_left(&range!(-3..2)));
457
1
        assert!(!range!(1..2).touches_left(&range!(4..5)));
458
1
        assert!(!range!(1..2).touches_left(&range!(2..3)));
459
1
        assert!(!range!(3..4).touches_left(&range!(1..2)));
460
1
    }
461

            
462
    #[test]
463
1
    pub fn test_range_touches_right() {
464
1
        assert!(range!(3..4).touches_right(&range!(1..2)));
465
1
        assert!(range!(3).touches_right(&range!(1..2)));
466
1
        assert!(range!(0..1).touches_right(&range!(-2..-1)));
467
1
        assert!(!range!(1..2).touches_right(&range!(3..4)));
468
1
        assert!(!range!(2..3).touches_right(&range!(1..2)));
469
1
        assert!(!range!(1..2).touches_right(&range!(1..2)));
470
1
    }
471

            
472
    #[test]
473
1
    pub fn test_range_is_before() {
474
1
        assert!(range!(1..2).is_before(&range!(4..5)));
475
1
        assert!(range!(1..2).is_before(&range!(4..)));
476
1
        assert!(!range!(1..2).is_before(&range!(3..)));
477
1
        assert!(!range!(1..2).is_before(&range!(..4)));
478
1
        assert!(!range!(1..2).is_before(&range!(2..4)));
479
1
        assert!(!range!(3..4).is_before(&range!(1..2)));
480
1
        assert!(!range!(1..2).is_before(&Range::Unbounded));
481
1
    }
482

            
483
    #[test]
484
1
    pub fn test_range_is_after() {
485
1
        assert!(range!(5..6).is_after(&range!(1..2)));
486
1
        assert!(range!(4..5).is_after(&range!(..2)));
487
1
        assert!(!range!(4..5).is_after(&range!(..3)));
488
1
        assert!(!range!(2..3).is_after(&range!(1..2)));
489
1
        assert!(!range!(1..2).is_after(&range!(3..4)));
490
1
        assert!(!range!(1..2).is_after(&Range::Unbounded));
491
1
    }
492

            
493
    #[test]
494
1
    pub fn test_range_squeeze() {
495
1
        let input = vec![range!(2..3), range!(4), range!(..1), range!(6..8)];
496
1
        let squeezed = Range::squeeze(&input);
497
1
        assert_eq!(squeezed, vec![range!(..4), range!(6..8)]);
498
1
    }
499

            
500
    #[test]
501
1
    pub fn test_range_spanning() {
502
1
        assert_eq!(Range::<Int>::spanning(&[]), Range::Unbounded);
503
1
        assert_eq!(Range::spanning(&[range!(1..2), range!(4..5)]), range!(1..5));
504
1
        assert_eq!(
505
1
            Range::spanning(&[range!(..0), range!(2..4)]),
506
1
            Range::UnboundedL(4)
507
        );
508
1
        assert_eq!(
509
1
            Range::spanning(&[range!(0), range!(2..3), range!(5..)]),
510
1
            Range::UnboundedR(0)
511
        );
512
1
        assert_eq!(
513
1
            Range::spanning(&[range!(..0), range!(2..)]),
514
1
            Range::Unbounded
515
        );
516
1
    }
517

            
518
    #[test]
519
1
    pub fn test_range_join() {
520
1
        assert_eq!(range!(1..3).join(&range!(2..4)), Some(range!(1..4)));
521
1
        assert_eq!(range!(1..3).join(&range!(3..4)), Some(range!(1..4)));
522
1
        assert_eq!(range!(1..3).join(&range!(4..5)), Some(range!(1..5)));
523
1
        assert_eq!(range!(..3).join(&range!(4..5)), Some(range!(..5)));
524
1
        assert_eq!(range!(1..3).join(&range!(4..)), Some(range!(1..)));
525
1
        assert_eq!(range!(4..).join(&range!(1..3)), Some(range!(1..)));
526
1
        assert_eq!(range!(..3).join(&range!(4..)), Some(Range::Unbounded));
527
1
        assert_eq!(range!(1..3).join(&range!(5..6)), None);
528
1
    }
529
}