1 //! A variant of `SortedMap` that preserves insertion order.
3 use std::hash::{Hash, Hasher};
4 use std::iter::FromIterator;
6 use crate::stable_hasher::{HashStable, StableHasher};
7 use rustc_index::vec::{Idx, IndexVec};
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.
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.
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.
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
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)>,
32 /// Indices of the items in the set, sorted by the item's key.
33 idx_sorted_by_item_key: Vec<I>,
36 impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
38 pub fn new() -> Self {
39 SortedIndexMultiMap { items: IndexVec::new(), idx_sorted_by_item_key: Vec::new() }
43 pub fn len(&self) -> usize {
48 pub fn is_empty(&self) -> bool {
52 /// Returns an iterator over the items in the map in insertion order.
54 pub fn into_iter(self) -> impl DoubleEndedIterator<Item = (K, V)> {
55 self.items.into_iter()
58 /// Returns an iterator over the items in the map in insertion order along with their indices.
60 pub fn into_iter_enumerated(self) -> impl DoubleEndedIterator<Item = (I, (K, V))> {
61 self.items.into_iter_enumerated()
64 /// Returns an iterator over the items in the map in insertion order.
66 pub fn iter(&self) -> impl '_ + DoubleEndedIterator<Item = (&K, &V)> {
67 self.items.iter().map(|(ref k, ref v)| (k, v))
70 /// Returns an iterator over the items in the map in insertion order along with their indices.
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)))
76 /// Returns the item in the map with the given index.
78 pub fn get(&self, idx: I) -> Option<&(K, V)> {
82 /// Returns an iterator over the items in the map that are equal to `key`.
84 /// If there are multiple items that are equivalent to `key`, they will be yielded in
87 pub fn get_by_key(&self, key: K) -> impl Iterator<Item = &V> + '_ {
88 self.get_by_key_enumerated(key).map(|(_, v)| v)
91 /// Returns an iterator over the items in the map that are equal to `key` along with their
94 /// If there are multiple items that are equivalent to `key`, they will be yielded in
97 pub fn get_by_key_enumerated(&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))
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
114 impl<I: Idx, K, V> Hash for SortedIndexMultiMap<I, K, V>
119 fn hash<H: Hasher>(&self, hasher: &mut H) {
120 self.items.hash(hasher)
123 impl<I: Idx, K, V, C> HashStable<C> for SortedIndexMultiMap<I, K, V>
128 fn hash_stable(&self, ctx: &mut C, hasher: &mut StableHasher) {
129 self.items.hash_stable(ctx, hasher)
133 impl<I: Idx, K: Ord, V> FromIterator<(K, V)> for SortedIndexMultiMap<I, K, V> {
134 fn from_iter<J>(iter: J) -> Self
136 J: IntoIterator<Item = (K, V)>,
138 let items = IndexVec::from_iter(iter);
139 let mut idx_sorted_by_item_key: Vec<_> = items.indices().collect();
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);
144 SortedIndexMultiMap { items, idx_sorted_by_item_key }
148 impl<I: Idx, K, V> std::ops::Index<I> for SortedIndexMultiMap<I, K, V> {
151 fn index(&self, idx: I) -> &Self::Output {