rustc_data_structures/sorted_map/
index_map.rs

1//! A variant of `SortedMap` that preserves insertion order.
2
3use std::hash::{Hash, Hasher};
4
5use rustc_index::{Idx, IndexVec};
6
7use crate::stable_hasher::{HashStable, StableHasher};
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)]
28pub 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
36impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
37    #[inline]
38    pub fn new() -> Self {
39        SortedIndexMultiMap { items: IndexVec::new(), idx_sorted_by_item_key: Vec::new() }
40    }
41
42    #[inline]
43    pub fn len(&self) -> usize {
44        self.items.len()
45    }
46
47    #[inline]
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.
53    #[inline]
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]
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]
66    pub fn iter(&self) -> impl '_ + DoubleEndedIterator<Item = (&K, &V)> {
67        self.items.iter().map(|(k, v)| (k, v))
68    }
69
70    /// Returns an iterator over the items in the map in insertion order along with their indices.
71    #[inline]
72    pub fn iter_enumerated(&self) -> impl '_ + DoubleEndedIterator<Item = (I, (&K, &V))> {
73        self.items.iter_enumerated().map(|(i, (k, v))| (i, (k, v)))
74    }
75
76    /// Returns the item in the map with the given index.
77    #[inline]
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]
87    pub fn get_by_key(&self, key: K) -> impl Iterator<Item = &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]
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))
102        })
103    }
104
105    #[inline]
106    pub fn contains_key(&self, key: K) -> bool {
107        self.get_by_key(key).next().is_some()
108    }
109}
110
111impl<I: Idx, K: Eq, V: Eq> Eq for SortedIndexMultiMap<I, K, V> {}
112impl<I: Idx, K: PartialEq, V: PartialEq> PartialEq for SortedIndexMultiMap<I, K, V> {
113    fn eq(&self, other: &Self) -> bool {
114        // No need to compare the sorted index. If the items are the same, the index will be too.
115        self.items == other.items
116    }
117}
118
119impl<I: Idx, K, V> Hash for SortedIndexMultiMap<I, K, V>
120where
121    K: Hash,
122    V: Hash,
123{
124    fn hash<H: Hasher>(&self, hasher: &mut H) {
125        self.items.hash(hasher)
126    }
127}
128
129impl<I: Idx, K, V, C> HashStable<C> for SortedIndexMultiMap<I, K, V>
130where
131    K: HashStable<C>,
132    V: HashStable<C>,
133{
134    fn hash_stable(&self, ctx: &mut C, hasher: &mut StableHasher) {
135        let SortedIndexMultiMap {
136            items,
137            // We can ignore this field because it is not observable from the outside.
138            idx_sorted_by_item_key: _,
139        } = self;
140
141        items.hash_stable(ctx, hasher)
142    }
143}
144
145impl<I: Idx, K: Ord, V> FromIterator<(K, V)> for SortedIndexMultiMap<I, K, V> {
146    fn from_iter<J>(iter: J) -> Self
147    where
148        J: IntoIterator<Item = (K, V)>,
149    {
150        let items = IndexVec::<I, _>::from_iter(iter);
151        let mut idx_sorted_by_item_key: Vec<_> = items.indices().collect();
152
153        // `sort_by_key` is stable, so insertion order is preserved for duplicate items.
154        idx_sorted_by_item_key.sort_by_key(|&idx| &items[idx].0);
155
156        SortedIndexMultiMap { items, idx_sorted_by_item_key }
157    }
158}
159
160impl<I: Idx, K, V> std::ops::Index<I> for SortedIndexMultiMap<I, K, V> {
161    type Output = V;
162
163    fn index(&self, idx: I) -> &Self::Output {
164        &self.items[idx].1
165    }
166}
rustc_data_structures/sorted_map/index_map.rs

rustc_data_structures/sorted_map/
index_map.rs
