]>
Commit | Line | Data |
---|---|---|
1 | //! A variant of `SortedMap` that preserves insertion order. | |
2 | ||
3 | use std::hash::{Hash, Hasher}; | |
4 | ||
5 | use crate::stable_hasher::{HashStable, StableHasher}; | |
6 | use rustc_index::{Idx, IndexVec}; | |
7 | ||
8 | /// An indexed multi-map that preserves insertion order while permitting both *O*(log *n*) lookup of | |
9 | /// an item by key and *O*(1) lookup by index. | |
10 | /// | |
11 | /// This data structure is a hybrid of an [`IndexVec`] and a [`SortedMap`]. Like `IndexVec`, | |
12 | /// `SortedIndexMultiMap` assigns a typed index to each item while preserving insertion order. | |
13 | /// Like `SortedMap`, `SortedIndexMultiMap` has efficient lookup of items by key. However, this | |
14 | /// is accomplished by sorting an array of item indices instead of the items themselves. | |
15 | /// | |
16 | /// Unlike `SortedMap`, this data structure can hold multiple equivalent items at once, so the | |
17 | /// `get_by_key` method and its variants return an iterator instead of an `Option`. Equivalent | |
18 | /// items will be yielded in insertion order. | |
19 | /// | |
20 | /// Unlike a general-purpose map like `BTreeSet` or `HashSet`, `SortedMap` and | |
21 | /// `SortedIndexMultiMap` require *O*(*n*) time to insert a single item. This is because we may need | |
22 | /// to insert into the middle of the sorted array. Users should avoid mutating this data structure | |
23 | /// in-place. | |
24 | /// | |
25 | /// [`SortedMap`]: super::SortedMap | |
26 | #[derive(Clone, Debug)] | |
27 | pub struct SortedIndexMultiMap<I: Idx, K, V> { | |
28 | /// The elements of the map in insertion order. | |
29 | items: IndexVec<I, (K, V)>, | |
30 | ||
31 | /// Indices of the items in the set, sorted by the item's key. | |
32 | idx_sorted_by_item_key: Vec<I>, | |
33 | } | |
34 | ||
35 | impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> { | |
36 | #[inline] | |
37 | pub fn new() -> Self { | |
38 | SortedIndexMultiMap { items: IndexVec::new(), idx_sorted_by_item_key: Vec::new() } | |
39 | } | |
40 | ||
41 | #[inline] | |
42 | pub fn len(&self) -> usize { | |
43 | self.items.len() | |
44 | } | |
45 | ||
46 | #[inline] | |
47 | pub fn is_empty(&self) -> bool { | |
48 | self.items.is_empty() | |
49 | } | |
50 | ||
51 | /// Returns an iterator over the items in the map in insertion order. | |
52 | #[inline] | |
53 | pub fn into_iter(self) -> impl DoubleEndedIterator<Item = (K, V)> { | |
54 | self.items.into_iter() | |
55 | } | |
56 | ||
57 | /// Returns an iterator over the items in the map in insertion order along with their indices. | |
58 | #[inline] | |
59 | pub fn into_iter_enumerated(self) -> impl DoubleEndedIterator<Item = (I, (K, V))> { | |
60 | self.items.into_iter_enumerated() | |
61 | } | |
62 | ||
63 | /// Returns an iterator over the items in the map in insertion order. | |
64 | #[inline] | |
65 | pub fn iter(&self) -> impl '_ + DoubleEndedIterator<Item = (&K, &V)> { | |
66 | self.items.iter().map(|(k, v)| (k, v)) | |
67 | } | |
68 | ||
69 | /// Returns an iterator over the items in the map in insertion order along with their indices. | |
70 | #[inline] | |
71 | pub fn iter_enumerated(&self) -> impl '_ + DoubleEndedIterator<Item = (I, (&K, &V))> { | |
72 | self.items.iter_enumerated().map(|(i, (k, v))| (i, (k, v))) | |
73 | } | |
74 | ||
75 | /// Returns the item in the map with the given index. | |
76 | #[inline] | |
77 | pub fn get(&self, idx: I) -> Option<&(K, V)> { | |
78 | self.items.get(idx) | |
79 | } | |
80 | ||
81 | /// Returns an iterator over the items in the map that are equal to `key`. | |
82 | /// | |
83 | /// If there are multiple items that are equivalent to `key`, they will be yielded in | |
84 | /// insertion order. | |
85 | #[inline] | |
86 | pub fn get_by_key(&self, key: K) -> impl Iterator<Item = &V> + '_ { | |
87 | self.get_by_key_enumerated(key).map(|(_, v)| v) | |
88 | } | |
89 | ||
90 | /// Returns an iterator over the items in the map that are equal to `key` along with their | |
91 | /// indices. | |
92 | /// | |
93 | /// If there are multiple items that are equivalent to `key`, they will be yielded in | |
94 | /// insertion order. | |
95 | #[inline] | |
96 | pub fn get_by_key_enumerated(&self, key: K) -> impl Iterator<Item = (I, &V)> + '_ { | |
97 | let lower_bound = self.idx_sorted_by_item_key.partition_point(|&i| self.items[i].0 < key); | |
98 | self.idx_sorted_by_item_key[lower_bound..].iter().map_while(move |&i| { | |
99 | let (k, v) = &self.items[i]; | |
100 | (k == &key).then_some((i, v)) | |
101 | }) | |
102 | } | |
103 | ||
104 | #[inline] | |
105 | pub fn contains_key(&self, key: K) -> bool { | |
106 | self.get_by_key(key).next().is_some() | |
107 | } | |
108 | } | |
109 | ||
110 | impl<I: Idx, K: Eq, V: Eq> Eq for SortedIndexMultiMap<I, K, V> {} | |
111 | impl<I: Idx, K: PartialEq, V: PartialEq> PartialEq for SortedIndexMultiMap<I, K, V> { | |
112 | fn eq(&self, other: &Self) -> bool { | |
113 | // No need to compare the sorted index. If the items are the same, the index will be too. | |
114 | self.items == other.items | |
115 | } | |
116 | } | |
117 | ||
118 | impl<I: Idx, K, V> Hash for SortedIndexMultiMap<I, K, V> | |
119 | where | |
120 | K: Hash, | |
121 | V: Hash, | |
122 | { | |
123 | fn hash<H: Hasher>(&self, hasher: &mut H) { | |
124 | self.items.hash(hasher) | |
125 | } | |
126 | } | |
127 | ||
128 | impl<I: Idx, K, V, C> HashStable<C> for SortedIndexMultiMap<I, K, V> | |
129 | where | |
130 | K: HashStable<C>, | |
131 | V: HashStable<C>, | |
132 | { | |
133 | fn hash_stable(&self, ctx: &mut C, hasher: &mut StableHasher) { | |
134 | let SortedIndexMultiMap { | |
135 | items, | |
136 | // We can ignore this field because it is not observable from the outside. | |
137 | idx_sorted_by_item_key: _, | |
138 | } = self; | |
139 | ||
140 | items.hash_stable(ctx, hasher) | |
141 | } | |
142 | } | |
143 | ||
144 | impl<I: Idx, K: Ord, V> FromIterator<(K, V)> for SortedIndexMultiMap<I, K, V> { | |
145 | fn from_iter<J>(iter: J) -> Self | |
146 | where | |
147 | J: IntoIterator<Item = (K, V)>, | |
148 | { | |
149 | let items = IndexVec::from_iter(iter); | |
150 | let mut idx_sorted_by_item_key: Vec<_> = items.indices().collect(); | |
151 | ||
152 | // `sort_by_key` is stable, so insertion order is preserved for duplicate items. | |
153 | idx_sorted_by_item_key.sort_by_key(|&idx| &items[idx].0); | |
154 | ||
155 | SortedIndexMultiMap { items, idx_sorted_by_item_key } | |
156 | } | |
157 | } | |
158 | ||
159 | impl<I: Idx, K, V> std::ops::Index<I> for SortedIndexMultiMap<I, K, V> { | |
160 | type Output = V; | |
161 | ||
162 | fn index(&self, idx: I) -> &Self::Output { | |
163 | &self.items[idx].1 | |
164 | } | |
165 | } |