Skip to main content

btree_range_map/generic/
multimap.rs

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
12/// Multi map.
13///
14/// In a multi map, each key is associated to a set of values.
15/// The type parameter `S` is the set type. It can be replaced by anything
16/// implementing the [`cc_traits::SetMut`] trait, such as the standard
17/// [`BTreeSet`](std::collections::BTreeSet) and [`HashSet`](std::collections::HashSet).
18pub 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	/// # Complexity
63	///
64	/// `O((k + 1) log n + |S|)`, where `n` is the number of ranges in the map,
65	/// `k` is the number of existing ranges that overlap, or are connected
66	/// to, `key` (see [`RangeMap::update`]), and `|S|` is the size of the
67	/// value set attached to the affected range, which is cloned and mutated
68	/// by this operation.
69	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	/// # Complexity
87	///
88	/// `O((k + 1) log n + |S|)`, where `n` is the number of ranges in the map,
89	/// `k` is the number of existing ranges that overlap, or are connected
90	/// to, `key` (see [`RangeMap::update`]), and `|S|` is the size of the
91	/// value set attached to the affected range, which is cloned and mutated
92	/// by this operation.
93	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		// 1. One range is set to a given set of values.
156		map.insert(0..10, "a");
157
158		// 2. A value is added to a sub-range: splits the range in three.
159		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	/// Regression test: adding a value to a sub-range and then removing it
172	/// again must merge the map back into a single range, instead of leaving
173	/// three ranges that all happen to carry the same value set.
174	#[test]
175	fn insert_then_remove_sub_range_merges_back() {
176		let mut map: RangeMultiMap<i32, BTreeSet<&str>> = RangeMultiMap::new();
177
178		// 1. One range is set to a given set of values (1 range).
179		map.insert(0..10, "a");
180		assert_eq!(map.iter().count(), 1);
181
182		// 2. A value is added to a sub-range (3 ranges).
183		map.insert(3..6, "b");
184		assert_eq!(map.iter().count(), 3);
185
186		// 3. The value is removed from that range: back to 1 range.
187		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}