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: 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		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		// 1. One range is set to a given set of values.
155		map.insert(0..10, "a");
156
157		// 2. A value is added to a sub-range: splits the range in three.
158		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	/// Regression test: adding a value to a sub-range and then removing it
171	/// again must merge the map back into a single range, instead of leaving
172	/// three ranges that all happen to carry the same value set.
173	#[test]
174	fn insert_then_remove_sub_range_merges_back() {
175		let mut map: RangeMultiMap<i32, BTreeSet<&str>> = RangeMultiMap::new();
176
177		// 1. One range is set to a given set of values (1 range).
178		map.insert(0..10, "a");
179		assert_eq!(map.iter().count(), 1);
180
181		// 2. A value is added to a sub-range (3 ranges).
182		map.insert(3..6, "b");
183		assert_eq!(map.iter().count(), 3);
184
185		// 3. The value is removed from that range: back to 1 range.
186		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}