range_map/
classic.rs

1// Copyright 2021 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use std::borrow::Borrow;
6use std::cmp::Ordering;
7use std::collections::BTreeMap;
8use std::ops::Range;
9
10/// Keys for the map inside RangeMap.
11///
12/// This object holds a Range but implements the ordering traits according to
13/// the start of the range. Using this struct lets us store both ends of the
14/// range in the BTreeMap and recover ranges by querying for their start point.
15#[derive(Clone, Debug)]
16struct RangeStart<T> {
17    range: Range<T>,
18}
19
20impl<T: Copy> RangeStart<T> {
21    /// Wrap the given range in a RangeStart.
22    ///
23    /// Used in the BTreeMap to order the entries by the start of the range but
24    /// also remember the end of the range.
25    fn new(range: Range<T>) -> Self {
26        RangeStart { range }
27    }
28
29    /// An empty range with both endpoints at the start.
30    ///
31    /// Used for queries into the BTreeMap, but never stored in the BTreeMap.
32    fn from_point(point: T) -> Self {
33        RangeStart { range: Range { start: point, end: point } }
34    }
35}
36
37/// PartialEq according to the start of the Range.
38impl<T: PartialEq> PartialEq for RangeStart<T> {
39    fn eq(&self, other: &Self) -> bool {
40        self.range.start.eq(&other.range.start)
41    }
42}
43
44/// Eq according to the start of the Range.
45impl<T: Eq> Eq for RangeStart<T> {}
46
47/// PartialOrd according to the start of the Range.
48impl<T: Ord> PartialOrd for RangeStart<T> {
49    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
50        Some(self.cmp(other))
51    }
52}
53
54/// Ord according to the start of the Range.
55impl<T: Ord> Ord for RangeStart<T> {
56    fn cmp(&self, other: &Self) -> Ordering {
57        self.range.start.cmp(&other.range.start)
58    }
59}
60
61/// A map from a range of keys to values.
62///
63/// At any given time, the map contains a set of non-overlapping, non-empty
64/// ranges of type K that are associated with values of type V.
65///
66/// A given range can be split into two separate ranges if some of the
67/// intermediate values are removed from the map of if another value is
68/// inserted over the intermediate values. When that happens, the value
69/// for the split range is cloned using the Copy trait.
70///
71/// Adjacent ranges are not merged. Even if the value is "the same" (for some
72/// definition of "the same"), the ranges are kept separately.
73///
74/// Querying a point in the map returns not only the value stored at that point
75/// but also the range that value occupies in the map.
76#[derive(Debug)]
77pub struct RangeMap<K: Ord + Clone, V: Clone + Eq> {
78    map: BTreeMap<RangeStart<K>, V>,
79}
80
81impl<K, V> Default for RangeMap<K, V>
82where
83    K: Ord + Copy,
84    V: Clone + Eq,
85{
86    /// By default, a RangeMap is empty.
87    fn default() -> Self {
88        Self { map: Default::default() }
89    }
90}
91
92impl<K, V> RangeMap<K, V>
93where
94    K: Ord + Copy,
95    V: Clone + Eq,
96{
97    /// Returns the range (and associated value) that contains the given point,
98    /// if any.
99    ///
100    /// At most one range and value can exist at a given point because the
101    /// ranges in the map are non-overlapping.
102    ///
103    /// Empty ranges do not contain any points and therefore cannot be found
104    /// using this method. Rather than being stored in the map, values
105    /// associated with empty ranges are dropped.
106    pub fn get(&self, point: K) -> Option<(&Range<K>, &V)> {
107        self.map
108            .range(..=RangeStart::from_point(point))
109            .next_back()
110            .filter(|(k, _)| k.range.contains(&point))
111            .map(|(k, v)| (&k.range, v))
112    }
113
114    /// Inserts a range with the given value.
115    ///
116    /// The keys included in the given range are now associated with the given
117    /// value. If those keys were previously associated with another value,
118    /// are no longer associated with that previous value.
119    ///
120    /// This method can cause one or more values in the map to be dropped if
121    /// the all of the keys associated with those values are contained within
122    /// the given range.
123    ///
124    /// If the inserted range is directly adjacent to another range with an equal value, the
125    /// inserted range will be merged with the adjacent ranges.
126    pub fn insert(&mut self, mut range: Range<K>, value: V) {
127        if range.end <= range.start {
128            return;
129        }
130        self.remove(range.clone());
131
132        // Check for a range directly before this one. If it exists, it will be the last range with
133        // start < range.start.
134        if let Some((prev_range, prev_value)) =
135            self.map.range(..RangeStart::from_point(range.start)).next_back()
136        {
137            let prev_range = &prev_range.range;
138            if prev_range.end == range.start && &value == prev_value {
139                range.start = prev_range.start.clone();
140                self.remove_exact_range(prev_range.clone());
141            }
142        }
143        // Check for a range directly after. If it exists, we can look it up by exact start value
144        // of range.end.
145        if let Some((next_range, next_value)) =
146            self.map.get_key_value(&RangeStart::from_point(range.end))
147        {
148            let next_range = &next_range.range;
149            if next_range.start == range.end && &value == next_value {
150                range.end = next_range.end.clone();
151                self.remove_exact_range(next_range.clone());
152            }
153        }
154
155        self.insert_into_empty_range(range, value);
156    }
157
158    /// Remove the given range from the map.
159    ///
160    /// The keys included in the given range are no longer associated with any
161    /// values.
162    ///
163    /// This method can cause one or more values in the map to be dropped if all of the keys
164    /// associated with those values are contained within the given range.
165    ///
166    /// Returns any removed values.
167    pub fn remove(&mut self, range: Range<K>) -> Vec<V> {
168        let mut removed_values = vec![];
169        // If the given range is empty, there is nothing to do.
170        if range.end <= range.start {
171            return removed_values;
172        }
173
174        // Find the range (if any) that intersects the start of range.
175        //
176        // There can be at most one such range because we maintain the
177        // invariant that the ranges stored in the map are non-overlapping.
178        if let Some((old_range, v)) =
179            self.get(range.start).map(|(range, v)| (range.clone(), v.clone()))
180        {
181            // Remove that range from the map.
182            if let Some(value) = self.remove_exact_range(old_range.clone()) {
183                removed_values.push(value);
184            }
185
186            // If the removed range extends after the end of the given range,
187            // re-insert the part of the old range that extends beyond the end
188            // of the given range.
189            if old_range.end > range.end {
190                self.insert_into_empty_range(range.end.clone()..old_range.end, v.clone());
191            }
192
193            // If the removed range extends before the start of the given
194            // range, re-insert the part of the old range that extends before
195            // the start of the given range.
196            if old_range.start < range.start {
197                self.insert_into_empty_range(old_range.start..range.start.clone(), v);
198            }
199
200            // Notice that we can end up splitting the old range into two
201            // separate ranges if the old range extends both beyond the given
202            // range and before the given range.
203        }
204
205        // Find the range (if any) that intersects the end of range.
206        //
207        // There can be at most one such range because we maintain the
208        // invariant that the ranges stored in the map are non-overlapping.
209        //
210        // We exclude the end of the given range because a range that starts
211        // exactly at the end of the given range does not overalp the given
212        // range.
213        if let Some((old_range, v)) = self
214            .map
215            .range(RangeStart::from_point(range.start)..RangeStart::from_point(range.end))
216            .next_back()
217            .filter(|(k, _)| k.range.contains(&range.end))
218            .map(|(k, v)| (k.range.clone(), v.clone()))
219        {
220            // Remove that range from the map.
221            if let Some(value) = self.remove_exact_range(old_range.clone()) {
222                removed_values.push(value);
223            }
224
225            // If the removed range extends after the end of the given range,
226            // re-insert the part of the old range that extends beyond the end
227            // of the given range.
228            if old_range.end > range.end {
229                self.insert_into_empty_range(range.end.clone()..old_range.end, v);
230            }
231        }
232
233        // Remove any remaining ranges that are contained within the range.
234        //
235        // These ranges cannot possibly extend beyond the given range because
236        // we will have already removed them from the map at this point.
237        //
238        // We collect the doomed keys into a Vec to avoid mutating the map
239        // during the iteration.
240        let doomed: Vec<_> = self
241            .map
242            .range(RangeStart::from_point(range.start)..RangeStart::from_point(range.end))
243            .map(|(k, _)| k.clone())
244            .collect();
245
246        for key in &doomed {
247            if let Some(removed_value) = self.map.remove(key) {
248                removed_values.push(removed_value);
249            }
250        }
251        removed_values
252    }
253
254    pub fn is_empty(&self) -> bool {
255        self.map.is_empty()
256    }
257
258    /// Iterate over the ranges in the map.
259    pub fn iter(&self) -> impl Iterator<Item = (&Range<K>, &V)> {
260        self.map.iter().map(|(k, value)| (&k.range, value))
261    }
262
263    /// Iterate over the ranges in the map, starting at the first range starting after or at the given point.
264    pub fn iter_starting_at(&self, point: K) -> impl Iterator<Item = (&Range<K>, &V)> {
265        self.map.range(RangeStart::from_point(point)..).map(|(k, value)| (&k.range, value))
266    }
267
268    /// Iterate over the ranges in the map, starting at the last range starting before or at the given point.
269    pub fn iter_ending_at(&self, point: K) -> impl DoubleEndedIterator<Item = (&Range<K>, &V)> {
270        self.map.range(..RangeStart::from_point(point)).map(|(k, value)| (&k.range, value))
271    }
272
273    /// Iterate over the ranges in the map that intersect the requested range.
274    pub fn intersection<R>(&self, range: R) -> impl DoubleEndedIterator<Item = (&Range<K>, &V)>
275    where
276        R: Borrow<Range<K>>,
277    {
278        let range = range.borrow();
279        let start = self.get(range.start).map(|(r, _)| r.start).unwrap_or(range.start);
280        self.map
281            .range(RangeStart::from_point(start)..RangeStart::from_point(range.end))
282            .map(|(k, value)| (&k.range, value))
283    }
284
285    /// Returns the final range in the map.
286    pub fn last_range(&self) -> Option<Range<K>> {
287        if let Some((r, _)) = self.map.last_key_value() {
288            Some(r.range.clone())
289        } else {
290            None
291        }
292    }
293
294    /// Associate the keys in the given range with the given value.
295    ///
296    /// Callers must ensure that the keys in the given range are not already
297    /// associated with any values.
298    fn insert_into_empty_range(&mut self, range: Range<K>, value: V) {
299        self.map.insert(RangeStart::new(range), value);
300    }
301
302    /// Remove the given range from the map.
303    ///
304    /// Callers must ensure that the exact range provided as an argument is
305    /// contained in the map.
306    fn remove_exact_range(&mut self, range: Range<K>) -> Option<V> {
307        self.map.remove(&RangeStart::new(range))
308    }
309}
310
311#[cfg(test)]
312mod test {
313    use super::*;
314
315    #[::fuchsia::test]
316    fn test_empty() {
317        let mut map = RangeMap::<u32, i32>::default();
318
319        assert!(map.get(12).is_none());
320        map.remove(10..34);
321        // This is a test to make sure we can handle reversed ranges
322        #[allow(clippy::reversed_empty_ranges)]
323        map.remove(34..10);
324    }
325
326    #[::fuchsia::test]
327    fn test_insert_into_empty() {
328        let mut map = RangeMap::<u32, i32>::default();
329
330        map.insert(10..34, -14);
331
332        assert_eq!((&(10..34), &-14), map.get(12).unwrap());
333        assert_eq!((&(10..34), &-14), map.get(10).unwrap());
334        assert!(map.get(9).is_none());
335        assert_eq!((&(10..34), &-14), map.get(33).unwrap());
336        assert!(map.get(34).is_none());
337    }
338
339    #[::fuchsia::test]
340    fn test_iter() {
341        let mut map = RangeMap::<u32, i32>::default();
342
343        map.insert(10..34, -14);
344        map.insert(74..92, -12);
345
346        let mut iter = map.iter();
347
348        assert_eq!(iter.next().expect("missing elem"), (&(10..34), &-14));
349        assert_eq!(iter.next().expect("missing elem"), (&(74..92), &-12));
350
351        assert!(iter.next().is_none());
352
353        let mut iter = map.iter_starting_at(10);
354        assert_eq!(iter.next().expect("missing elem"), (&(10..34), &-14));
355        let mut iter = map.iter_starting_at(11);
356        assert_eq!(iter.next().expect("missing elem"), (&(74..92), &-12));
357        let mut iter = map.iter_starting_at(74);
358        assert_eq!(iter.next().expect("missing elem"), (&(74..92), &-12));
359        let mut iter = map.iter_starting_at(75);
360        assert_eq!(iter.next(), None);
361
362        assert_eq!(map.iter_ending_at(9).collect::<Vec<_>>(), vec![]);
363        assert_eq!(map.iter_ending_at(34).collect::<Vec<_>>(), vec![(&(10..34), &-14)]);
364        assert_eq!(map.iter_ending_at(74).collect::<Vec<_>>(), vec![(&(10..34), &-14)]);
365        assert_eq!(
366            map.iter_ending_at(75).collect::<Vec<_>>(),
367            vec![(&(10..34), &-14), (&(74..92), &-12)]
368        );
369        assert_eq!(
370            map.iter_ending_at(91).collect::<Vec<_>>(),
371            vec![(&(10..34), &-14), (&(74..92), &-12)]
372        );
373        assert_eq!(
374            map.iter_ending_at(92).collect::<Vec<_>>(),
375            vec![(&(10..34), &-14), (&(74..92), &-12)]
376        );
377    }
378
379    #[::fuchsia::test]
380    fn test_remove_overlapping_edge() {
381        let mut map = RangeMap::<u32, i32>::default();
382
383        map.insert(10..34, -14);
384
385        map.remove(2..11);
386        assert_eq!((&(11..34), &-14), map.get(11).unwrap());
387
388        map.remove(33..42);
389        assert_eq!((&(11..33), &-14), map.get(12).unwrap());
390    }
391
392    #[::fuchsia::test]
393    fn test_remove_middle_splits_range() {
394        let mut map = RangeMap::<u32, i32>::default();
395
396        map.insert(10..34, -14);
397        map.remove(15..18);
398
399        assert_eq!((&(10..15), &-14), map.get(12).unwrap());
400        assert_eq!((&(18..34), &-14), map.get(20).unwrap());
401    }
402
403    #[::fuchsia::test]
404    fn test_remove_upper_half_of_split_range_leaves_lower_range() {
405        let mut map = RangeMap::<u32, i32>::default();
406
407        map.insert(10..34, -14);
408        map.remove(15..18);
409        map.insert(2..7, -21);
410        map.remove(20..42);
411
412        assert_eq!((&(2..7), &-21), map.get(5).unwrap());
413        assert_eq!((&(10..15), &-14), map.get(12).unwrap());
414    }
415
416    #[::fuchsia::test]
417    fn test_range_map_overlapping_insert() {
418        let mut map = RangeMap::<u32, i32>::default();
419
420        map.insert(2..7, -21);
421        map.insert(5..9, -42);
422        map.insert(1..3, -43);
423        map.insert(6..8, -44);
424
425        assert_eq!((&(1..3), &-43), map.get(2).unwrap());
426        assert_eq!((&(3..5), &-21), map.get(4).unwrap());
427        assert_eq!((&(5..6), &-42), map.get(5).unwrap());
428        assert_eq!((&(6..8), &-44), map.get(7).unwrap());
429    }
430
431    #[::fuchsia::test]
432    fn test_intersect_single() {
433        let mut map = RangeMap::<u32, i32>::default();
434
435        map.insert(2..7, -10);
436
437        let mut iter = map.intersection(3..4);
438        assert_eq!(iter.next(), Some((&(2..7), &-10)));
439        assert_eq!(iter.next(), None);
440
441        let mut iter = map.intersection(2..3);
442        assert_eq!(iter.next(), Some((&(2..7), &-10)));
443        assert_eq!(iter.next(), None);
444
445        let mut iter = map.intersection(1..4);
446        assert_eq!(iter.next(), Some((&(2..7), &-10)));
447        assert_eq!(iter.next(), None);
448
449        let mut iter = map.intersection(1..2);
450        assert_eq!(iter.next(), None);
451
452        let mut iter = map.intersection(6..7);
453        assert_eq!(iter.next(), Some((&(2..7), &-10)));
454        assert_eq!(iter.next(), None);
455    }
456
457    #[::fuchsia::test]
458    fn test_intersect_multiple() {
459        let mut map = RangeMap::<u32, i32>::default();
460
461        map.insert(2..7, -10);
462        map.insert(7..9, -20);
463        map.insert(10..11, -30);
464
465        let mut iter = map.intersection(3..8);
466        assert_eq!(iter.next(), Some((&(2..7), &-10)));
467        assert_eq!(iter.next(), Some((&(7..9), &-20)));
468        assert_eq!(iter.next(), None);
469
470        let mut iter = map.intersection(3..11);
471        assert_eq!(iter.next(), Some((&(2..7), &-10)));
472        assert_eq!(iter.next(), Some((&(7..9), &-20)));
473        assert_eq!(iter.next(), Some((&(10..11), &-30)));
474        assert_eq!(iter.next(), None);
475    }
476
477    #[::fuchsia::test]
478    fn test_intersect_no_gaps() {
479        let mut map = RangeMap::<u32, i32>::default();
480
481        map.insert(0..1, -10);
482        map.insert(1..2, -20);
483        map.insert(2..3, -30);
484
485        let mut iter = map.intersection(0..3);
486        assert_eq!(iter.next(), Some((&(0..1), &-10)));
487        assert_eq!(iter.next(), Some((&(1..2), &-20)));
488        assert_eq!(iter.next(), Some((&(2..3), &-30)));
489        assert_eq!(iter.next(), None);
490    }
491
492    #[test]
493    fn test_merging() {
494        let mut map = RangeMap::<u32, i32>::default();
495
496        map.insert(1..2, -10);
497        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(1..2), &-10)]);
498        map.insert(3..4, -10);
499        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(1..2), &-10), (&(3..4), &-10)]);
500        map.insert(2..3, -10);
501        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(1..4), &-10)]);
502        map.insert(0..1, -10);
503        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(0..4), &-10)]);
504        map.insert(4..5, -10);
505        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(0..5), &-10)]);
506        map.insert(2..3, -20);
507        assert_eq!(
508            map.iter().collect::<Vec<_>>(),
509            vec![(&(0..2), &-10), (&(2..3), &-20), (&(3..5), &-10)]
510        );
511    }
512
513    #[test]
514    fn test_remove_multiple_ranges_exact_query() {
515        let first = 15..21;
516        let second = first.end..29;
517
518        let mut map = RangeMap::default();
519        map.insert(first.clone(), 1);
520        map.insert(second.clone(), 2);
521
522        assert_eq!(map.remove(first.start..second.end), &[1, 2]);
523    }
524}
range_map/classic.rs

range_map/
classic.rs
