defids are indexmapped

4 files changed, 19 insertions, 11 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 082bb4be93c..2882ec0682e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4504,6 +4504,7 @@ dependencies = [
 name = "rustc_smir"
 version = "0.0.0"
 dependencies = [
+ "rustc_data_structures",
  "rustc_driver",
  "rustc_hir",
  "rustc_interface",
diff --git a/compiler/rustc_smir/Cargo.toml b/compiler/rustc_smir/Cargo.toml
index 3e0d6baab6a..2b77044d6bf 100644
--- a/compiler/rustc_smir/Cargo.toml
+++ b/compiler/rustc_smir/Cargo.toml
@@ -4,6 +4,7 @@ version = "0.0.0"
 edition = "2021"
 
 [dependencies]
+rustc_data_structures = { path = "../rustc_data_structures" }
 rustc_driver = { path = "../rustc_driver" }
 rustc_hir = { path = "../rustc_hir" }
 rustc_interface = { path = "../rustc_interface" }
diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 36eb2247253..2f4bc5a0499 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -7,6 +7,7 @@ use std::ops::{ControlFlow, Index};
 
 use crate::rustc_internal;
 use crate::rustc_smir::Tables;
+use rustc_data_structures::fx;
 use rustc_driver::{Callbacks, Compilation, RunCompiler};
 use rustc_interface::{interface, Queries};
 use rustc_middle::mir::interpret::AllocId;
@@ -20,7 +21,7 @@ impl<'tcx> Index<stable_mir::DefId> for Tables<'tcx> {
 
     #[inline(always)]
     fn index(&self, index: stable_mir::DefId) -> &Self::Output {
-        &self.def_ids[index.0]
+        &self.def_ids.get_index(index.0).unwrap().0
     }
 }
 
@@ -95,15 +96,13 @@ impl<'tcx> Tables<'tcx> {
     }
 
     fn create_def_id(&mut self, did: DefId) -> stable_mir::DefId {
-        // FIXME: this becomes inefficient when we have too many ids
-        for (i, &d) in self.def_ids.iter().enumerate() {
-            if d == did {
-                return stable_mir::DefId(i);
-            }
+        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)
         }
-        let id = self.def_ids.len();
-        self.def_ids.push(did);
-        stable_mir::DefId(id)
     }
 
     fn create_alloc_id(&mut self, aid: AllocId) -> stable_mir::AllocId {
@@ -134,7 +133,13 @@ pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
 
 pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) {
     stable_mir::run(
-        Tables { tcx, def_ids: vec![], alloc_ids: vec![], spans: vec![], types: vec![] },
+        Tables {
+            tcx,
+            def_ids: fx::FxIndexMap::default(),
+            alloc_ids: vec![],
+            spans: vec![],
+            types: vec![],
+        },
         f,
     );
 }
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 7d1122c2fd2..144f1625b04 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -9,6 +9,7 @@
 
 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};
@@ -194,7 +195,7 @@ impl<S, R: PartialEq> PartialEq<R> for MaybeStable<S, R> {
 
 pub struct Tables<'tcx> {
     pub tcx: TyCtxt<'tcx>,
-    pub def_ids: Vec<DefId>,
+    pub def_ids: FxIndexMap<DefId, stable_mir::DefId>,
     pub alloc_ids: Vec<AllocId>,
     pub spans: Vec<rustc_span::Span>,
     pub types: Vec<MaybeStable<stable_mir::ty::TyKind, Ty<'tcx>>>,