1 //! A variant of `SortedMap` that preserves insertion order. 2 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 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. 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 22 /// `SortedIndexMultiMap` require *O*(*n*) time to insert a single item. This is because we may need 23 /// to insert into the middle of the sorted array. Users should avoid mutating this data structure 24 /// in-place. 25 /// 26 /// [`SortedMap`]: super::SortedMap 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> { 37 #[inline] new() -> Self38 pub fn new() -> Self { 39 SortedIndexMultiMap { items: IndexVec::new(), idx_sorted_by_item_key: Vec::new() } 40 } 41 42 #[inline] len(&self) -> usize43 pub fn len(&self) -> usize { 44 self.items.len() 45 } 46 47 #[inline] is_empty(&self) -> bool48 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. 53 #[inline] into_iter(self) -> impl DoubleEndedIterator<Item = (K, V)>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. 59 #[inline] into_iter_enumerated(self) -> impl DoubleEndedIterator<Item = (I, (K, V))>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. 65 #[inline] iter(&self) -> impl '_ + DoubleEndedIterator<Item = (&K, &V)>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. 71 #[inline] iter_enumerated(&self) -> impl '_ + DoubleEndedIterator<Item = (I, (&K, &V))>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. 77 #[inline] get(&self, idx: I) -> Option<&(K, V)>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. 86 #[inline] get_by_key(&'a self, key: K) -> impl 'a + Iterator<Item = &'a V>87 pub fn get_by_key(&'a self, key: K) -> impl 'a + Iterator<Item = &'a V> { 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. 96 #[inline] get_by_key_enumerated(&'a self, key: K) -> impl '_ + Iterator<Item = (I, &V)>97 pub fn get_by_key_enumerated(&'a self, key: K) -> impl '_ + Iterator<Item = (I, &V)> { 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 }) 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> { eq(&self, other: &Self) -> bool108 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 { hash<H: Hasher>(&self, hasher: &mut H)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 { hash_stable(&self, ctx: &mut C, hasher: &mut StableHasher)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> { from_iter<J>(iter: J) -> Self where J: IntoIterator<Item = (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 index(&self, idx: I) -> &Self::Output151 fn index(&self, idx: I) -> &Self::Output { 152 &self.items[idx].1 153 } 154 } 155 156 #[cfg(tests)] 157 mod tests; 158