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