From 17f6df9c63695d8206f30e7a23e56d676b95c392 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Tue, 24 Oct 2023 10:38:05 -0700
Subject: [PATCH] Use IndexMap for handling stable Ty

---
 .../rustc_smir/src/rustc_internal/internal.rs |  2 +-
 compiler/rustc_smir/src/rustc_internal/mod.rs |  2 +-
 compiler/rustc_smir/src/rustc_smir/mod.rs     | 11 ++----
 compiler/stable_mir/src/ty.rs                 | 37 +++++++++----------
 4 files changed, 23 insertions(+), 29 deletions(-)

diff --git a/compiler/rustc_smir/src/rustc_internal/internal.rs b/compiler/rustc_smir/src/rustc_internal/internal.rs
index 2f2dac3c06d5f..78144524ac508 100644
--- a/compiler/rustc_smir/src/rustc_internal/internal.rs
+++ b/compiler/rustc_smir/src/rustc_internal/internal.rs
@@ -46,7 +46,7 @@ impl<'tcx> RustcInternal<'tcx> for Region {
 impl<'tcx> RustcInternal<'tcx> for Ty {
     type T = InternalTy<'tcx>;
     fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
-        tables.types[self.0]
+        tables.types[*self]
     }
 }
 
diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 7c0d811fa484b..4d2a518226d89 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -167,7 +167,7 @@ pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) {
         def_ids: IndexMap::default(),
         alloc_ids: IndexMap::default(),
         spans: IndexMap::default(),
-        types: vec![],
+        types: IndexMap::default(),
         instances: IndexMap::default(),
         constants: IndexMap::default(),
     }));
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 3498b6a2f6e39..0bf0677be3bf0 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -148,7 +148,7 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
 
     fn ty_kind(&self, ty: stable_mir::ty::Ty) -> TyKind {
         let mut tables = self.0.borrow_mut();
-        tables.types[ty.0].kind().stable(&mut *tables)
+        tables.types[ty].kind().stable(&mut *tables)
     }
 
     fn generics_of(&self, def_id: stable_mir::DefId) -> stable_mir::ty::Generics {
@@ -252,19 +252,14 @@ pub struct Tables<'tcx> {
     pub(crate) def_ids: IndexMap<DefId, stable_mir::DefId>,
     pub(crate) alloc_ids: IndexMap<AllocId, stable_mir::AllocId>,
     pub(crate) spans: IndexMap<rustc_span::Span, Span>,
-    pub(crate) types: Vec<Ty<'tcx>>,
+    pub(crate) types: IndexMap<Ty<'tcx>, stable_mir::ty::Ty>,
     pub(crate) instances: IndexMap<ty::Instance<'tcx>, InstanceDef>,
     pub(crate) constants: IndexMap<mir::Const<'tcx>, ConstId>,
 }
 
 impl<'tcx> Tables<'tcx> {
     fn intern_ty(&mut self, ty: Ty<'tcx>) -> stable_mir::ty::Ty {
-        if let Some(id) = self.types.iter().position(|t| *t == ty) {
-            return stable_mir::ty::Ty(id);
-        }
-        let id = self.types.len();
-        self.types.push(ty);
-        stable_mir::ty::Ty(id)
+        self.types.create_or_fetch(ty)
     }
 
     fn intern_const(&mut self, constant: mir::Const<'tcx>) -> ConstId {
diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs
index f1af09fac66d0..0dbf6fe23aaa5 100644
--- a/compiler/stable_mir/src/ty.rs
+++ b/compiler/stable_mir/src/ty.rs
@@ -6,7 +6,7 @@ use super::{
 use crate::{Filename, Opaque};
 use std::fmt::{self, Debug, Formatter};
 
-#[derive(Copy, Clone)]
+#[derive(Copy, Clone, Eq, PartialEq, Hash)]
 pub struct Ty(pub usize);
 
 impl Debug for Ty {
@@ -52,15 +52,6 @@ impl Const {
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct ConstId(pub usize);
 
-impl IndexedVal for ConstId {
-    fn to_val(index: usize) -> Self {
-        ConstId(index)
-    }
-    fn to_index(&self) -> usize {
-        self.0
-    }
-}
-
 type Ident = Opaque;
 
 #[derive(Debug, Clone)]
@@ -136,15 +127,6 @@ pub struct LineInfo {
     pub end_col: usize,
 }
 
-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),
@@ -631,3 +613,20 @@ pub trait IndexedVal {
 
     fn to_index(&self) -> usize;
 }
+
+macro_rules! index_impl {
+    ($name:ident) => {
+        impl IndexedVal for $name {
+            fn to_val(index: usize) -> Self {
+                $name(index)
+            }
+            fn to_index(&self) -> usize {
+                self.0
+            }
+        }
+    };
+}
+
+index_impl!(ConstId);
+index_impl!(Ty);
+index_impl!(Span);