Rollup merge of #116560 - ouz-a:efficient_ids, r=oli-obk

In smir use `FxIndexMap` to store indexed ids

Previously we used `vec` for storing indexed types, which is fine for small cases but will lead to huge performance issues when we use `smir` for real world cases.

Addresses https://github.com/rust-lang/project-stable-mir/issues/35

r? ``@oli-obk``

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

@@ -3,24 +3,27 @@
 //! 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;
 use rustc_driver::{Callbacks, Compilation, RunCompiler};
 use rustc_interface::{interface, Queries};
 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[index.0]
+        &self.def_ids[index]
     }
 }
 
@@ -29,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[index.0]
+        &self.spans[index]
     }
 }
 
@@ -95,36 +98,15 @@ 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);
-            }
-        }
-        let id = self.def_ids.len();
-        self.def_ids.push(did);
-        stable_mir::DefId(id)
+        self.def_ids.create_or_fetch(did)
     }
 
     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)
+        self.alloc_ids.create_or_fetch(aid)
     }
 
     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);
-            }
-        }
-        let id = self.spans.len();
-        self.spans.push(span);
-        stable_mir::ty::Span(id)
+        self.spans.create_or_fetch(span)
     }
 }
 
@@ -134,7 +116,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: 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,
     );
 }
@@ -197,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,6 +7,7 @@
 //!
 //! 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_hir as hir;
@@ -201,9 +202,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: 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;
+}