Add support to get virtual table allocation

Author: celinval

compiler/rustc_smir/src/rustc_internal/internal.rs

@@ -10,8 +10,8 @@ use rustc_span::Symbol;
 use stable_mir::mir::alloc::AllocId;
 use stable_mir::mir::mono::{Instance, MonoItem, StaticDef};
 use stable_mir::ty::{
-    Binder, BoundRegionKind, BoundTyKind, BoundVariableKind, ClosureKind, Const, GenericArgKind,
-    GenericArgs, Region, TraitRef, Ty,
+    AdtDef, Binder, BoundRegionKind, BoundTyKind, BoundVariableKind, ClosureKind, Const,
+    ExistentialTraitRef, GenericArgKind, GenericArgs, Region, TraitRef, Ty,
 };
 use stable_mir::{CrateItem, DefId};
 
@@ -153,6 +153,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>;
 
@@ -184,6 +195,17 @@ impl<'tcx> RustcInternal<'tcx> for ClosureKind {
     }
 }
 
+impl<'tcx> RustcInternal<'tcx> for AdtDef {
+    type T = rustc_ty::AdtDef<'tcx>;
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        let ty = tables.tcx.type_of(self.0.internal(&mut *tables)).instantiate_identity().kind();
+        let rustc_ty::TyKind::Adt(def, _) = ty else {
+            panic!("Expected an ADT definition, but found: {ty:?}")
+        };
+        *def
+    }
+}
+
 impl<'tcx, T> RustcInternal<'tcx> for &T
 where
     T: RustcInternal<'tcx>,
@@ -194,3 +216,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))
+    }
+}

compiler/rustc_smir/src/rustc_internal/mod.rs

@@ -117,7 +117,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)
     }
 

compiler/rustc_smir/src/rustc_smir/mod.rs

@@ -20,7 +20,11 @@ use rustc_target::abi::FieldIdx;
 use stable_mir::mir::alloc::GlobalAlloc;
 use stable_mir::mir::mono::{InstanceDef, StaticDef};
 use stable_mir::mir::{Body, CopyNonOverlapping, Statement, UserTypeProjection, VariantIdx};
-use stable_mir::ty::{AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, ConstId, ConstantKind, FloatTy, FnDef, GenericArgs, GenericParamDef, IntTy, LineInfo, Movability, RigidTy, Span, TyKind, UintTy, EarlyParamRegion};
+use stable_mir::ty::{
+    AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, ConstId, ConstantKind,
+    EarlyParamRegion, FloatTy, FnDef, GenericArgs, GenericParamDef, IntTy, LineInfo, Movability,
+    RigidTy, Span, TyKind, UintTy,
+};
 use stable_mir::{self, opaque, Context, CrateItem, Error, Filename, ItemKind};
 use std::cell::RefCell;
 use tracing::debug;
@@ -312,6 +316,18 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         let alloc_id = alloc.internal(&mut *tables);
         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))
+    }
 }
 
 pub(crate) struct TablesWrapper<'tcx>(pub(crate) RefCell<Tables<'tcx>>);
@@ -1536,6 +1552,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;
 

compiler/stable_mir/src/lib.rs

@@ -268,6 +268,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

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