compiler/rustc_data_structures/src/sorted_map/index_map.rs

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