osb/utils/
interval_map.rs

1/// Data structure to associate keys of an interval type to a certain value
2pub struct IntervalMap<K, V> {
3    points: Vec<(K, Vec<V>)>,
4}
5
6impl<K, V> Default for IntervalMap<K, V> {
7    fn default() -> Self {
8        Self { points: Vec::new() }
9    }
10}
11
12use std::cmp::{Ord, Ordering::*};
13use std::ops::Range;
14
15impl<K, V> IntervalMap<K, V>
16where
17    K: Ord,
18    V: Clone,
19{
20    /// Initializes a `IntervalMap`
21    pub fn new() -> Self {
22        Self::default()
23    }
24
25    /// Adds a value to our `IntervalMap`.
26    ///
27    /// In the following example, our value is of integer type, but
28    /// it can be anything type that implements the trait `Clone`.
29    ///
30    /// Usage:
31    /// ```
32    /// use osb::utils::IntervalMap;
33    /// let mut interval_map = IntervalMap::new();
34    /// interval_map.push(10..50, 1);
35    /// ```
36    pub fn push(&mut self, range: Range<K>, value: V) {
37        let position = match self
38            .points
39            .binary_search_by(|&(ref point, _)| point.cmp(&range.start))
40        {
41            Ok(position) => position,
42            Err(position) => {
43                self.points.insert(
44                    position,
45                    (
46                        range.start,
47                        self.points
48                            .get(position.wrapping_sub(1))
49                            .map(|(_, values)| values.clone())
50                            .unwrap_or_default(),
51                    ),
52                );
53                position
54            }
55        };
56        for (i, point) in self.points.iter_mut().enumerate().skip(position) {
57            match point.0.cmp(&range.end) {
58                Less => point.1.push(value.clone()),
59                Equal => return,
60                Greater => {
61                    let mut new_point = (range.end, point.1.clone());
62                    new_point.1.push(value);
63                    self.points.insert(i, new_point);
64                    return;
65                }
66            }
67        }
68        self.points.push((range.end, Vec::new()))
69    }
70
71    /// Retrieve all of the values that is inside an interval
72    ///
73    /// Usage:
74    /// ```
75    /// use osb::utils::IntervalMap;
76    /// let mut interval_map = IntervalMap::new();
77    /// interval_map.push(10..50, 1);
78    /// interval_map.push(30..50, 42);
79    ///
80    /// let mut result1 = interval_map.get(&20);
81    /// assert_eq!(result1.next(), Some(&1));
82    /// assert_eq!(result1.next(), None);
83    ///
84    /// let mut result2 = interval_map.get(&40);
85    /// assert_eq!(result2.next(), Some(&1));
86    /// assert_eq!(result2.next(), Some(&42));
87    /// assert_eq!(result2.next(), None);
88    /// ```
89    pub fn get(&self, key: &K) -> std::slice::Iter<V> {
90        let index = match self
91            .points
92            .binary_search_by(|&(ref point, _)| point.cmp(key))
93        {
94            Err(index) => {
95                if index == 0 {
96                    return (&[]).iter();
97                } else {
98                    index - 1
99                }
100            }
101            Ok(index) => index,
102        };
103        self.points
104            .get(index)
105            .map(|point| &point.1[..])
106            .unwrap_or(&[])
107            .iter()
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::IntervalMap;
114
115    #[test]
116    fn basic() {
117        let mut interval_map = IntervalMap::new();
118
119        interval_map.push(10..50, 1);
120        interval_map.push(30..55, 2);
121
122        assert_eq!(interval_map.get(&0).next(), None);
123        assert_eq!(interval_map.get(&100).next(), None);
124
125        let mut result = interval_map.get(&20);
126        assert_eq!(result.next(), Some(&1));
127        assert_eq!(result.next(), None);
128
129        let mut result = interval_map.get(&53);
130        assert_eq!(result.next(), Some(&2));
131        assert_eq!(result.next(), None);
132
133        let mut result = interval_map.get(&40);
134        assert_eq!(result.next(), Some(&1));
135        assert_eq!(result.next(), Some(&2));
136        assert_eq!(result.next(), None);
137    }
138}
osb/utils/interval_map.rs

osb/utils/
interval_map.rs
