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 pub fn insert(&mut self, location: ChunkOffset, value: V) {
19 self.btm.insert(location, value);
20 }
21 pub fn remove(&mut self, location: &ChunkOffset) -> Option<V> {
23 self.btm.remove(location)
24 }
25 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}