From 5b3cf6610b7c94d5eea715386946ba0bf564690e Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Fri, 17 Nov 2023 08:05:40 -0800
Subject: [PATCH] Add support to get virtual table allocation

---
 .../rustc_smir/src/rustc_internal/internal.rs | 26 +++++++++++++++++--
 compiler/rustc_smir/src/rustc_internal/mod.rs |  2 +-
 compiler/rustc_smir/src/rustc_smir/mod.rs     | 19 ++++++++++++++
 compiler/stable_mir/src/lib.rs                |  3 +++
 compiler/stable_mir/src/mir/alloc.rs          |  7 +++++
 5 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/compiler/rustc_smir/src/rustc_internal/internal.rs b/compiler/rustc_smir/src/rustc_internal/internal.rs
index 988d1d7226a..fe226ef60ed 100644
--- a/compiler/rustc_smir/src/rustc_internal/internal.rs
+++ b/compiler/rustc_smir/src/rustc_internal/internal.rs
@@ -10,8 +10,9 @@ use rustc_span::Symbol;
 use stable_mir::mir::alloc::AllocId;
 use stable_mir::mir::mono::{Instance, MonoItem, StaticDef};
 use stable_mir::ty::{
-    AdtDef, Binder, BoundRegionKind, BoundTyKind, BoundVariableKind, ClosureKind, Const, FloatTy,
-    GenericArgKind, GenericArgs, IntTy, Region, RigidTy, TraitRef, Ty, UintTy,
+    AdtDef, Binder, BoundRegionKind, BoundTyKind, BoundVariableKind, ClosureKind, Const,
+    ExistentialTraitRef, FloatTy, GenericArgKind, GenericArgs, IntTy, Region, RigidTy, TraitRef,
+    Ty, UintTy,
 };
 use stable_mir::{CrateItem, DefId};
 
@@ -229,6 +230,17 @@ impl<'tcx> RustcInternal<'tcx> for BoundVariableKind {
     }
 }
 
+impl<'tcx> RustcInternal<'tcx> for ExistentialTraitRef {
+    type T = rustc_ty::ExistentialTraitRef<'tcx>;
+
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        rustc_ty::ExistentialTraitRef {
+            def_id: self.def_id.0.internal(tables),
+            args: self.generic_args.internal(tables),
+        }
+    }
+}
+
 impl<'tcx> RustcInternal<'tcx> for TraitRef {
     type T = rustc_ty::TraitRef<'tcx>;
 
@@ -277,3 +289,13 @@ where
         (*self).internal(tables)
     }
 }
+impl<'tcx, T> RustcInternal<'tcx> for Option<T>
+where
+    T: RustcInternal<'tcx>,
+{
+    type T = Option<T::T>;
+
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        self.as_ref().map(|inner| inner.internal(tables))
+    }
+}
diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index b37c9a1f632..fed21f16a4e 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -118,7 +118,7 @@ impl<'tcx> Tables<'tcx> {
         self.def_ids.create_or_fetch(did)
     }
 
-    fn create_alloc_id(&mut self, aid: AllocId) -> stable_mir::mir::alloc::AllocId {
+    pub(crate) fn create_alloc_id(&mut self, aid: AllocId) -> stable_mir::mir::alloc::AllocId {
         self.alloc_ids.create_or_fetch(aid)
     }
 
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 509706387db..403c7703deb 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -331,6 +331,18 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         tables.tcx.global_alloc(alloc_id).stable(&mut *tables)
     }
 
+    fn vtable_allocation(
+        &self,
+        global_alloc: &GlobalAlloc,
+    ) -> Option<stable_mir::mir::alloc::AllocId> {
+        let mut tables = self.0.borrow_mut();
+        let GlobalAlloc::VTable(ty, trait_ref) = global_alloc else { return None };
+        let alloc_id = tables
+            .tcx
+            .vtable_allocation((ty.internal(&mut *tables), trait_ref.internal(&mut *tables)));
+        Some(alloc_id.stable(&mut *tables))
+    }
+
     fn usize_to_const(&self, val: u64) -> Result<Const, Error> {
         let mut tables = self.0.borrow_mut();
         let ty = tables.tcx.types.usize;
@@ -1623,6 +1635,13 @@ impl<'tcx> Stable<'tcx> for mir::interpret::Allocation {
     }
 }
 
+impl<'tcx> Stable<'tcx> for mir::interpret::AllocId {
+    type T = stable_mir::mir::alloc::AllocId;
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        tables.create_alloc_id(*self)
+    }
+}
+
 impl<'tcx> Stable<'tcx> for mir::interpret::GlobalAlloc<'tcx> {
     type T = GlobalAlloc;
 
diff --git a/compiler/stable_mir/src/lib.rs b/compiler/stable_mir/src/lib.rs
index d5c4881c89d..6c1b723a8da 100644
--- a/compiler/stable_mir/src/lib.rs
+++ b/compiler/stable_mir/src/lib.rs
@@ -291,6 +291,9 @@ pub trait Context {
 
     /// Retrieve global allocation for the given allocation ID.
     fn global_alloc(&self, id: AllocId) -> GlobalAlloc;
+
+    /// Retrieve the id for the virtual table.
+    fn vtable_allocation(&self, global_alloc: &GlobalAlloc) -> Option<AllocId>;
 }
 
 // A thread local variable that stores a pointer to the tables mapping between TyCtxt
diff --git a/compiler/stable_mir/src/mir/alloc.rs b/compiler/stable_mir/src/mir/alloc.rs
index ca24653585c..ef9053b78ec 100644
--- a/compiler/stable_mir/src/mir/alloc.rs
+++ b/compiler/stable_mir/src/mir/alloc.rs
@@ -9,6 +9,7 @@ pub enum GlobalAlloc {
     /// The alloc ID is used as a function pointer.
     Function(Instance),
     /// This alloc ID points to a symbolic (not-reified) vtable.
+    /// The `None` trait ref is used to represent auto traits.
     VTable(Ty, Option<Binder<ExistentialTraitRef>>),
     /// The alloc ID points to a "lazy" static variable that did not get computed (yet).
     /// This is also used to break the cycle in recursive statics.
@@ -23,6 +24,12 @@ impl From<AllocId> for GlobalAlloc {
     }
 }
 
+impl GlobalAlloc {
+    pub fn vtable_allocation(&self) -> Option<AllocId> {
+        with(|cx| cx.vtable_allocation(self))
+    }
+}
+
 /// A unique identification number for each provenance
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub struct AllocId(usize);