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: PartialEq + 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 V: PartialEq + Clone,
97 S: SetMut<V> + PartialEq + Clone + Default,
98 {
99 self.map.update(key, |set_opt| match set_opt {
100 Some(set) => {
101 let mut result = set.clone();
102 result.remove(value);
103 if result.is_empty() {
104 None
105 } else {
106 Some(result)
107 }
108 }
109 None => None,
110 })
111 }
112}
113
114impl<K: Clone + PartialOrd + Measure, S, C: Storage<Item<AnyRange<K>, S>>> IntoIterator
115 for RangeMultiMap<K, S, C>
116{
117 type Item = (AnyRange<K>, S);
118 type IntoIter = IntoIter<K, S, C>;
119
120 fn into_iter(self) -> Self::IntoIter {
121 self.map.into_iter()
122 }
123}
124
125#[cfg(test)]
126mod tests {
127 use crate::{AnyRange, RangeMultiMap};
128 use std::collections::BTreeSet;
129
130 fn set<const N: usize>(values: [&'static str; N]) -> BTreeSet<&'static str> {
131 values.into_iter().collect()
132 }
133
134 fn entries(
135 map: &RangeMultiMap<i32, BTreeSet<&'static str>>,
136 ) -> Vec<(AnyRange<i32>, BTreeSet<&'static str>)> {
137 map.iter()
138 .map(|(range, values)| (*range, values.clone()))
139 .collect()
140 }
141
142 #[test]
143 fn insert_single_range() {
144 let mut map: RangeMultiMap<i32, BTreeSet<&str>> = RangeMultiMap::new();
145
146 map.insert(0..10, "a");
147
148 assert_eq!(entries(&map), vec![(AnyRange::from(0..10), set(["a"]))]);
149 }
150
151 #[test]
152 fn insert_creates_sub_range() {
153 let mut map: RangeMultiMap<i32, BTreeSet<&str>> = RangeMultiMap::new();
154
155 map.insert(0..10, "a");
157
158 map.insert(3..6, "b");
160
161 assert_eq!(
162 entries(&map),
163 vec![
164 (AnyRange::from(0..3), set(["a"])),
165 (AnyRange::from(3..6), set(["a", "b"])),
166 (AnyRange::from(6..10), set(["a"])),
167 ]
168 );
169 }
170
171 #[test]
175 fn insert_then_remove_sub_range_merges_back() {
176 let mut map: RangeMultiMap<i32, BTreeSet<&str>> = RangeMultiMap::new();
177
178 map.insert(0..10, "a");
180 assert_eq!(map.iter().count(), 1);
181
182 map.insert(3..6, "b");
184 assert_eq!(map.iter().count(), 3);
185
186 map.remove(3..6, &"b");
188
189 assert_eq!(
190 entries(&map),
191 vec![(AnyRange::from(0..10), set(["a"]))],
192 "the three ranges must be merged back into one, not left as three \
193 identical ranges"
194 );
195 }
196
197 #[test]
198 fn remove_absent_value_is_noop() {
199 let mut map: RangeMultiMap<i32, BTreeSet<&str>> = RangeMultiMap::new();
200
201 map.insert(0..10, "a");
202 map.remove(0..10, &"z");
203
204 assert_eq!(entries(&map), vec![(AnyRange::from(0..10), set(["a"]))]);
205 }
206
207 #[test]
208 fn remove_last_value_drops_range() {
209 let mut map: RangeMultiMap<i32, BTreeSet<&str>> = RangeMultiMap::new();
210
211 map.insert(0..10, "a");
212 map.remove(0..10, &"a");
213
214 assert_eq!(entries(&map), vec![]);
215 assert_eq!(map.iter().count(), 0);
216 }
217}