1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
use std::{collections::BTreeMap, convert::TryFrom};

use thiserror::Error;

/// Error type for [`RangeSet`]
#[derive(Clone, Error, Debug, PartialEq, Eq)]
pub enum RangeSetError {
    /// Used when an invalid range is provided (e.g. End is before start).
    #[error("Invalid range provided")]
    InvalidRange,
    /// Used when an insertion would result in an overlap.
    #[error("Overlapping ranges")]
    RangeOverlap,
}

/// A set of non-overlapping ranges.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RangeSet<RK: Ord + Copy> {
    ranges: BTreeMap<RK, RK>, // Maps start -> end
}

impl<RK: Ord + Copy> RangeSet<RK> {
    /// Creates a new empty `RangeSet`.
    pub fn new() -> Self {
        Self {
            ranges: BTreeMap::new(),
        }
    }

    /// Inserts a new range `[start, end)`, ensuring no overlaps.
    pub fn insert(&mut self, start: RK, end: RK) -> Result<(), RangeSetError> {
        if start >= end {
            return Err(RangeSetError::InvalidRange);
        }

        // Find adjacent or overlapping ranges
        if let Some((&_prev_start, &prev_end)) = self.ranges.range(..=start).next_back() {
            if prev_end >= start {
                return Err(RangeSetError::RangeOverlap);
            }
        }
        if let Some((&next_start, _)) = self.ranges.range(start..).next() {
            if next_start < end {
                return Err(RangeSetError::RangeOverlap);
            }
        }

        self.ranges.insert(start, end);
        Ok(())
    }

    /// Checks if a value exists in any range.
    pub fn contains(&self, value: RK) -> bool {
        if let Some((&_start, &end)) = self.ranges.range(..=value).next_back() {
            return value < end;
        }
        false
    }

    /// Removes a range if it exists.
    pub fn remove(&mut self, start: RK, end: RK) -> bool {
        if let Some(range_end) = self.ranges.get(&start) {
            if *range_end == end {
                self.ranges.remove(&start);
                return true;
            }
        }

        false
    }

    /// Returns all stored ranges, sorted by key.
    pub fn iter(&self) -> impl Iterator<Item = (RK, RK)> + '_ {
        self.ranges.iter().map(|(&s, &e)| (s, e))
    }

    /// Returns the number of elements in the set.
    pub fn len(&self) -> usize {
        self.ranges.len()
    }

    /// Returns whether this set is empty.
    pub fn is_empty(&self) -> bool {
        self.ranges.is_empty()
    }
}

impl<RK: Copy + Ord> Default for RangeSet<RK> {
    fn default() -> Self {
        RangeSet::new()
    }
}

impl<RK: Copy + Ord> TryFrom<Vec<(RK, RK)>> for RangeSet<RK> {
    type Error = RangeSetError;

    fn try_from(value: Vec<(RK, RK)>) -> Result<Self, Self::Error> {
        let mut range_set = RangeSet::new();

        for range in value {
            range_set.insert(range.0, range.1)?;
        }

        Ok(range_set)
    }
}

#[cfg(test)]
mod tests {
    use std::convert::TryInto;

    use crate::rangeset::{RangeSet, RangeSetError};

    #[test]
    fn test_insertions() {
        let mut rs = RangeSet::new();

        assert_eq!(rs.insert(10, 20), Ok(()));
        assert_eq!(rs.insert(30, 40), Ok(()));
        assert_eq!(rs.insert(15, 25), Err(RangeSetError::RangeOverlap));
        assert!(rs.contains(15));
        assert!(!rs.contains(25));
    }

    #[test]
    fn test_removal() {
        let mut rs = RangeSet::new();

        assert_eq!(rs.insert(10, 20), Ok(()));
        assert_eq!(rs.insert(30, 40), Ok(()));
        assert!(!rs.remove(10, 11));
        assert!(rs.remove(10, 20));
        assert!(rs.remove(30, 40));
        assert!(!rs.remove(30, 40));
    }

    #[test]
    fn test_try_from_vec() {
        let range_set: RangeSet<u32> = vec![(0, 10), (100, 150)].try_into().unwrap();

        assert_eq!(range_set.len(), 2);

        let invalid_result: Result<RangeSet<u32>, RangeSetError> =
            vec![(10, 30), (11, 23)].try_into();

        assert_eq!(invalid_result, Err(RangeSetError::RangeOverlap))
    }
}