add a `VecMap` data structure

3 files changed, 114 insertions, 0 deletions

diff --git a/src/librustc_data_structures/lib.rs b/src/librustc_data_structures/lib.rs
index a1d7ab8856d..1829da25087 100644
--- a/src/librustc_data_structures/lib.rs
+++ b/src/librustc_data_structures/lib.rs
@@ -95,6 +95,7 @@ pub mod tiny_list;
 pub mod thin_vec;
 pub mod transitive_relation;
 pub use ena::unify;
+pub mod vec_map;
 pub mod vec_linked_list;
 pub mod work_queue;
 pub mod fingerprint;
diff --git a/src/librustc_data_structures/vec_map/mod.rs b/src/librustc_data_structures/vec_map/mod.rs
new file mode 100644
index 00000000000..9d9e54bfc7b
--- /dev/null
+++ b/src/librustc_data_structures/vec_map/mod.rs
@@ -0,0 +1,85 @@
+#[cfg(test)]
+mod test;
+
+/// A (multi-)map based on a sorted vector. This uses binary search to
+/// find the starting index for a given element and can be a fairly
+/// efficient data structure, particularly for small-ish sets of data.
+///
+/// To use, you supply the starting vector along with a "key fn" that
+/// extracts the key from each element.
+pub struct VecMap<E, KeyFn> {
+    data: Vec<E>,
+    key_fn: KeyFn,
+}
+
+impl<E, K, KeyFn> VecMap<E, KeyFn>
+where
+    KeyFn: Fn(&E) -> K,
+    K: Ord + std::fmt::Debug,
+{
+    pub fn new(
+        mut data: Vec<E>,
+        key_fn: KeyFn,
+    ) -> Self {
+        data.sort_by_key(&key_fn);
+        Self { data, key_fn }
+    }
+
+    /// Extract the first index for the given key using binary search.
+    /// Returns `None` if there is no such index.
+    fn get_range(&self, key: &K) -> Option<(usize, usize)> {
+        match self.data.binary_search_by_key(key, &self.key_fn) {
+            Ok(mid) => {
+                // We get back *some* element with the given key -- so
+                // search backwards to find the *first* one.
+                //
+                // (It'd be more efficient to use a "galloping" search
+                // here, but it's not really worth it for small-ish
+                // amounts of data.)
+                let mut start = mid;
+                while start > 0 {
+                    if (self.key_fn)(&self.data[start - 1]) == *key {
+                        start -= 1;
+                    } else {
+                        break;
+                    }
+                }
+
+                // Now search forward to find the *last* one.
+                //
+                // (It'd be more efficient to use a "galloping" search
+                // here, but it's not really worth it for small-ish
+                // amounts of data.)
+                let mut end = mid + 1;
+                let max = self.data.len();
+                while end < max {
+                    if (self.key_fn)(&self.data[end]) == *key {
+                        end += 1;
+                    } else {
+                        break;
+                    }
+                }
+
+                Some((start, end))
+            }
+            Err(_) => None,
+        }
+    }
+
+    /// Gets the (first) value associated with a given key.
+    pub fn get_first(&self, key: &K) -> Option<&E> {
+        let (start, _) = self.get_range(key)?;
+        Some(&self.data[start])
+    }
+
+    /// Gets a slice of values associated with the given key.
+    pub fn get_all(&self, key: &K) -> &[E] {
+        let (start, end) = self.get_range(key).unwrap_or((0, 0));
+        &self.data[start..end]
+    }
+
+    /// Gets a slice of values associated with the given key.
+    pub fn get_iter<'k>(&'k self, key: &'k K) -> impl Iterator<Item = &'k E> {
+        self.get_all(key).iter()
+    }
+}
diff --git a/src/librustc_data_structures/vec_map/test.rs b/src/librustc_data_structures/vec_map/test.rs
new file mode 100644
index 00000000000..9e4f581a8cf
--- /dev/null
+++ b/src/librustc_data_structures/vec_map/test.rs
@@ -0,0 +1,28 @@
+use super::*;
+
+type Element = (usize, &'static str);
+
+fn test_map() -> VecMap<Element, impl Fn(&Element) -> usize> {
+    let data = vec![(3, "three-a"), (0, "zero"), (3, "three-b"), (22, "twenty-two")];
+    VecMap::new(data, |&(key, _)| key)
+}
+
+#[test]
+fn get_first() {
+    let map = test_map();
+    assert_eq!(map.get_first(&0), Some(&(0, "zero")));
+    assert_eq!(map.get_first(&1), None);
+    assert_eq!(map.get_first(&3), Some(&(3, "three-a")));
+    assert_eq!(map.get_first(&22), Some(&(22, "twenty-two")));
+    assert_eq!(map.get_first(&23), None);
+}
+
+#[test]
+fn get_all() {
+    let map = test_map();
+    assert_eq!(map.get_all(&0), &[(0, "zero")]);
+    assert_eq!(map.get_all(&1), &[]);
+    assert_eq!(map.get_all(&3), &[(3, "three-a"), (3, "three-b")]);
+    assert_eq!(map.get_all(&22), &[(22, "twenty-two")]);
+    assert_eq!(map.get_all(&23), &[]);
+}