Add support to get virtual table allocation

2 files changed, 10 insertions, 0 deletions

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);