1use crate::{
2 AnyRange, AsRange,
3 generic::{
4 RangeMap,
5 map::{IntoIter, Iter},
6 },
7};
8use cc_traits::SetMut;
9use range_traits::{Measure, PartialEnum};
10use raw_btree::{Item, Storage};
11
12pub struct RangeMultiMap<K, S, C: Storage<Item<AnyRange<K>, S>>> {
19 map: RangeMap<K, S, C>,
20}
21
22impl<K: Clone, S: Clone, C: Storage<Item<AnyRange<K>, S>>> Clone for RangeMultiMap<K, S, C> {
23 fn clone(&self) -> Self {
24 RangeMultiMap {
25 map: self.map.clone(),
26 }
27 }
28}
29
30impl<K, S, C: Storage<Item<AnyRange<K>, S>>> RangeMultiMap<K, S, C> {
31 pub fn new() -> RangeMultiMap<K, S, C> {
32 RangeMultiMap {
33 map: RangeMap::new(),
34 }
35 }
36}
37
38impl<K, S, C: Storage<Item<AnyRange<K>, S>>> Default for RangeMultiMap<K, S, C> {
39 fn default() -> Self {
40 Self::new()
41 }
42}
43
44impl<K, S, C: Storage<Item<AnyRange<K>, S>>> RangeMultiMap<K, S, C> {
45 pub fn iter(&self) -> Iter<'_, K, S, C> {
46 self.map.iter()
47 }
48}
49
50impl<'a, K: Clone + PartialOrd + Measure, S, C: Storage<Item<AnyRange<K>, S>>> IntoIterator
51 for &'a RangeMultiMap<K, S, C>
52{
53 type Item = (&'a AnyRange<K>, &'a S);
54 type IntoIter = Iter<'a, K, S, C>;
55
56 fn into_iter(self) -> Self::IntoIter {
57 self.iter()
58 }
59}
60
61impl<K, S, C: Storage<Item<AnyRange<K>, S>>> RangeMultiMap<K, S, C> {
62 pub fn insert<R: AsRange<Item = K>, V>(&mut self, key: R, value: V)
70 where
71 K: Clone + PartialEnum + Measure,
72 V: Clone,
73 S: SetMut<V> + PartialEq + Clone + Default,
74 {
75 self.map.update(key, |set_opt| {
76 let mut result = match set_opt {
77 Some(set) => set.clone(),
78 None => S::default(),
79 };
80
81 result.insert(value.clone());
82 Some(result)
83 })
84 }
85
86 pub fn remove<R: AsRange<Item = K>, V>(&mut self, key: R, value: &V)
94 where
95 K: Clone + PartialEnum + Measure,
96 S: SetMut<V> + PartialEq + Clone,
97 {
98 self.map.update(key, |set_opt| match set_opt {
99 Some(set) => {
100 let mut result = set.clone();
101 result.remove(value);
102 if result.is_empty() {
103 None
104 } else {
105 Some(result)
106 }
107 }
108 None => None,
109 })
110 }
111}
112
113impl<K: Clone + PartialOrd + Measure, S, C: Storage<Item<AnyRange<K>, S>>> IntoIterator
114 for RangeMultiMap<K, S, C>
115{
116 type Item = (AnyRange<K>, S);
117 type IntoIter = IntoIter<K, S, C>;
118
119 fn into_iter(self) -> Self::IntoIter {
120 self.map.into_iter()
121 }
122}
123
124#[cfg(test)]
125mod tests {
126 use crate::{AnyRange, RangeMultiMap};
127 use std::collections::BTreeSet;
128
129 fn set<const N: usize>(values: [&'static str; N]) -> BTreeSet<&'static str> {
130 values.into_iter().collect()
131 }
132
133 fn entries(
134 map: &RangeMultiMap<i32, BTreeSet<&'static str>>,
135 ) -> Vec<(AnyRange<i32>, BTreeSet<&'static str>)> {
136 map.iter()
137 .map(|(range, values)| (*range, values.clone()))
138 .collect()
139 }
140
141 #[test]
142 fn insert_single_range() {
143 let mut map: RangeMultiMap<i32, BTreeSet<&str>> = RangeMultiMap::new();
144
145 map.insert(0..10, "a");
146
147 assert_eq!(entries(&map), vec![(AnyRange::from(0..10), set(["a"]))]);
148 }
149
150 #[test]
151 fn insert_creates_sub_range() {
152 let mut map: RangeMultiMap<i32, BTreeSet<&str>> = RangeMultiMap::new();
153
154 map.insert(0..10, "a");
156
157 map.insert(3..6, "b");
159
160 assert_eq!(
161 entries(&map),
162 vec![
163 (AnyRange::from(0..3), set(["a"])),
164 (AnyRange::from(3..6), set(["a", "b"])),
165 (AnyRange::from(6..10), set(["a"])),
166 ]
167 );
168 }
169
170 #[test]
174 fn insert_then_remove_sub_range_merges_back() {
175 let mut map: RangeMultiMap<i32, BTreeSet<&str>> = RangeMultiMap::new();
176
177 map.insert(0..10, "a");
179 assert_eq!(map.iter().count(), 1);
180
181 map.insert(3..6, "b");
183 assert_eq!(map.iter().count(), 3);
184
185 map.remove(3..6, &"b");
187
188 assert_eq!(
189 entries(&map),
190 vec![(AnyRange::from(0..10), set(["a"]))],
191 "the three ranges must be merged back into one, not left as three \
192 identical ranges"
193 );
194 }
195
196 #[test]
197 fn remove_absent_value_is_noop() {
198 let mut map: RangeMultiMap<i32, BTreeSet<&str>> = RangeMultiMap::new();
199
200 map.insert(0..10, "a");
201 map.remove(0..10, &"z");
202
203 assert_eq!(entries(&map), vec![(AnyRange::from(0..10), set(["a"]))]);
204 }
205
206 #[test]
207 fn remove_last_value_drops_range() {
208 let mut map: RangeMultiMap<i32, BTreeSet<&str>> = RangeMultiMap::new();
209
210 map.insert(0..10, "a");
211 map.remove(0..10, &"a");
212
213 assert_eq!(entries(&map), vec![]);
214 assert_eq!(map.iter().count(), 0);
215 }
216}