add new wrapper for FxIndexMap

compiler/rustc_smir/src/rustc_internal/mod.rs

@@ -3,8 +3,6 @@
 //! For that, we define APIs that will temporarily be public to 3P that exposes rustc internal APIs
 //! until stable MIR is complete.
 
-use std::ops::{ControlFlow, Index};
-
 use crate::rustc_internal;
 use crate::rustc_smir::Tables;
 use rustc_data_structures::fx;
@@ -14,14 +12,18 @@ use rustc_middle::mir::interpret::AllocId;
 use rustc_middle::ty::TyCtxt;
 use rustc_span::def_id::{CrateNum, DefId};
 use rustc_span::Span;
+use stable_mir::ty::IndexedVal;
 use stable_mir::CompilerError;
+use std::fmt::Debug;
+use std::hash::Hash;
+use std::ops::{ControlFlow, Index};
 
 impl<'tcx> Index<stable_mir::DefId> for Tables<'tcx> {
     type Output = DefId;
 
     #[inline(always)]
     fn index(&self, index: stable_mir::DefId) -> &Self::Output {
-        &self.def_ids.get_index(index.0).unwrap().0
+        &self.def_ids[index]
     }
 }
 
@@ -30,7 +32,7 @@ impl<'tcx> Index<stable_mir::ty::Span> for Tables<'tcx> {
 
     #[inline(always)]
     fn index(&self, index: stable_mir::ty::Span) -> &Self::Output {
-        &self.spans.get_index(index.0).unwrap().0
+        &self.spans[index]
     }
 }
 
@@ -96,33 +98,15 @@ impl<'tcx> Tables<'tcx> {
     }
 
     fn create_def_id(&mut self, did: DefId) -> stable_mir::DefId {
-        if let Some(i) = self.def_ids.get(&did) {
-            return *i;
-        } else {
-            let id = self.def_ids.len();
-            self.def_ids.insert(did, stable_mir::DefId(id));
-            stable_mir::DefId(id)
-        }
+        self.def_ids.create_or_fetch(did)
     }
 
     fn create_alloc_id(&mut self, aid: AllocId) -> stable_mir::AllocId {
-        if let Some(i) = self.alloc_ids.get(&aid) {
-            return *i;
-        } else {
-            let id = self.def_ids.len();
-            self.alloc_ids.insert(aid, stable_mir::AllocId(id));
-            stable_mir::AllocId(id)
-        }
+        self.alloc_ids.create_or_fetch(aid)
     }
 
     pub(crate) fn create_span(&mut self, span: Span) -> stable_mir::ty::Span {
-        if let Some(i) = self.spans.get(&span) {
-            return *i;
-        } else {
-            let id = self.spans.len();
-            self.spans.insert(span, stable_mir::ty::Span(id));
-            stable_mir::ty::Span(id)
-        }
+        self.spans.create_or_fetch(span)
     }
 }
 
@@ -134,9 +118,9 @@ pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) {
     stable_mir::run(
         Tables {
             tcx,
-            def_ids: fx::FxIndexMap::default(),
-            alloc_ids: fx::FxIndexMap::default(),
-            spans: fx::FxIndexMap::default(),
+            def_ids: rustc_internal::IndexMap { index_map: fx::FxIndexMap::default() },
+            alloc_ids: rustc_internal::IndexMap { index_map: fx::FxIndexMap::default() },
+            spans: rustc_internal::IndexMap { index_map: fx::FxIndexMap::default() },
             types: vec![],
         },
         f,
@@ -201,3 +185,29 @@ where
         })
     }
 }
+
+/// Simmilar to rustc's `FxIndexMap`, `IndexMap` with extra
+/// safety features added.
+pub struct IndexMap<K, V> {
+    index_map: fx::FxIndexMap<K, V>,
+}
+
+impl<K: PartialEq + Hash + Eq, V: Copy + Debug + PartialEq + IndexedVal> IndexMap<K, V> {
+    pub fn create_or_fetch(&mut self, key: K) -> V {
+        let len = self.index_map.len();
+        let v = self.index_map.entry(key).or_insert(V::to_val(len));
+        *v
+    }
+}
+
+impl<K: PartialEq + Hash + Eq, V: Copy + Debug + PartialEq + IndexedVal> Index<V>
+    for IndexMap<K, V>
+{
+    type Output = K;
+
+    fn index(&self, index: V) -> &Self::Output {
+        let (k, v) = self.index_map.get_index(index.to_index()).unwrap();
+        assert_eq!(*v, index, "Provided value doesn't match with indexed value");
+        k
+    }
+}

compiler/rustc_smir/src/rustc_smir/mod.rs

@@ -7,9 +7,9 @@
 //!
 //! For now, we are developing everything inside `rustc`, thus, we keep this module private.
 
+use crate::rustc_internal::IndexMap;
 use crate::rustc_smir::hir::def::DefKind;
 use crate::rustc_smir::stable_mir::ty::{BoundRegion, EarlyBoundRegion, Region};
-use rustc_data_structures::fx::FxIndexMap;
 use rustc_hir as hir;
 use rustc_middle::mir;
 use rustc_middle::mir::interpret::{alloc_range, AllocId};
@@ -195,9 +195,9 @@ impl<S, R: PartialEq> PartialEq<R> for MaybeStable<S, R> {
 
 pub struct Tables<'tcx> {
     pub tcx: TyCtxt<'tcx>,
-    pub def_ids: FxIndexMap<DefId, stable_mir::DefId>,
-    pub alloc_ids: FxIndexMap<AllocId, stable_mir::AllocId>,
-    pub spans: FxIndexMap<rustc_span::Span, Span>,
+    pub def_ids: IndexMap<DefId, stable_mir::DefId>,
+    pub alloc_ids: IndexMap<AllocId, stable_mir::AllocId>,
+    pub spans: IndexMap<rustc_span::Span, Span>,
     pub types: Vec<MaybeStable<stable_mir::ty::TyKind, Ty<'tcx>>>,
 }
 

compiler/stable_mir/src/lib.rs

@@ -22,7 +22,8 @@ use std::fmt;
 use std::fmt::Debug;
 
 use self::ty::{
-    GenericPredicates, Generics, ImplDef, ImplTrait, Span, TraitDecl, TraitDef, Ty, TyKind,
+    GenericPredicates, Generics, ImplDef, ImplTrait, IndexedVal, Span, TraitDecl, TraitDef, Ty,
+    TyKind,
 };
 
 #[macro_use]
@@ -41,7 +42,7 @@ pub type CrateNum = usize;
 
 /// A unique identification number for each item accessible for the current compilation unit.
 #[derive(Clone, Copy, PartialEq, Eq)]
-pub struct DefId(pub usize);
+pub struct DefId(usize);
 
 impl Debug for DefId {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -52,9 +53,28 @@ impl Debug for DefId {
     }
 }
 
+impl IndexedVal for DefId {
+    fn to_val(index: usize) -> Self {
+        DefId(index)
+    }
+
+    fn to_index(&self) -> usize {
+        self.0
+    }
+}
+
 /// A unique identification number for each provenance
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
-pub struct AllocId(pub usize);
+pub struct AllocId(usize);
+
+impl IndexedVal for AllocId {
+    fn to_val(index: usize) -> Self {
+        AllocId(index)
+    }
+    fn to_index(&self) -> usize {
+        self.0
+    }
+}
 
 /// A list of crate items.
 pub type CrateItems = Vec<CrateItem>;

compiler/stable_mir/src/ty.rs

@@ -75,7 +75,7 @@ pub struct Placeholder<T> {
 }
 
 #[derive(Clone, Copy, PartialEq, Eq)]
-pub struct Span(pub usize);
+pub struct Span(usize);
 
 impl Debug for Span {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
@@ -86,6 +86,15 @@ impl Debug for Span {
     }
 }
 
+impl IndexedVal for Span {
+    fn to_val(index: usize) -> Self {
+        Span(index)
+    }
+    fn to_index(&self) -> usize {
+        self.0
+    }
+}
+
 #[derive(Clone, Debug)]
 pub enum TyKind {
     RigidTy(RigidTy),
@@ -565,3 +574,9 @@ pub enum ImplPolarity {
     Negative,
     Reservation,
 }
+
+pub trait IndexedVal {
+    fn to_val(index: usize) -> Self;
+
+    fn to_index(&self) -> usize;
+}