1use std::borrow::Borrow;
2use std::cmp::Ordering;
3use std::fmt::Debug;
4use std::mem;
5use std::ops::{Bound, Index, IndexMut, RangeBounds};
6
7use rustc_macros::{Decodable_NoContext, Encodable_NoContext};
8
9use crate::stable_hasher::{HashStable, StableHasher, StableOrd};
10
11mod index_map;
12
13pub use index_map::SortedIndexMultiMap;
14
15#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Encodable_NoContext, Decodable_NoContext)]
24pub struct SortedMap<K, V> {
25 data: Vec<(K, V)>,
26}
27
28impl<K, V> Default for SortedMap<K, V> {
29 #[inline]
30 fn default() -> SortedMap<K, V> {
31 SortedMap { data: Vec::new() }
32 }
33}
34
35impl<K, V> SortedMap<K, V> {
36 #[inline]
37 pub const fn new() -> SortedMap<K, V> {
38 SortedMap { data: Vec::new() }
39 }
40}
41
42impl<K: Ord, V> SortedMap<K, V> {
43 #[inline]
49 pub fn from_presorted_elements(elements: Vec<(K, V)>) -> SortedMap<K, V> {
50 debug_assert!(elements.array_windows().all(|[fst, snd]| fst.0 < snd.0));
51
52 SortedMap { data: elements }
53 }
54
55 #[inline]
56 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
57 match self.lookup_index_for(&key) {
58 Ok(index) => {
59 let slot = unsafe { self.data.get_unchecked_mut(index) };
60 Some(mem::replace(&mut slot.1, value))
61 }
62 Err(index) => {
63 self.data.insert(index, (key, value));
64 None
65 }
66 }
67 }
68
69 #[inline]
70 pub fn remove(&mut self, key: &K) -> Option<V> {
71 match self.lookup_index_for(key) {
72 Ok(index) => Some(self.data.remove(index).1),
73 Err(_) => None,
74 }
75 }
76
77 #[inline]
78 pub fn get<Q>(&self, key: &Q) -> Option<&V>
79 where
80 K: Borrow<Q>,
81 Q: Ord + ?Sized,
82 {
83 match self.lookup_index_for(key) {
84 Ok(index) => unsafe { Some(&self.data.get_unchecked(index).1) },
85 Err(_) => None,
86 }
87 }
88
89 #[inline]
90 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
91 where
92 K: Borrow<Q>,
93 Q: Ord + ?Sized,
94 {
95 match self.lookup_index_for(key) {
96 Ok(index) => unsafe { Some(&mut self.data.get_unchecked_mut(index).1) },
97 Err(_) => None,
98 }
99 }
100
101 #[inline]
103 pub fn get_mut_or_insert_default(&mut self, key: K) -> &mut V
104 where
105 K: Eq,
106 V: Default,
107 {
108 let index = match self.lookup_index_for(&key) {
109 Ok(index) => index,
110 Err(index) => {
111 self.data.insert(index, (key, V::default()));
112 index
113 }
114 };
115 unsafe { &mut self.data.get_unchecked_mut(index).1 }
116 }
117
118 #[inline]
119 pub fn clear(&mut self) {
120 self.data.clear();
121 }
122
123 #[inline]
125 pub fn iter(&self) -> std::slice::Iter<'_, (K, V)> {
126 self.data.iter()
127 }
128
129 #[inline]
131 pub fn keys(&self) -> impl ExactSizeIterator<Item = &K> + DoubleEndedIterator {
132 self.data.iter().map(|(k, _)| k)
133 }
134
135 #[inline]
137 pub fn values(&self) -> impl ExactSizeIterator<Item = &V> + DoubleEndedIterator {
138 self.data.iter().map(|(_, v)| v)
139 }
140
141 #[inline]
142 pub fn len(&self) -> usize {
143 self.data.len()
144 }
145
146 #[inline]
147 pub fn is_empty(&self) -> bool {
148 self.len() == 0
149 }
150
151 #[inline]
152 pub fn range<R>(&self, range: R) -> &[(K, V)]
153 where
154 R: RangeBounds<K>,
155 {
156 let (start, end) = self.range_slice_indices(range);
157 &self.data[start..end]
158 }
159
160 #[inline]
162 pub fn range_is_empty<R>(&self, range: R) -> bool
163 where
164 R: RangeBounds<K>,
165 {
166 self.data
171 .binary_search_by(|(x, _)| {
172 match range.start_bound() {
174 Bound::Included(start) if x < start => return Ordering::Less,
175 Bound::Excluded(start) if x <= start => return Ordering::Less,
176 _ => {}
177 };
178
179 match range.end_bound() {
181 Bound::Included(end) if x > end => return Ordering::Greater,
182 Bound::Excluded(end) if x >= end => return Ordering::Greater,
183 _ => {}
184 };
185
186 Ordering::Equal
188 })
189 .is_err()
190 }
191
192 #[inline]
193 pub fn remove_range<R>(&mut self, range: R)
194 where
195 R: RangeBounds<K>,
196 {
197 let (start, end) = self.range_slice_indices(range);
198 self.data.splice(start..end, std::iter::empty());
199 }
200
201 #[inline]
204 pub fn offset_keys<F>(&mut self, f: F)
205 where
206 F: Fn(&mut K),
207 {
208 self.data.iter_mut().map(|(k, _)| k).for_each(f);
209 }
210
211 #[inline]
218 pub fn insert_presorted(&mut self, elements: Vec<(K, V)>) {
219 if elements.is_empty() {
220 return;
221 }
222
223 debug_assert!(elements.array_windows().all(|[fst, snd]| fst.0 < snd.0));
224
225 let start_index = self.lookup_index_for(&elements[0].0);
226
227 let elements = match start_index {
228 Ok(index) => {
229 let mut elements = elements.into_iter();
230 self.data[index] = elements.next().unwrap();
231 elements
232 }
233 Err(index) => {
234 if index == self.data.len() || elements.last().unwrap().0 < self.data[index].0 {
235 self.data.splice(index..index, elements);
238 return;
239 }
240
241 let mut elements = elements.into_iter();
242 self.data.insert(index, elements.next().unwrap());
243 elements
244 }
245 };
246
247 for (k, v) in elements {
249 self.insert(k, v);
250 }
251 }
252
253 #[inline(always)]
255 fn lookup_index_for<Q>(&self, key: &Q) -> Result<usize, usize>
256 where
257 K: Borrow<Q>,
258 Q: Ord + ?Sized,
259 {
260 self.data.binary_search_by(|(x, _)| x.borrow().cmp(key))
261 }
262
263 #[inline]
264 fn range_slice_indices<R>(&self, range: R) -> (usize, usize)
265 where
266 R: RangeBounds<K>,
267 {
268 let start = match range.start_bound() {
269 Bound::Included(k) => match self.lookup_index_for(k) {
270 Ok(index) | Err(index) => index,
271 },
272 Bound::Excluded(k) => match self.lookup_index_for(k) {
273 Ok(index) => index + 1,
274 Err(index) => index,
275 },
276 Bound::Unbounded => 0,
277 };
278
279 let end = match range.end_bound() {
280 Bound::Included(k) => match self.lookup_index_for(k) {
281 Ok(index) => index + 1,
282 Err(index) => index,
283 },
284 Bound::Excluded(k) => match self.lookup_index_for(k) {
285 Ok(index) | Err(index) => index,
286 },
287 Bound::Unbounded => self.data.len(),
288 };
289
290 (start, end)
291 }
292
293 #[inline]
294 pub fn contains_key<Q>(&self, key: &Q) -> bool
295 where
296 K: Borrow<Q>,
297 Q: Ord + ?Sized,
298 {
299 self.get(key).is_some()
300 }
301}
302
303impl<K: Ord, V> IntoIterator for SortedMap<K, V> {
304 type Item = (K, V);
305 type IntoIter = std::vec::IntoIter<(K, V)>;
306
307 fn into_iter(self) -> Self::IntoIter {
308 self.data.into_iter()
309 }
310}
311
312impl<'a, K, Q, V> Index<&'a Q> for SortedMap<K, V>
313where
314 K: Ord + Borrow<Q>,
315 Q: Ord + ?Sized,
316{
317 type Output = V;
318
319 fn index(&self, key: &Q) -> &Self::Output {
320 self.get(key).expect("no entry found for key")
321 }
322}
323
324impl<'a, K, Q, V> IndexMut<&'a Q> for SortedMap<K, V>
325where
326 K: Ord + Borrow<Q>,
327 Q: Ord + ?Sized,
328{
329 fn index_mut(&mut self, key: &Q) -> &mut Self::Output {
330 self.get_mut(key).expect("no entry found for key")
331 }
332}
333
334impl<K: Ord, V> FromIterator<(K, V)> for SortedMap<K, V> {
335 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
336 let mut data: Vec<(K, V)> = iter.into_iter().collect();
337
338 data.sort_unstable_by(|(k1, _), (k2, _)| k1.cmp(k2));
339 data.dedup_by(|(k1, _), (k2, _)| k1 == k2);
340
341 SortedMap { data }
342 }
343}
344
345impl<K: HashStable<CTX> + StableOrd, V: HashStable<CTX>, CTX> HashStable<CTX> for SortedMap<K, V> {
346 #[inline]
347 fn hash_stable(&self, ctx: &mut CTX, hasher: &mut StableHasher) {
348 self.data.hash_stable(ctx, hasher);
349 }
350}
351
352impl<K: Debug, V: Debug> Debug for SortedMap<K, V> {
353 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
354 f.debug_map().entries(self.data.iter().map(|(a, b)| (a, b))).finish()
355 }
356}
357
358#[cfg(test)]
359mod tests;