summaryrefslogtreecommitdiffstats
path: root/compiler/rustc_data_structures/src/sorted_map
diff options
context:
space:
mode:
Diffstat (limited to 'compiler/rustc_data_structures/src/sorted_map')
-rw-r--r--compiler/rustc_data_structures/src/sorted_map/index_map.rs154
-rw-r--r--compiler/rustc_data_structures/src/sorted_map/tests.rs222
2 files changed, 376 insertions, 0 deletions
diff --git a/compiler/rustc_data_structures/src/sorted_map/index_map.rs b/compiler/rustc_data_structures/src/sorted_map/index_map.rs
new file mode 100644
index 000000000..0ec32dc43
--- /dev/null
+++ b/compiler/rustc_data_structures/src/sorted_map/index_map.rs
@@ -0,0 +1,154 @@
+//! A variant of `SortedMap` that preserves insertion order.
+
+use std::hash::{Hash, Hasher};
+use std::iter::FromIterator;
+
+use crate::stable_hasher::{HashStable, StableHasher};
+use rustc_index::vec::{Idx, IndexVec};
+
+/// An indexed multi-map that preserves insertion order while permitting both *O*(log *n*) lookup of
+/// an item by key and *O*(1) lookup by index.
+///
+/// This data structure is a hybrid of an [`IndexVec`] and a [`SortedMap`]. Like `IndexVec`,
+/// `SortedIndexMultiMap` assigns a typed index to each item while preserving insertion order.
+/// Like `SortedMap`, `SortedIndexMultiMap` has efficient lookup of items by key. However, this
+/// is accomplished by sorting an array of item indices instead of the items themselves.
+///
+/// Unlike `SortedMap`, this data structure can hold multiple equivalent items at once, so the
+/// `get_by_key` method and its variants return an iterator instead of an `Option`. Equivalent
+/// items will be yielded in insertion order.
+///
+/// Unlike a general-purpose map like `BTreeSet` or `HashSet`, `SortedMap` and
+/// `SortedIndexMultiMap` require *O*(*n*) time to insert a single item. This is because we may need
+/// to insert into the middle of the sorted array. Users should avoid mutating this data structure
+/// in-place.
+///
+/// [`SortedMap`]: super::SortedMap
+#[derive(Clone, Debug)]
+pub struct SortedIndexMultiMap<I: Idx, K, V> {
+ /// The elements of the map in insertion order.
+ items: IndexVec<I, (K, V)>,
+
+ /// Indices of the items in the set, sorted by the item's key.
+ idx_sorted_by_item_key: Vec<I>,
+}
+
+impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
+ #[inline]
+ pub fn new() -> Self {
+ SortedIndexMultiMap { items: IndexVec::new(), idx_sorted_by_item_key: Vec::new() }
+ }
+
+ #[inline]
+ pub fn len(&self) -> usize {
+ self.items.len()
+ }
+
+ #[inline]
+ pub fn is_empty(&self) -> bool {
+ self.items.is_empty()
+ }
+
+ /// Returns an iterator over the items in the map in insertion order.
+ #[inline]
+ pub fn into_iter(self) -> impl DoubleEndedIterator<Item = (K, V)> {
+ self.items.into_iter()
+ }
+
+ /// Returns an iterator over the items in the map in insertion order along with their indices.
+ #[inline]
+ pub fn into_iter_enumerated(self) -> impl DoubleEndedIterator<Item = (I, (K, V))> {
+ self.items.into_iter_enumerated()
+ }
+
+ /// Returns an iterator over the items in the map in insertion order.
+ #[inline]
+ pub fn iter(&self) -> impl '_ + DoubleEndedIterator<Item = (&K, &V)> {
+ self.items.iter().map(|(ref k, ref v)| (k, v))
+ }
+
+ /// Returns an iterator over the items in the map in insertion order along with their indices.
+ #[inline]
+ pub fn iter_enumerated(&self) -> impl '_ + DoubleEndedIterator<Item = (I, (&K, &V))> {
+ self.items.iter_enumerated().map(|(i, (ref k, ref v))| (i, (k, v)))
+ }
+
+ /// Returns the item in the map with the given index.
+ #[inline]
+ pub fn get(&self, idx: I) -> Option<&(K, V)> {
+ self.items.get(idx)
+ }
+
+ /// Returns an iterator over the items in the map that are equal to `key`.
+ ///
+ /// If there are multiple items that are equivalent to `key`, they will be yielded in
+ /// insertion order.
+ #[inline]
+ pub fn get_by_key(&self, key: K) -> impl Iterator<Item = &V> + '_ {
+ self.get_by_key_enumerated(key).map(|(_, v)| v)
+ }
+
+ /// Returns an iterator over the items in the map that are equal to `key` along with their
+ /// indices.
+ ///
+ /// If there are multiple items that are equivalent to `key`, they will be yielded in
+ /// insertion order.
+ #[inline]
+ pub fn get_by_key_enumerated(&self, key: K) -> impl Iterator<Item = (I, &V)> + '_ {
+ let lower_bound = self.idx_sorted_by_item_key.partition_point(|&i| self.items[i].0 < key);
+ self.idx_sorted_by_item_key[lower_bound..].iter().map_while(move |&i| {
+ let (k, v) = &self.items[i];
+ (k == &key).then_some((i, v))
+ })
+ }
+}
+
+impl<I: Idx, K: Eq, V: Eq> Eq for SortedIndexMultiMap<I, K, V> {}
+impl<I: Idx, K: PartialEq, V: PartialEq> PartialEq for SortedIndexMultiMap<I, K, V> {
+ fn eq(&self, other: &Self) -> bool {
+ // No need to compare the sorted index. If the items are the same, the index will be too.
+ self.items == other.items
+ }
+}
+
+impl<I: Idx, K, V> Hash for SortedIndexMultiMap<I, K, V>
+where
+ K: Hash,
+ V: Hash,
+{
+ fn hash<H: Hasher>(&self, hasher: &mut H) {
+ self.items.hash(hasher)
+ }
+}
+impl<I: Idx, K, V, C> HashStable<C> for SortedIndexMultiMap<I, K, V>
+where
+ K: HashStable<C>,
+ V: HashStable<C>,
+{
+ fn hash_stable(&self, ctx: &mut C, hasher: &mut StableHasher) {
+ self.items.hash_stable(ctx, hasher)
+ }
+}
+
+impl<I: Idx, K: Ord, V> FromIterator<(K, V)> for SortedIndexMultiMap<I, K, V> {
+ fn from_iter<J>(iter: J) -> Self
+ where
+ J: IntoIterator<Item = (K, V)>,
+ {
+ let items = IndexVec::from_iter(iter);
+ let mut idx_sorted_by_item_key: Vec<_> = items.indices().collect();
+
+ // `sort_by_key` is stable, so insertion order is preserved for duplicate items.
+ idx_sorted_by_item_key.sort_by_key(|&idx| &items[idx].0);
+
+ SortedIndexMultiMap { items, idx_sorted_by_item_key }
+ }
+}
+
+impl<I: Idx, K, V> std::ops::Index<I> for SortedIndexMultiMap<I, K, V> {
+ type Output = V;
+
+ fn index(&self, idx: I) -> &Self::Output {
+ &self.items[idx].1
+ }
+}
diff --git a/compiler/rustc_data_structures/src/sorted_map/tests.rs b/compiler/rustc_data_structures/src/sorted_map/tests.rs
new file mode 100644
index 000000000..1e977d709
--- /dev/null
+++ b/compiler/rustc_data_structures/src/sorted_map/tests.rs
@@ -0,0 +1,222 @@
+use super::{SortedIndexMultiMap, SortedMap};
+
+#[test]
+fn test_sorted_index_multi_map() {
+ let entries: Vec<_> = vec![(2, 0), (1, 0), (2, 1), (3, 0), (2, 2)];
+ let set: SortedIndexMultiMap<usize, _, _> = entries.iter().copied().collect();
+
+ // Insertion order is preserved.
+ assert!(entries.iter().map(|(ref k, ref v)| (k, v)).eq(set.iter()));
+
+ // Indexing
+ for (i, expect) in entries.iter().enumerate() {
+ assert_eq!(set[i], expect.1);
+ }
+
+ // `get_by_key` works.
+ assert_eq!(set.get_by_key(3).copied().collect::<Vec<_>>(), vec![0]);
+ assert!(set.get_by_key(4).next().is_none());
+
+ // `get_by_key` returns items in insertion order.
+ let twos: Vec<_> = set.get_by_key_enumerated(2).collect();
+ let idxs: Vec<usize> = twos.iter().map(|(i, _)| *i).collect();
+ let values: Vec<usize> = twos.iter().map(|(_, &v)| v).collect();
+
+ assert_eq!(idxs, vec![0, 2, 4]);
+ assert_eq!(values, vec![0, 1, 2]);
+}
+
+#[test]
+fn test_insert_and_iter() {
+ let mut map = SortedMap::new();
+ let mut expected = Vec::new();
+
+ for x in 0..100 {
+ assert_eq!(map.iter().cloned().collect::<Vec<_>>(), expected);
+
+ let x = 1000 - x * 2;
+ map.insert(x, x);
+ expected.insert(0, (x, x));
+ }
+}
+
+#[test]
+fn test_get_and_index() {
+ let mut map = SortedMap::new();
+ let mut expected = Vec::new();
+
+ for x in 0..100 {
+ let x = 1000 - x;
+ if x & 1 == 0 {
+ map.insert(x, x);
+ }
+ expected.push(x);
+ }
+
+ for mut x in expected {
+ if x & 1 == 0 {
+ assert_eq!(map.get(&x), Some(&x));
+ assert_eq!(map.get_mut(&x), Some(&mut x));
+ assert_eq!(map[&x], x);
+ assert_eq!(&mut map[&x], &mut x);
+ } else {
+ assert_eq!(map.get(&x), None);
+ assert_eq!(map.get_mut(&x), None);
+ }
+ }
+}
+
+#[test]
+fn test_range() {
+ let mut map = SortedMap::new();
+ map.insert(1, 1);
+ map.insert(3, 3);
+ map.insert(6, 6);
+ map.insert(9, 9);
+
+ let keys = |s: &[(_, _)]| s.into_iter().map(|e| e.0).collect::<Vec<u32>>();
+
+ for start in 0..11 {
+ for end in 0..11 {
+ if end < start {
+ continue;
+ }
+
+ let mut expected = vec![1, 3, 6, 9];
+ expected.retain(|&x| x >= start && x < end);
+
+ assert_eq!(keys(map.range(start..end)), expected, "range = {}..{}", start, end);
+ }
+ }
+}
+
+#[test]
+fn test_offset_keys() {
+ let mut map = SortedMap::new();
+ map.insert(1, 1);
+ map.insert(3, 3);
+ map.insert(6, 6);
+
+ map.offset_keys(|k| *k += 1);
+
+ let mut expected = SortedMap::new();
+ expected.insert(2, 1);
+ expected.insert(4, 3);
+ expected.insert(7, 6);
+
+ assert_eq!(map, expected);
+}
+
+fn keys(s: SortedMap<u32, u32>) -> Vec<u32> {
+ s.into_iter().map(|(k, _)| k).collect::<Vec<u32>>()
+}
+
+fn elements(s: SortedMap<u32, u32>) -> Vec<(u32, u32)> {
+ s.into_iter().collect::<Vec<(u32, u32)>>()
+}
+
+#[test]
+fn test_remove_range() {
+ let mut map = SortedMap::new();
+ map.insert(1, 1);
+ map.insert(3, 3);
+ map.insert(6, 6);
+ map.insert(9, 9);
+
+ for start in 0..11 {
+ for end in 0..11 {
+ if end < start {
+ continue;
+ }
+
+ let mut expected = vec![1, 3, 6, 9];
+ expected.retain(|&x| x < start || x >= end);
+
+ let mut map = map.clone();
+ map.remove_range(start..end);
+
+ assert_eq!(keys(map), expected, "range = {}..{}", start, end);
+ }
+ }
+}
+
+#[test]
+fn test_remove() {
+ let mut map = SortedMap::new();
+ let mut expected = Vec::new();
+
+ for x in 0..10 {
+ map.insert(x, x);
+ expected.push((x, x));
+ }
+
+ for x in 0..10 {
+ let mut map = map.clone();
+ let mut expected = expected.clone();
+
+ assert_eq!(map.remove(&x), Some(x));
+ expected.remove(x as usize);
+
+ assert_eq!(map.iter().cloned().collect::<Vec<_>>(), expected);
+ }
+}
+
+#[test]
+fn test_insert_presorted_non_overlapping() {
+ let mut map = SortedMap::new();
+ map.insert(2, 0);
+ map.insert(8, 0);
+
+ map.insert_presorted(vec![(3, 0), (7, 0)]);
+
+ let expected = vec![2, 3, 7, 8];
+ assert_eq!(keys(map), expected);
+}
+
+#[test]
+fn test_insert_presorted_first_elem_equal() {
+ let mut map = SortedMap::new();
+ map.insert(2, 2);
+ map.insert(8, 8);
+
+ map.insert_presorted(vec![(2, 0), (7, 7)]);
+
+ let expected = vec![(2, 0), (7, 7), (8, 8)];
+ assert_eq!(elements(map), expected);
+}
+
+#[test]
+fn test_insert_presorted_last_elem_equal() {
+ let mut map = SortedMap::new();
+ map.insert(2, 2);
+ map.insert(8, 8);
+
+ map.insert_presorted(vec![(3, 3), (8, 0)]);
+
+ let expected = vec![(2, 2), (3, 3), (8, 0)];
+ assert_eq!(elements(map), expected);
+}
+
+#[test]
+fn test_insert_presorted_shuffle() {
+ let mut map = SortedMap::new();
+ map.insert(2, 2);
+ map.insert(7, 7);
+
+ map.insert_presorted(vec![(1, 1), (3, 3), (8, 8)]);
+
+ let expected = vec![(1, 1), (2, 2), (3, 3), (7, 7), (8, 8)];
+ assert_eq!(elements(map), expected);
+}
+
+#[test]
+fn test_insert_presorted_at_end() {
+ let mut map = SortedMap::new();
+ map.insert(1, 1);
+ map.insert(2, 2);
+
+ map.insert_presorted(vec![(3, 3), (8, 8)]);
+
+ let expected = vec![(1, 1), (2, 2), (3, 3), (8, 8)];
+ assert_eq!(elements(map), expected);
+}
generated by cgit v1.2.3 (git 2.39.1) at 2024-07-10 13:26:31 +0000