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

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

            
18
impl<A> Range<A> {
19
    /// Whether the range is **bounded** on either side. A bounded range may still be infinite.
20
    /// See also: [Range::is_finite].
21
132198245
    pub fn is_lower_or_upper_bounded(&self) -> bool {
22
132198245
        match &self {
23
            Range::Single(_)
24
            | Range::Bounded(_, _)
25
            | Range::UnboundedL(_)
26
132198244
            | Range::UnboundedR(_) => true,
27
1
            Range::Unbounded => false,
28
        }
29
132198245
    }
30

            
31
    /// Whether the range is **unbounded** on both sides.
32
    pub fn is_unbounded(&self) -> bool {
33
        !self.is_lower_or_upper_bounded()
34
    }
35

            
36
    /// Whether the range is **finite**. See also: [Range::is_lower_or_upper_bounded].
37
5
    pub fn is_finite(&self) -> bool {
38
5
        match &self {
39
2
            Range::Single(_) | Range::Bounded(_, _) => true,
40
3
            Range::Unbounded | Range::UnboundedL(_) | Range::UnboundedR(_) => false,
41
        }
42
5
    }
43
}
44

            
45
impl<A: Ord> Range<A> {
46
618213
    pub fn contains(&self, val: &A) -> bool {
47
618213
        match self {
48
2590
            Range::Single(x) => x == val,
49
615302
            Range::Bounded(x, y) => x <= val && val <= y,
50
            Range::UnboundedR(x) => x <= val,
51
            Range::UnboundedL(x) => val <= x,
52
321
            Range::Unbounded => true,
53
        }
54
618213
    }
55

            
56
    /// Returns the lower bound of the range, if it has one
57
18198696
    pub fn low(&self) -> Option<&A> {
58
18198696
        match self {
59
9664322
            Range::Single(a) => Some(a),
60
8534095
            Range::Bounded(a, _) => Some(a),
61
194
            Range::UnboundedR(a) => Some(a),
62
85
            Range::UnboundedL(_) | Range::Unbounded => None,
63
        }
64
18198696
    }
65

            
66
    /// Returns the upper bound of the range, if it has one
67
18197375
    pub fn high(&self) -> Option<&A> {
68
18197375
        match self {
69
9489227
            Range::Single(a) => Some(a),
70
8707870
            Range::Bounded(_, a) => Some(a),
71
28
            Range::UnboundedL(a) => Some(a),
72
250
            Range::UnboundedR(_) | Range::Unbounded => None,
73
        }
74
18197375
    }
75
}
76

            
77
impl<A: Ord + Clone> Range<A> {
78
    /// Create a new range with a lower and upper bound
79
1838993
    pub fn new(lo: Option<A>, hi: Option<A>) -> Range<A> {
80
1838993
        match (lo, hi) {
81
2
            (None, None) => Range::Unbounded,
82
3
            (Some(l), None) => Range::UnboundedR(l),
83
3
            (None, Some(r)) => Range::UnboundedL(r),
84
1838985
            (Some(l), Some(r)) => {
85
1838985
                if l == r {
86
164568
                    Range::Single(l)
87
                } else {
88
1674417
                    let min = Ord::min(&l, &r).clone();
89
1674417
                    let max = Ord::max(l, r);
90
1674417
                    Range::Bounded(min, max)
91
                }
92
            }
93
        }
94
1838993
    }
95

            
96
    /// Given a slice of ranges, create a single range that spans from the start
97
    /// of the leftmost range to the end of the rightmost range.
98
    /// An empty slice is considered equivalent to `Range::unbounded`.
99
201417
    pub fn spanning(rngs: &[Range<A>]) -> Range<A> {
100
201417
        if rngs.is_empty() {
101
1
            return Range::Unbounded;
102
201416
        }
103

            
104
201416
        let mut lo = rngs[0].low();
105
201416
        let mut hi = rngs[0].high();
106
216585
        for rng in rngs {
107
216585
            lo = match (lo, rng.low()) {
108
216581
                (Some(curr), Some(new)) => Some(curr.min(new)),
109
4
                _ => None,
110
            };
111
216585
            hi = match (hi, rng.high()) {
112
216583
                (Some(curr), Some(new)) => Some(curr.max(new)),
113
2
                _ => None,
114
            };
115
        }
116
201416
        Range::new(lo.cloned(), hi.cloned())
117
201417
    }
118
    /// Find the range such that:
119
    /// - the lower bound is the maximum of the lower bounds
120
    /// - the upper bound is the minimum of the upper bounds
121
    /// - **ranges must not be disjoint**
122
    ///
123
    /// * `DomainopError::ConflictingArgs`: if given disjoint ranges; e.g. (2..4) (6..8)
124
168
    pub fn minimal(rngs: &[Range<A>]) -> Result<Range<A>, DomainOpError> {
125
168
        if rngs.is_empty() {
126
            return Ok(Range::Unbounded);
127
168
        }
128
168
        let mut lo = rngs[0].low();
129
168
        let mut hi = rngs[0].high();
130
336
        for rng in rngs {
131
336
            lo = match (lo, rng.low()) {
132
272
                (Some(curr), Some(new)) => Some(curr.max(new)),
133
32
                (None, Some(new)) => Some(new),
134
                (Some(curr), None) => Some(curr),
135
32
                _ => None,
136
            };
137
336
            hi = match (hi, rng.high()) {
138
144
                (Some(curr), Some(new)) => Some(curr.min(new)),
139
96
                (None, Some(new)) => Some(new),
140
                (Some(curr), None) => Some(curr),
141
96
                _ => None,
142
            };
143
336
            if let (Some(l), Some(h)) = (lo, hi)
144
236
                && l > h
145
            {
146
8
                return Err(DomainOpError::ConflictingAttrs);
147
328
            }
148
        }
149
160
        Ok(Range::new(lo.cloned(), hi.cloned()))
150
168
    }
151
}
152

            
153
impl<A: Num + Ord + Clone> Range<A> {
154
7034
    pub fn length(&self) -> Option<A> {
155
7034
        match self {
156
2
            Range::Single(_) => Some(A::one()),
157
7027
            Range::Bounded(i, j) => Some(j.clone() - i.clone() + A::one()),
158
5
            Range::UnboundedR(_) | Range::UnboundedL(_) | Range::Unbounded => None,
159
        }
160
7034
    }
161

            
162
    /// Returns true if this interval overlaps another one, i.e. at least one
163
    /// number is part of both `self` and `other`
164
    /// E.g:
165
    /// - [0, 2] overlaps [2, 4]
166
    /// - [1, 3] overlaps [2, 4]
167
    /// - [4, 6] overlaps [2, 4]
168
3554743
    pub fn overlaps(&self, other: &Range<A>) -> bool {
169
3554743
        self.low()
170
3554743
            .is_none_or(|la| other.high().is_none_or(|rb| la <= rb))
171
3543133
            && self
172
3543133
                .high()
173
3543133
                .is_none_or(|ra| other.low().is_none_or(|lb| ra >= lb))
174
3554743
    }
175

            
176
    /// Returns true if this interval touches another one on the left
177
    /// E.g: [1, 2] touches_left  [3, 4]
178
3702602
    pub fn touches_left(&self, other: &Range<A>) -> bool {
179
3702602
        self.high().is_some_and(|ra| {
180
3702601
            let ra = ra.clone() + A::one();
181
3702601
            other.low().is_some_and(|lb| ra.eq(lb))
182
3702601
        })
183
3702602
    }
184

            
185
    /// Returns true if this interval touches another one on the right
186
    /// E.g: [3, 4] touches_right  [1, 2]
187
2175903
    pub fn touches_right(&self, other: &Range<A>) -> bool {
188
2175903
        self.low().is_some_and(|la| {
189
2175902
            let la = la.clone() - A::one();
190
2175902
            other.high().is_some_and(|rb| la.eq(rb))
191
2175902
        })
192
2175903
    }
193

            
194
    /// Returns true if this interval overlaps or touches another one
195
    /// E.g:
196
    /// - [1, 3] joins [4, 6]
197
    /// - [2, 4] joins [4, 6]
198
    /// - [3, 5] joins [4, 6]
199
    /// - [6, 8] joins [4, 6]
200
    /// - [7, 8] joins [4, 6]
201
3702596
    pub fn joins(&self, other: &Range<A>) -> bool {
202
3702596
        self.touches_left(other) || self.overlaps(other) || self.touches_right(other)
203
3702596
    }
204

            
205
    /// Returns true if this interval is strictly before another one
206
7
    pub fn is_before(&self, other: &Range<A>) -> bool {
207
7
        self.high()
208
7
            .is_some_and(|ra| other.low().is_some_and(|lb| ra < &(lb.clone() - A::one())))
209
7
    }
210

            
211
    /// Returns true if this interval is strictly after another one
212
6
    pub fn is_after(&self, other: &Range<A>) -> bool {
213
6
        self.low()
214
6
            .is_some_and(|la| other.high().is_some_and(|rb| la > &(rb.clone() + A::one())))
215
6
    }
216

            
217
    /// If the two ranges join, return a new range which spans both
218
3702596
    pub fn join(&self, other: &Range<A>) -> Option<Range<A>> {
219
3702596
        if self.joins(other) {
220
1537229
            let lo = Ord::min(self.low(), other.low());
221
1537229
            let hi = match (self.high(), other.high()) {
222
1537226
                (Some(a), Some(b)) => Some(Ord::max(a, b)),
223
3
                _ => None,
224
            };
225
1537229
            return Some(Range::new(lo.cloned(), hi.cloned()));
226
2165367
        }
227
2165367
        None
228
3702596
    }
229

            
230
    /// Merge all joining ranges in the list, and return a new vec of disjoint ranges.
231
    /// E.g:
232
    /// ```ignore
233
    /// [(2..3), (4), (..1), (6..8)] -> [(..4), (6..8)]
234
    /// ```
235
    ///
236
    /// # Performance
237
    /// Currently uses a naive O(n^2) algorithm.
238
    /// A more optimal approach based on interval trees is planned.
239
6697801
    pub fn squeeze(rngs: &[Range<A>]) -> Vec<Range<A>> {
240
6697801
        let mut ans = Vec::from(rngs);
241

            
242
6697801
        if ans.is_empty() {
243
264
            return ans;
244
6697537
        }
245

            
246
        loop {
247
8234759
            let mut merged = false;
248

            
249
            // Check every pair of ranges and join them if possible
250
8658788
            'outer: for i in 0..ans.len() {
251
8658788
                for j in (i + 1)..ans.len() {
252
3702588
                    if let Some(joined) = ans[i].join(&ans[j]) {
253
1537222
                        ans[i] = joined;
254
                        // Safe to delete here because we restart the outer loop immediately
255
1537222
                        ans.remove(j);
256
1537222
                        merged = true;
257
1537222
                        break 'outer;
258
2165366
                    }
259
                }
260
            }
261

            
262
            // If no merges occurred, we're done
263
8234759
            if !merged {
264
6697537
                break;
265
1537222
            }
266
        }
267

            
268
6697537
        ans
269
6697801
    }
270

            
271
    /// If this range is bounded, returns a lazy iterator over all values within the range.
272
    /// Otherwise, returns None.
273
1108056
    pub fn iter(&self) -> Option<RangeIterator<A>> {
274
1108056
        match self {
275
17524
            Range::Single(val) => Some(RangeIterator::Single(Some(val.clone()))),
276
1090532
            Range::Bounded(start, end) => Some(RangeIterator::Bounded {
277
1090532
                current: start.clone(),
278
1090532
                end: end.clone(),
279
1090532
            }),
280
            Range::UnboundedL(_) | Range::UnboundedR(_) | Range::Unbounded => None,
281
        }
282
1108056
    }
283

            
284
    /// Lazily iterate all values in a list of ranges
285
1076
    pub fn values(rngs: &[Range<A>]) -> Option<impl Iterator<Item = A>> {
286
1076
        let itrs = rngs
287
1076
            .iter()
288
1076
            .map(Range::iter)
289
1076
            .collect::<Option<Vec<RangeIterator<A>>>>()?;
290
1076
        Some(itrs.into_iter().flatten())
291
1076
    }
292

            
293
    /// True if the list of ranges is contiguous
294
    pub fn is_contiguous(rngs: &[Range<A>]) -> bool {
295
        Self::squeeze(rngs).len() <= 1
296
    }
297

            
298
    /// Lowest value from a sequence of ranges
299
    pub fn low_of(rngs: &[Range<A>]) -> Option<&A> {
300
        let mut low = rngs.first()?.low()?;
301
        for rng in rngs {
302
            low = Ord::min(low, rng.low()?);
303
        }
304
        Some(low)
305
    }
306

            
307
    /// Lowest value from a sequence of ranges
308
    pub fn high_of(rngs: &[Range<A>]) -> Option<&A> {
309
        let mut hi = rngs.first()?.high()?;
310
        for rng in rngs {
311
            hi = Ord::min(hi, rng.low()?);
312
        }
313
        Some(hi)
314
    }
315

            
316
    /// Total number of values across a slice of ranges.
317
    /// Returns `None` if any range is unbounded.
318
    pub fn total_length(rngs: &[Range<A>]) -> Option<A> {
319
        rngs.iter()
320
            .try_fold(A::zero(), |acc, r| Some(acc + r.length()?))
321
    }
322
}
323

            
324
/// Iterator for Range<A> that yields values lazily
325
pub enum RangeIterator<A> {
326
    Single(Option<A>),
327
    Bounded { current: A, end: A },
328
}
329

            
330
impl<A: Num + Ord + Clone> Iterator for RangeIterator<A> {
331
    type Item = A;
332

            
333
5539448
    fn next(&mut self) -> Option<Self::Item> {
334
5539448
        match self {
335
35048
            RangeIterator::Single(val) => val.take(),
336
5504400
            RangeIterator::Bounded { current, end } => {
337
5504400
                if current > end {
338
1090532
                    return None;
339
4413868
                }
340

            
341
4413868
                let result = current.clone();
342
4413868
                *current = current.clone() + A::one();
343

            
344
4413868
                Some(result)
345
            }
346
        }
347
5539448
    }
348
}
349

            
350
impl<A: Display> Display for Range<A> {
351
132198240
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
352
132198240
        match self {
353
5123888
            Range::Single(i) => write!(f, "{i}"),
354
124474292
            Range::Bounded(i, j) => write!(f, "{i}..{j}"),
355
2600000
            Range::UnboundedR(i) => write!(f, "{i}.."),
356
60
            Range::UnboundedL(i) => write!(f, "..{i}"),
357
            Range::Unbounded => write!(f, ""),
358
        }
359
132198240
    }
360
}
361

            
362
#[allow(unused_imports)]
363
mod test {
364
    use super::*;
365
    use crate::range;
366

            
367
    #[test]
368
1
    pub fn test_range_macros() {
369
1
        assert_eq!(range!(1..3), Range::Bounded(1, 3));
370
1
        assert_eq!(range!(1..), Range::UnboundedR(1));
371
1
        assert_eq!(range!(..3), Range::UnboundedL(3));
372
1
        assert_eq!(range!(1), Range::Single(1));
373
1
    }
374

            
375
    #[test]
376
1
    pub fn test_range_low() {
377
1
        assert_eq!(range!(1..3).low(), Some(&1));
378
1
        assert_eq!(range!(1..).low(), Some(&1));
379
1
        assert_eq!(range!(1).low(), Some(&1));
380
1
        assert_eq!(range!(..3).low(), None);
381
1
        assert_eq!(Range::<Int>::Unbounded.low(), None);
382
1
    }
383

            
384
    #[test]
385
1
    pub fn test_range_high() {
386
1
        assert_eq!(range!(1..3).high(), Some(&3));
387
1
        assert_eq!(range!(1..).high(), None);
388
1
        assert_eq!(range!(1).high(), Some(&1));
389
1
        assert_eq!(range!(..3).high(), Some(&3));
390
1
        assert_eq!(Range::<Int>::Unbounded.high(), None);
391
1
    }
392

            
393
    #[test]
394
1
    pub fn test_range_is_finite() {
395
1
        assert!(range!(1..3).is_finite());
396
1
        assert!(range!(1).is_finite());
397
1
        assert!(!range!(1..).is_finite());
398
1
        assert!(!range!(..3).is_finite());
399
1
        assert!(!Range::<Int>::Unbounded.is_finite());
400
1
    }
401

            
402
    #[test]
403
1
    pub fn test_range_bounded() {
404
1
        assert!(range!(1..3).is_lower_or_upper_bounded());
405
1
        assert!(range!(1).is_lower_or_upper_bounded());
406
1
        assert!(range!(1..).is_lower_or_upper_bounded());
407
1
        assert!(range!(..3).is_lower_or_upper_bounded());
408
1
        assert!(!Range::<Int>::Unbounded.is_lower_or_upper_bounded());
409
1
    }
410

            
411
    #[test]
412
1
    pub fn test_range_length() {
413
1
        assert_eq!(range!(1..3).length(), Some(3));
414
1
        assert_eq!(range!(1).length(), Some(1));
415
1
        assert_eq!(range!(1..).length(), None);
416
1
        assert_eq!(range!(..3).length(), None);
417
1
        assert_eq!(Range::<Int>::Unbounded.length(), None);
418
1
    }
419

            
420
    #[test]
421
1
    pub fn test_range_contains_value() {
422
1
        assert!(range!(1..3).contains(&2));
423
1
        assert!(!range!(1..3).contains(&4));
424
1
        assert!(range!(1).contains(&1));
425
1
        assert!(!range!(1).contains(&2));
426
1
        assert!(Range::Unbounded.contains(&42));
427
1
    }
428

            
429
    #[test]
430
1
    pub fn test_range_overlaps() {
431
1
        assert!(range!(1..3).overlaps(&range!(2..4)));
432
1
        assert!(range!(1..3).overlaps(&range!(3..5)));
433
1
        assert!(!range!(1..3).overlaps(&range!(4..6)));
434
1
        assert!(Range::Unbounded.overlaps(&range!(1..3)));
435
1
    }
436

            
437
    #[test]
438
1
    pub fn test_range_touches_left() {
439
1
        assert!(range!(1..2).touches_left(&range!(3..4)));
440
1
        assert!(range!(1..2).touches_left(&range!(3)));
441
1
        assert!(range!(-5..-4).touches_left(&range!(-3..2)));
442
1
        assert!(!range!(1..2).touches_left(&range!(4..5)));
443
1
        assert!(!range!(1..2).touches_left(&range!(2..3)));
444
1
        assert!(!range!(3..4).touches_left(&range!(1..2)));
445
1
    }
446

            
447
    #[test]
448
1
    pub fn test_range_touches_right() {
449
1
        assert!(range!(3..4).touches_right(&range!(1..2)));
450
1
        assert!(range!(3).touches_right(&range!(1..2)));
451
1
        assert!(range!(0..1).touches_right(&range!(-2..-1)));
452
1
        assert!(!range!(1..2).touches_right(&range!(3..4)));
453
1
        assert!(!range!(2..3).touches_right(&range!(1..2)));
454
1
        assert!(!range!(1..2).touches_right(&range!(1..2)));
455
1
    }
456

            
457
    #[test]
458
1
    pub fn test_range_is_before() {
459
1
        assert!(range!(1..2).is_before(&range!(4..5)));
460
1
        assert!(range!(1..2).is_before(&range!(4..)));
461
1
        assert!(!range!(1..2).is_before(&range!(3..)));
462
1
        assert!(!range!(1..2).is_before(&range!(..4)));
463
1
        assert!(!range!(1..2).is_before(&range!(2..4)));
464
1
        assert!(!range!(3..4).is_before(&range!(1..2)));
465
1
        assert!(!range!(1..2).is_before(&Range::Unbounded));
466
1
    }
467

            
468
    #[test]
469
1
    pub fn test_range_is_after() {
470
1
        assert!(range!(5..6).is_after(&range!(1..2)));
471
1
        assert!(range!(4..5).is_after(&range!(..2)));
472
1
        assert!(!range!(4..5).is_after(&range!(..3)));
473
1
        assert!(!range!(2..3).is_after(&range!(1..2)));
474
1
        assert!(!range!(1..2).is_after(&range!(3..4)));
475
1
        assert!(!range!(1..2).is_after(&Range::Unbounded));
476
1
    }
477

            
478
    #[test]
479
1
    pub fn test_range_squeeze() {
480
1
        let input = vec![range!(2..3), range!(4), range!(..1), range!(6..8)];
481
1
        let squeezed = Range::squeeze(&input);
482
1
        assert_eq!(squeezed, vec![range!(..4), range!(6..8)]);
483
1
    }
484

            
485
    #[test]
486
1
    pub fn test_range_spanning() {
487
1
        assert_eq!(Range::<Int>::spanning(&[]), Range::Unbounded);
488
1
        assert_eq!(Range::spanning(&[range!(1..2), range!(4..5)]), range!(1..5));
489
1
        assert_eq!(
490
1
            Range::spanning(&[range!(..0), range!(2..4)]),
491
            Range::UnboundedL(4)
492
        );
493
1
        assert_eq!(
494
1
            Range::spanning(&[range!(0), range!(2..3), range!(5..)]),
495
            Range::UnboundedR(0)
496
        );
497
1
        assert_eq!(
498
1
            Range::spanning(&[range!(..0), range!(2..)]),
499
            Range::Unbounded
500
        );
501
1
    }
502

            
503
    #[test]
504
1
    pub fn test_range_join() {
505
1
        assert_eq!(range!(1..3).join(&range!(2..4)), Some(range!(1..4)));
506
1
        assert_eq!(range!(1..3).join(&range!(3..4)), Some(range!(1..4)));
507
1
        assert_eq!(range!(1..3).join(&range!(4..5)), Some(range!(1..5)));
508
1
        assert_eq!(range!(..3).join(&range!(4..5)), Some(range!(..5)));
509
1
        assert_eq!(range!(1..3).join(&range!(4..)), Some(range!(1..)));
510
1
        assert_eq!(range!(4..).join(&range!(1..3)), Some(range!(1..)));
511
1
        assert_eq!(range!(..3).join(&range!(4..)), Some(Range::Unbounded));
512
1
        assert_eq!(range!(1..3).join(&range!(5..6)), None);
513
1
    }
514
}