In smir use FxIndexMap to store indexed ids

Cargo.lock

@@ -4504,6 +4504,7 @@ dependencies = [
 name = "rustc_smir"
 version = "0.0.0"
 dependencies = [
+ "rustc_data_structures",
  "rustc_driver",
  "rustc_hir",
  "rustc_interface",

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" }

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
     }
 }
 
@@ -29,7 +30,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[index.0]
+        &self.spans.get_index(index.0).unwrap().0
     }
 }
 
@@ -95,36 +96,33 @@ 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 {
-        // FIXME: this becomes inefficient when we have too many ids
-        if let Some(i) = self.alloc_ids.iter().position(|a| *a == aid) {
-            return stable_mir::AllocId(i);
-        };
-        let id = self.def_ids.len();
-        self.alloc_ids.push(aid);
-        stable_mir::AllocId(id)
+        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)
+        }
     }
 
     pub(crate) fn create_span(&mut self, span: Span) -> stable_mir::ty::Span {
-        for (i, &sp) in self.spans.iter().enumerate() {
-            if sp == span {
-                return stable_mir::ty::Span(i);
-            }
+        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)
         }
-        let id = self.spans.len();
-        self.spans.push(span);
-        stable_mir::ty::Span(id)
     }
 }
 
@@ -134,7 +132,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: fx::FxIndexMap::default(),
+            spans: fx::FxIndexMap::default(),
+            types: vec![],
+        },
         f,
     );
 }

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,9 +195,9 @@ impl<S, R: PartialEq> PartialEq<R> for MaybeStable<S, R> {
 
 pub struct Tables<'tcx> {
     pub tcx: TyCtxt<'tcx>,
-    pub def_ids: Vec<DefId>,
-    pub alloc_ids: Vec<AllocId>,
-    pub spans: Vec<rustc_span::Span>,
+    pub def_ids: FxIndexMap<DefId, stable_mir::DefId>,
+    pub alloc_ids: FxIndexMap<AllocId, stable_mir::AllocId>,
+    pub spans: FxIndexMap<rustc_span::Span, Span>,
     pub types: Vec<MaybeStable<stable_mir::ty::TyKind, Ty<'tcx>>>,
 }