From d50ca3cbeef4204ec7f163813f3ff7253c44b5c2 Mon Sep 17 00:00:00 2001
From: Camille GILLOT <gillot.camille@gmail.com>
Date: Tue, 19 May 2020 12:05:19 +0200
Subject: [PATCH] Introduce HirIdVec.

---
 compiler/rustc_hir/src/hir_id.rs | 53 ++++++++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/compiler/rustc_hir/src/hir_id.rs b/compiler/rustc_hir/src/hir_id.rs
index e24eb5e4490..dc7cf40bd90 100644
--- a/compiler/rustc_hir/src/hir_id.rs
+++ b/compiler/rustc_hir/src/hir_id.rs
@@ -1,4 +1,5 @@
 use crate::def_id::{LocalDefId, CRATE_DEF_INDEX};
+use rustc_index::vec::IndexVec;
 use std::fmt;
 
 /// Uniquely identifies a node in the HIR of the current crate. It is
@@ -61,3 +62,55 @@ pub const CRATE_HIR_ID: HirId = HirId {
     owner: LocalDefId { local_def_index: CRATE_DEF_INDEX },
     local_id: ItemLocalId::from_u32(0),
 };
+
+#[derive(Clone, Default, Debug, Encodable, Decodable)]
+pub struct HirIdVec<T> {
+    map: IndexVec<LocalDefId, IndexVec<ItemLocalId, T>>,
+}
+
+impl<T> HirIdVec<T> {
+    pub fn push_owner(&mut self, id: LocalDefId) {
+        self.map.ensure_contains_elem(id, IndexVec::new);
+    }
+
+    pub fn push(&mut self, id: HirId, value: T) {
+        if id.local_id == ItemLocalId::from_u32(0) {
+            self.push_owner(id.owner);
+        }
+        let submap = &mut self.map[id.owner];
+        let _ret_id = submap.push(value);
+        debug_assert_eq!(_ret_id, id.local_id);
+    }
+
+    pub fn get(&self, id: HirId) -> Option<&T> {
+        self.map.get(id.owner)?.get(id.local_id)
+    }
+
+    pub fn get_owner(&self, id: LocalDefId) -> &IndexVec<ItemLocalId, T> {
+        &self.map[id]
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = &T> {
+        self.map.iter().flat_map(|la| la.iter())
+    }
+
+    pub fn iter_enumerated(&self) -> impl Iterator<Item = (HirId, &T)> {
+        self.map.iter_enumerated().flat_map(|(owner, la)| {
+            la.iter_enumerated().map(move |(local_id, attr)| (HirId { owner, local_id }, attr))
+        })
+    }
+}
+
+impl<T> std::ops::Index<HirId> for HirIdVec<T> {
+    type Output = T;
+
+    fn index(&self, id: HirId) -> &T {
+        &self.map[id.owner][id.local_id]
+    }
+}
+
+impl<T> std::ops::IndexMut<HirId> for HirIdVec<T> {
+    fn index_mut(&mut self, id: HirId) -> &mut T {
+        &mut self.map[id.owner][id.local_id]
+    }
+}