bitar/
chunk_location_map.rs

1use std::collections::BTreeMap;
2use std::ops::Bound;
3
4use crate::ChunkOffset;
5
6#[derive(Default)]
7pub(crate) struct ChunkLocationMap<V> {
8    btm: BTreeMap<ChunkOffset, V>,
9}
10
11impl<V> ChunkLocationMap<V> {
12    pub fn new() -> Self {
13        Self {
14            btm: BTreeMap::new(),
15        }
16    }
17    // Insert at location. Does not verify if the given location overlaps.
18    pub fn insert(&mut self, location: ChunkOffset, value: V) {
19        self.btm.insert(location, value);
20    }
21    // Remove at exact location.
22    pub fn remove(&mut self, location: &ChunkOffset) -> Option<V> {
23        self.btm.remove(location)
24    }
25    // Iterate all locations which the given location overlaps.
26    pub fn iter_overlapping(
27        &self,
28        location: ChunkOffset,
29    ) -> impl Iterator<Item = (&ChunkOffset, &V)> {
30        let location_offset = location.offset;
31        self.btm
32            .range((
33                Bound::Unbounded,
34                Bound::Excluded(ChunkOffset::new(location.end(), 0)),
35            ))
36            .rev()
37            .take_while(move |(loc, _v)| location_offset < loc.end())
38    }
39}
40
41#[cfg(test)]
42mod tests {
43    use super::*;
44
45    #[test]
46    fn insert_and_iter_in_order() {
47        let mut clm = ChunkLocationMap::new();
48        clm.insert(ChunkOffset::new(0, 10), 1);
49        clm.insert(ChunkOffset::new(15, 7), 3);
50        clm.insert(ChunkOffset::new(10, 5), 2);
51        let mut iter = clm.btm.iter();
52        assert_eq!(iter.next(), Some((&ChunkOffset::new(0, 10), &1)));
53        assert_eq!(iter.next(), Some((&ChunkOffset::new(10, 5), &2)));
54        assert_eq!(iter.next(), Some((&ChunkOffset::new(15, 7), &3)));
55        assert_eq!(iter.next(), None);
56    }
57    #[test]
58    fn remove_one() {
59        let mut clm = ChunkLocationMap::new();
60        clm.insert(ChunkOffset::new(0, 10), 1);
61        clm.insert(ChunkOffset::new(15, 7), 3);
62        clm.insert(ChunkOffset::new(10, 5), 2);
63        assert_eq!(clm.remove(&ChunkOffset::new(15, 7)), Some(3));
64        assert_eq!(clm.btm.len(), 2);
65    }
66    #[test]
67    fn remove_first() {
68        let mut clm = ChunkLocationMap::new();
69        clm.insert(ChunkOffset::new(0, 10), 1);
70        clm.insert(ChunkOffset::new(15, 7), 2);
71        assert_eq!(clm.remove(&ChunkOffset::new(0, 10)), Some(1));
72        assert_eq!(clm.btm.len(), 1);
73    }
74    #[test]
75    fn remove_last() {
76        let mut clm = ChunkLocationMap::new();
77        clm.insert(ChunkOffset::new(0, 10), 1);
78        clm.insert(ChunkOffset::new(15, 7), 2);
79        assert_eq!(clm.remove(&ChunkOffset::new(15, 7)), Some(2));
80        assert_eq!(clm.btm.len(), 1);
81    }
82    #[test]
83    fn no_remove() {
84        let mut clm = ChunkLocationMap::new();
85        clm.insert(ChunkOffset::new(0, 10), 1);
86        clm.insert(ChunkOffset::new(15, 7), 2);
87        assert_eq!(clm.remove(&ChunkOffset::new(0, 1)), None);
88        assert_eq!(clm.remove(&ChunkOffset::new(0, 9)), None);
89        assert_eq!(clm.remove(&ChunkOffset::new(1, 10)), None);
90        assert_eq!(clm.btm.len(), 2);
91    }
92    #[test]
93    fn some_overlap() {
94        let mut clm = ChunkLocationMap::new();
95        clm.insert(ChunkOffset::new(0, 10), 1);
96        clm.insert(ChunkOffset::new(10, 5), 2);
97        clm.insert(ChunkOffset::new(15, 7), 3);
98        clm.insert(ChunkOffset::new(30, 10), 4);
99        clm.insert(ChunkOffset::new(40, 10), 5);
100
101        let mut iter = clm.iter_overlapping(ChunkOffset::new(10, 30));
102        assert_eq!(iter.next(), Some((&ChunkOffset::new(30, 10), &4)));
103        assert_eq!(iter.next(), Some((&ChunkOffset::new(15, 7), &3)));
104        assert_eq!(iter.next(), Some((&ChunkOffset::new(10, 5), &2)));
105        assert_eq!(iter.next(), None);
106    }
107    #[test]
108    fn exact_overlap() {
109        let mut clm = ChunkLocationMap::new();
110        clm.insert(ChunkOffset::new(0, 5), 1);
111        clm.insert(ChunkOffset::new(5, 5), 2);
112        clm.insert(ChunkOffset::new(10, 5), 3);
113        clm.insert(ChunkOffset::new(15, 5), 4);
114
115        let mut iter = clm.iter_overlapping(ChunkOffset::new(5, 5));
116        assert_eq!(iter.next(), Some((&ChunkOffset::new(5, 5), &2)));
117        assert_eq!(iter.next(), None);
118    }
119    #[test]
120    fn exact_overlap_plus_one() {
121        let mut clm = ChunkLocationMap::new();
122        clm.insert(ChunkOffset::new(0, 5), 1);
123        clm.insert(ChunkOffset::new(5, 5), 2);
124        clm.insert(ChunkOffset::new(10, 5), 3);
125        clm.insert(ChunkOffset::new(15, 5), 4);
126
127        let mut iter = clm.iter_overlapping(ChunkOffset::new(5, 6));
128        assert_eq!(iter.next(), Some((&ChunkOffset::new(10, 5), &3)));
129        assert_eq!(iter.next(), Some((&ChunkOffset::new(5, 5), &2)));
130        assert_eq!(iter.next(), None);
131    }
132    #[test]
133    fn exact_overlap_minus_one() {
134        let mut clm = ChunkLocationMap::new();
135        clm.insert(ChunkOffset::new(0, 5), 1);
136        clm.insert(ChunkOffset::new(5, 5), 2);
137        clm.insert(ChunkOffset::new(10, 5), 3);
138        clm.insert(ChunkOffset::new(15, 5), 4);
139
140        let mut iter = clm.iter_overlapping(ChunkOffset::new(4, 6));
141        assert_eq!(iter.next(), Some((&ChunkOffset::new(5, 5), &2)));
142        assert_eq!(iter.next(), Some((&ChunkOffset::new(0, 5), &1)));
143        assert_eq!(iter.next(), None);
144    }
145    #[test]
146    fn above_no_overlap() {
147        let mut clm = ChunkLocationMap::new();
148        clm.insert(ChunkOffset::new(0, 5), 1);
149        clm.insert(ChunkOffset::new(5, 5), 2);
150        clm.insert(ChunkOffset::new(10, 5), 3);
151        clm.insert(ChunkOffset::new(15, 5), 4);
152
153        let mut iter = clm.iter_overlapping(ChunkOffset::new(20, 10));
154        assert_eq!(iter.next(), None);
155    }
156    #[test]
157    fn between_no_overlap() {
158        let mut clm = ChunkLocationMap::new();
159        clm.insert(ChunkOffset::new(15, 7), 1);
160        clm.insert(ChunkOffset::new(30, 10), 2);
161
162        let mut iter = clm.iter_overlapping(ChunkOffset::new(23, 6));
163        assert_eq!(iter.next(), None);
164    }
165}