compiler/rustc_data_structures/src/sorted_map/index_map.rs

   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     pub fn new() -> Self {
  38         SortedIndexMultiMap { items: IndexVec::new(), idx_sorted_by_item_key: Vec::new() }
  39     }
  40
  41     pub fn len(&self) -> usize {
  42         self.items.len()
  43     }
  44
  45     pub fn is_empty(&self) -> bool {
  46         self.items.is_empty()
  47     }
  48
  49     /// Returns an iterator over the items in the map in insertion order.
  50     pub fn into_iter(self) -> impl DoubleEndedIterator<Item = (K, V)> {
  51         self.items.into_iter()
  52     }
  53
  54     /// Returns an iterator over the items in the map in insertion order along with their indices.
  55     pub fn into_iter_enumerated(self) -> impl DoubleEndedIterator<Item = (I, (K, V))> {
  56         self.items.into_iter_enumerated()
  57     }
  58
  59     /// Returns an iterator over the items in the map in insertion order.
  60     pub fn iter(&self) -> impl '_ + DoubleEndedIterator<Item = (&K, &V)> {
  61         self.items.iter().map(|(ref k, ref v)| (k, v))
  62     }
  63
  64     /// Returns an iterator over the items in the map in insertion order along with their indices.
  65     pub fn iter_enumerated(&self) -> impl '_ + DoubleEndedIterator<Item = (I, (&K, &V))> {
  66         self.items.iter_enumerated().map(|(i, (ref k, ref v))| (i, (k, v)))
  67     }
  68
  69     /// Returns the item in the map with the given index.
  70     pub fn get(&self, idx: I) -> Option<&(K, V)> {
  71         self.items.get(idx)
  72     }
  73
  74     /// Returns an iterator over the items in the map that are equal to `key`.
  75     ///
  76     /// If there are multiple items that are equivalent to `key`, they will be yielded in
  77     /// insertion order.
  78     pub fn get_by_key(&'a self, key: K) -> impl 'a + Iterator<Item = &'a V> {
  79         self.get_by_key_enumerated(key).map(|(_, v)| v)
  80     }
  81
  82     /// Returns an iterator over the items in the map that are equal to `key` along with their
  83     /// indices.
  84     ///
  85     /// If there are multiple items that are equivalent to `key`, they will be yielded in
  86     /// insertion order.
  87     pub fn get_by_key_enumerated(&'a self, key: K) -> impl '_ + Iterator<Item = (I, &V)> {
  88         let lower_bound = self.idx_sorted_by_item_key.partition_point(|&i| self.items[i].0 < key);
  89         self.idx_sorted_by_item_key[lower_bound..].iter().map_while(move |&i| {
  90             let (k, v) = &self.items[i];
  91             (k == &key).then_some((i, v))
  92         })
  93     }
  94 }
  95
  96 impl<I: Idx, K: Eq, V: Eq> Eq for SortedIndexMultiMap<I, K, V> {}
  97 impl<I: Idx, K: PartialEq, V: PartialEq> PartialEq for SortedIndexMultiMap<I, K, V> {
  98     fn eq(&self, other: &Self) -> bool {
  99         // No need to compare the sorted index. If the items are the same, the index will be too.
 100         self.items == other.items
 101     }
 102 }
 103
 104 impl<I: Idx, K, V> Hash for SortedIndexMultiMap<I, K, V>
 105 where
 106     K: Hash,
 107     V: Hash,
 108 {
 109     fn hash<H: Hasher>(&self, hasher: &mut H) {
 110         self.items.hash(hasher)
 111     }
 112 }
 113 impl<I: Idx, K, V, C> HashStable<C> for SortedIndexMultiMap<I, K, V>
 114 where
 115     K: HashStable<C>,
 116     V: HashStable<C>,
 117 {
 118     fn hash_stable(&self, ctx: &mut C, hasher: &mut StableHasher) {
 119         self.items.hash_stable(ctx, hasher)
 120     }
 121 }
 122
 123 impl<I: Idx, K: Ord, V> FromIterator<(K, V)> for SortedIndexMultiMap<I, K, V> {
 124     fn from_iter<J>(iter: J) -> Self
 125     where
 126         J: IntoIterator<Item = (K, V)>,
 127     {
 128         let items = IndexVec::from_iter(iter);
 129         let mut idx_sorted_by_item_key: Vec<_> = items.indices().collect();
 130
 131         // `sort_by_key` is stable, so insertion order is preserved for duplicate items.
 132         idx_sorted_by_item_key.sort_by_key(|&idx| &items[idx].0);
 133
 134         SortedIndexMultiMap { items, idx_sorted_by_item_key }
 135     }
 136 }
 137
 138 impl<I: Idx, K, V> std::ops::Index<I> for SortedIndexMultiMap<I, K, V> {
 139     type Output = V;
 140
 141     fn index(&self, idx: I) -> &Self::Output {
 142         &self.items[idx].1
 143     }
 144 }
 145
 146 #[cfg(tests)]
 147 mod tests;