Add stable_mir::DefId as new type wrapper

3 files changed, 23 insertions, 15 deletions

diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index ebacb7cce83..f9d71d77efa 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -4,6 +4,7 @@
 //! until stable MIR is complete.
 
 use std::fmt::Debug;
+use std::ops::Index;
 use std::string::ToString;
 
 use crate::{
@@ -67,21 +68,30 @@ pub fn impl_def(did: DefId) -> stable_mir::ty::ImplDef {
     with_tables(|t| t.impl_def(did))
 }
 
+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[index.0]
+    }
+}
+
 impl<'tcx> Tables<'tcx> {
     pub fn item_def_id(&self, item: &stable_mir::CrateItem) -> DefId {
-        self.def_ids[item.0]
+        self[item.0]
     }
 
     pub fn trait_def_id(&self, trait_def: &stable_mir::ty::TraitDef) -> DefId {
-        self.def_ids[trait_def.0]
+        self[trait_def.0]
     }
 
     pub fn impl_trait_def_id(&self, impl_def: &stable_mir::ty::ImplDef) -> DefId {
-        self.def_ids[impl_def.0]
+        self[impl_def.0]
     }
 
     pub fn generic_def_id(&self, generic_def: &stable_mir::ty::GenericDef) -> DefId {
-        self.def_ids[generic_def.0]
+        self[generic_def.0]
     }
 
     pub fn crate_item(&mut self, did: DefId) -> stable_mir::CrateItem {
@@ -140,12 +150,12 @@ impl<'tcx> Tables<'tcx> {
         // FIXME: this becomes inefficient when we have too many ids
         for (i, &d) in self.def_ids.iter().enumerate() {
             if d == did {
-                return i;
+                return stable_mir::DefId(i);
             }
         }
         let id = self.def_ids.len();
         self.def_ids.push(did);
-        id
+        stable_mir::DefId(id)
     }
 }
 
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 2bf4fd23baf..0efd44dc4f0 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -10,8 +10,8 @@
 use crate::rustc_internal::{self, opaque};
 use crate::stable_mir::mir::{CopyNonOverlapping, UserTypeProjection, VariantIdx};
 use crate::stable_mir::ty::{
-    allocation_filter, new_allocation, Const, FloatTy, GenericParamDef, IntTy,
-    Movability, RigidTy, TyKind, UintTy,
+    allocation_filter, new_allocation, Const, FloatTy, GenericParamDef, IntTy, Movability, RigidTy,
+    TyKind, UintTy,
 };
 use crate::stable_mir::{self, Context};
 use rustc_hir as hir;
@@ -103,16 +103,13 @@ impl<'tcx> Context for Tables<'tcx> {
     }
 
     fn generics_of(&mut self, def_id: stable_mir::DefId) -> stable_mir::ty::Generics {
-        let def_id = self.def_ids[def_id];
+        let def_id = self[def_id];
         let generics = self.tcx.generics_of(def_id);
         generics.stable(self)
     }
 
-    fn predicates_of(
-        &mut self,
-        def_id: stable_mir::DefId,
-    ) -> stable_mir::GenericPredicates {
-        let def_id = self.def_ids[def_id];
+    fn predicates_of(&mut self, def_id: stable_mir::DefId) -> stable_mir::GenericPredicates {
+        let def_id = self[def_id];
         let ty::GenericPredicates { parent, predicates } = self.tcx.predicates_of(def_id);
         stable_mir::GenericPredicates {
             parent: parent.map(|did| self.trait_def(did)),
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index d891f883468..0eee5cc91c7 100644
--- a/compiler/rustc_smir/src/stable_mir/mod.rs
+++ b/compiler/rustc_smir/src/stable_mir/mod.rs
@@ -29,7 +29,8 @@ pub type Symbol = String;
 pub type CrateNum = usize;
 
 /// A unique identification number for each item accessible for the current compilation unit.
-pub type DefId = usize;
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub struct DefId(pub(crate) usize);
 
 /// A list of crate items.
 pub type CrateItems = Vec<CrateItem>;