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::IndexToVal;
 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.get_index_and_assert(index.0, 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.get_index_and_assert(index.0, 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.entry(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.entry(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.entry(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,26 @@ 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 + IndexToVal> IndexMap<K, V> {
+    /// Because we are using values as indexes, this first get's the
+    /// key using index of the value provided, then asserts that value
+    /// we got from IndexMap is equal to one we provided.
+    pub fn get_index_and_assert(&self, index: usize, value: V) -> &K {
+        let (k, v) = self.index_map.get_index(index).unwrap();
+        assert_eq!(*v, value, "Provided value doesn't match with indexed value");
+        k
+    }
+
+    pub fn entry(&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
+    }
+}

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

@@ -21,6 +21,8 @@ use std::cell::Cell;
 use std::fmt;
 use std::fmt::Debug;
 
+use ty::IndexToVal;
+
 use self::ty::{
     GenericPredicates, Generics, ImplDef, ImplTrait, Span, TraitDecl, TraitDef, Ty, TyKind,
 };
@@ -52,10 +54,27 @@ impl Debug for DefId {
     }
 }
 
+impl IndexToVal for DefId {
+    fn to_val(index: usize) -> Self
+    where
+        Self: std::marker::Sized,
+    {
+        DefId(index)
+    }
+}
 /// A unique identification number for each provenance
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub struct AllocId(pub usize);
 
+impl IndexToVal for AllocId {
+    fn to_val(index: usize) -> Self
+    where
+        Self: std::marker::Sized,
+    {
+        AllocId(index)
+    }
+}
+
 /// A list of crate items.
 pub type CrateItems = Vec<CrateItem>;
 

compiler/stable_mir/src/ty.rs

@@ -86,6 +86,15 @@ impl Debug for Span {
     }
 }
 
+impl IndexToVal for Span {
+    fn to_val(index: usize) -> Self
+    where
+        Self: std::marker::Sized,
+    {
+        Span(index)
+    }
+}
+
 #[derive(Clone, Debug)]
 pub enum TyKind {
     RigidTy(RigidTy),
@@ -565,3 +574,9 @@ pub enum ImplPolarity {
     Negative,
     Reservation,
 }
+
+pub trait IndexToVal {
+    fn to_val(index: usize) -> Self
+    where
+        Self: std::marker::Sized;
+}