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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
216
    /// If the two ranges join, return a new range which spans both
217
3360588
    pub fn join(&self, other: &Range<A>) -> Option<Range<A>> {
218
3360588
        if self.joins(other) {
219
1350052
            let lo = Ord::min(self.low(), other.low());
220
1350052
            let hi = Ord::max(self.high(), other.high());
221
1350052
            return Some(Range::new(lo.cloned(), hi.cloned()));
222
2010536
        }
223
2010536
        None
224
3360588
    }
225

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

            
238
6279685
        if ans.is_empty() {
239
192
            return ans;
240
6279493
        }
241

            
242
        loop {
243
7629545
            let mut merged = false;
244

            
245
            // Check every pair of ranges and join them if possible
246
7941472
            'outer: for i in 0..ans.len() {
247
7941472
                for j in (i + 1)..ans.len() {
248
3360588
                    if let Some(joined) = ans[i].join(&ans[j]) {
249
1350052
                        ans[i] = joined;
250
                        // Safe to delete here because we restart the outer loop immediately
251
1350052
                        ans.remove(j);
252
1350052
                        merged = true;
253
1350052
                        break 'outer;
254
2010536
                    }
255
                }
256
            }
257

            
258
            // If no merges occurred, we're done
259
7629545
            if !merged {
260
6279493
                break;
261
1350052
            }
262
        }
263

            
264
6279493
        ans
265
6279685
    }
266

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

            
280
1058
    pub fn values(rngs: &[Range<A>]) -> Option<impl Iterator<Item = A>> {
281
1058
        let itrs = rngs
282
1058
            .iter()
283
1058
            .map(Range::iter)
284
1058
            .collect::<Option<Vec<RangeIterator<A>>>>()?;
285
1058
        Some(itrs.into_iter().flatten())
286
1058
    }
287
}
288

            
289
/// Iterator for Range<A> that yields values lazily
290
pub enum RangeIterator<A> {
291
    Single(Option<A>),
292
    Bounded { current: A, end: A },
293
}
294

            
295
impl<A: Num + Ord + Clone> Iterator for RangeIterator<A> {
296
    type Item = A;
297

            
298
4710872
    fn next(&mut self) -> Option<Self::Item> {
299
4710872
        match self {
300
34804
            RangeIterator::Single(val) => val.take(),
301
4676068
            RangeIterator::Bounded { current, end } => {
302
4676068
                if current > end {
303
885550
                    return None;
304
3790518
                }
305

            
306
3790518
                let result = current.clone();
307
3790518
                *current = current.clone() + A::one();
308

            
309
3790518
                Some(result)
310
            }
311
        }
312
4710872
    }
313
}
314

            
315
impl<A: Display> Display for Range<A> {
316
131082438
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
317
131082438
        match self {
318
4904322
            Range::Single(i) => write!(f, "{i}"),
319
124132568
            Range::Bounded(i, j) => write!(f, "{i}..{j}"),
320
2045488
            Range::UnboundedR(i) => write!(f, "{i}.."),
321
60
            Range::UnboundedL(i) => write!(f, "..{i}"),
322
            Range::Unbounded => write!(f, ""),
323
        }
324
131082438
    }
325
}
326

            
327
#[allow(unused_imports)]
328
mod test {
329
    use super::*;
330
    use crate::range;
331

            
332
    #[test]
333
1
    pub fn test_range_macros() {
334
1
        assert_eq!(range!(1..3), Range::Bounded(1, 3));
335
1
        assert_eq!(range!(1..), Range::UnboundedR(1));
336
1
        assert_eq!(range!(..3), Range::UnboundedL(3));
337
1
        assert_eq!(range!(1), Range::Single(1));
338
1
    }
339

            
340
    #[test]
341
1
    pub fn test_range_low() {
342
1
        assert_eq!(range!(1..3).low(), Some(&1));
343
1
        assert_eq!(range!(1..).low(), Some(&1));
344
1
        assert_eq!(range!(1).low(), Some(&1));
345
1
        assert_eq!(range!(..3).low(), None);
346
1
        assert_eq!(Range::<Int>::Unbounded.low(), None);
347
1
    }
348

            
349
    #[test]
350
1
    pub fn test_range_high() {
351
1
        assert_eq!(range!(1..3).high(), Some(&3));
352
1
        assert_eq!(range!(1..).high(), None);
353
1
        assert_eq!(range!(1).high(), Some(&1));
354
1
        assert_eq!(range!(..3).high(), Some(&3));
355
1
        assert_eq!(Range::<Int>::Unbounded.high(), None);
356
1
    }
357

            
358
    #[test]
359
1
    pub fn test_range_is_finite() {
360
1
        assert!(range!(1..3).is_finite());
361
1
        assert!(range!(1).is_finite());
362
1
        assert!(!range!(1..).is_finite());
363
1
        assert!(!range!(..3).is_finite());
364
1
        assert!(!Range::<Int>::Unbounded.is_finite());
365
1
    }
366

            
367
    #[test]
368
1
    pub fn test_range_bounded() {
369
1
        assert!(range!(1..3).is_lower_or_upper_bounded());
370
1
        assert!(range!(1).is_lower_or_upper_bounded());
371
1
        assert!(range!(1..).is_lower_or_upper_bounded());
372
1
        assert!(range!(..3).is_lower_or_upper_bounded());
373
1
        assert!(!Range::<Int>::Unbounded.is_lower_or_upper_bounded());
374
1
    }
375

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

            
385
    #[test]
386
1
    pub fn test_range_contains_value() {
387
1
        assert!(range!(1..3).contains(&2));
388
1
        assert!(!range!(1..3).contains(&4));
389
1
        assert!(range!(1).contains(&1));
390
1
        assert!(!range!(1).contains(&2));
391
1
        assert!(Range::Unbounded.contains(&42));
392
1
    }
393

            
394
    #[test]
395
1
    pub fn test_range_overlaps() {
396
1
        assert!(range!(1..3).overlaps(&range!(2..4)));
397
1
        assert!(range!(1..3).overlaps(&range!(3..5)));
398
1
        assert!(!range!(1..3).overlaps(&range!(4..6)));
399
1
        assert!(Range::Unbounded.overlaps(&range!(1..3)));
400
1
    }
401

            
402
    #[test]
403
1
    pub fn test_range_touches_left() {
404
1
        assert!(range!(1..2).touches_left(&range!(3..4)));
405
1
        assert!(range!(1..2).touches_left(&range!(3)));
406
1
        assert!(range!(-5..-4).touches_left(&range!(-3..2)));
407
1
        assert!(!range!(1..2).touches_left(&range!(4..5)));
408
1
        assert!(!range!(1..2).touches_left(&range!(2..3)));
409
1
        assert!(!range!(3..4).touches_left(&range!(1..2)));
410
1
    }
411

            
412
    #[test]
413
1
    pub fn test_range_touches_right() {
414
1
        assert!(range!(3..4).touches_right(&range!(1..2)));
415
1
        assert!(range!(3).touches_right(&range!(1..2)));
416
1
        assert!(range!(0..1).touches_right(&range!(-2..-1)));
417
1
        assert!(!range!(1..2).touches_right(&range!(3..4)));
418
1
        assert!(!range!(2..3).touches_right(&range!(1..2)));
419
1
        assert!(!range!(1..2).touches_right(&range!(1..2)));
420
1
    }
421

            
422
    #[test]
423
1
    pub fn test_range_is_before() {
424
1
        assert!(range!(1..2).is_before(&range!(4..5)));
425
1
        assert!(range!(1..2).is_before(&range!(4..)));
426
1
        assert!(!range!(1..2).is_before(&range!(3..)));
427
1
        assert!(!range!(1..2).is_before(&range!(..4)));
428
1
        assert!(!range!(1..2).is_before(&range!(2..4)));
429
1
        assert!(!range!(3..4).is_before(&range!(1..2)));
430
1
        assert!(!range!(1..2).is_before(&Range::Unbounded));
431
1
    }
432

            
433
    #[test]
434
1
    pub fn test_range_is_after() {
435
1
        assert!(range!(5..6).is_after(&range!(1..2)));
436
1
        assert!(range!(4..5).is_after(&range!(..2)));
437
1
        assert!(!range!(4..5).is_after(&range!(..3)));
438
1
        assert!(!range!(2..3).is_after(&range!(1..2)));
439
1
        assert!(!range!(1..2).is_after(&range!(3..4)));
440
1
        assert!(!range!(1..2).is_after(&Range::Unbounded));
441
1
    }
442

            
443
    #[test]
444
1
    pub fn test_range_squeeze() {
445
1
        let input = vec![range!(2..3), range!(4), range!(..1), range!(6..8)];
446
1
        let squeezed = Range::squeeze(&input);
447
1
        assert_eq!(squeezed, vec![range!(..4), range!(6..8)]);
448
1
    }
449

            
450
    #[test]
451
1
    pub fn test_range_spanning() {
452
1
        assert_eq!(Range::<Int>::spanning(&[]), Range::Unbounded);
453
1
        assert_eq!(Range::spanning(&[range!(1..2), range!(4..5)]), range!(1..5));
454
1
        assert_eq!(
455
1
            Range::spanning(&[range!(..0), range!(2..4)]),
456
            Range::UnboundedL(4)
457
        );
458
1
        assert_eq!(
459
1
            Range::spanning(&[range!(0), range!(2..3), range!(5..)]),
460
            Range::UnboundedR(0)
461
        );
462
1
        assert_eq!(
463
1
            Range::spanning(&[range!(..0), range!(2..)]),
464
            Range::Unbounded
465
        );
466
1
    }
467
}