Repository: rust-lang/rust
Pull request: rustc_public: add `vtable_entries()` to `TraitRef`

Resolves: https://github.com/rust-lang/project-stable-mir/issues/103

Changed files:
filename: compiler/rustc_public/src/compiler_interface.rs
patch: @@ -20,7 +20,8 @@ use crate::ty::{
     AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, CoroutineDef, Discr, FieldDef, FnDef,
     ForeignDef, ForeignItemKind, ForeignModule, ForeignModuleDef, GenericArgs, GenericPredicates,
     Generics, ImplDef, ImplTrait, IntrinsicDef, LineInfo, MirConst, PolyFnSig, RigidTy, Span,
-    TraitDecl, TraitDef, Ty, TyConst, TyConstId, TyKind, UintTy, VariantDef, VariantIdx,
+    TraitDecl, TraitDef, TraitRef, Ty, TyConst, TyConstId, TyKind, UintTy, VariantDef, VariantIdx,
+    VtblEntry,
 };
 use crate::unstable::{RustcInternal, Stable, new_item_kind};
 use crate::{
@@ -838,6 +839,25 @@ impl<'tcx> CompilerInterface<'tcx> {
         let did = tables[def_id];
         cx.associated_items(did).iter().map(|assoc| assoc.stable(&mut *tables, cx)).collect()
     }
+
+    /// Get all vtable entries of a trait.
+    pub(crate) fn vtable_entries(&self, trait_ref: &TraitRef) -> Vec<VtblEntry> {
+        let mut tables = self.tables.borrow_mut();
+        let cx = &*self.cx.borrow();
+        cx.vtable_entries(trait_ref.internal(&mut *tables, cx.tcx))
+            .iter()
+            .map(|v| v.stable(&mut *tables, cx))
+            .collect()
+    }
+
+    /// Returns the vtable entry at the given index.
+    ///
+    /// Returns `None` if the index is out of bounds.
+    pub(crate) fn vtable_entry(&self, trait_ref: &TraitRef, idx: usize) -> Option<VtblEntry> {
+        let mut tables = self.tables.borrow_mut();
+        let cx = &*self.cx.borrow();
+        cx.vtable_entry(trait_ref.internal(&mut *tables, cx.tcx), idx).stable(&mut *tables, cx)
+    }
 }
 
 // A thread local variable that stores a pointer to [`CompilerInterface`].

filename: compiler/rustc_public/src/ty.rs
patch: @@ -9,7 +9,7 @@ use super::{DefId, Error, Symbol, with};
 use crate::abi::{FnAbi, Layout};
 use crate::crate_def::{CrateDef, CrateDefType};
 use crate::mir::alloc::{AllocId, read_target_int, read_target_uint};
-use crate::mir::mono::StaticDef;
+use crate::mir::mono::{Instance, StaticDef};
 use crate::target::MachineInfo;
 use crate::{AssocItems, Filename, IndexedVal, Opaque, ThreadLocalIndex};
 
@@ -1440,6 +1440,18 @@ impl TraitRef {
         };
         self_ty
     }
+
+    /// Retrieve all vtable entries.
+    pub fn vtable_entries(&self) -> Vec<VtblEntry> {
+        with(|cx| cx.vtable_entries(self))
+    }
+
+    /// Returns the vtable entry at the given index.
+    ///
+    /// Returns `None` if the index is out of bounds.
+    pub fn vtable_entry(&self, idx: usize) -> Option<VtblEntry> {
+        with(|cx| cx.vtable_entry(self, idx))
+    }
 }
 
 #[derive(Clone, Debug, Eq, PartialEq, Serialize)]
@@ -1656,3 +1668,19 @@ impl AssocItem {
         matches!(self.kind, AssocKind::Type { data: AssocTypeData::Rpitit(_) })
     }
 }
+
+#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
+pub enum VtblEntry {
+    /// destructor of this type (used in vtable header)
+    MetadataDropInPlace,
+    /// layout size of this type (used in vtable header)
+    MetadataSize,
+    /// layout align of this type (used in vtable header)
+    MetadataAlign,
+    /// non-dispatchable associated function that is excluded from trait object
+    Vacant,
+    /// dispatchable associated function
+    Method(Instance),
+    /// pointer to a separate supertrait vtable, can be used by trait upcasting coercion
+    TraitVPtr(TraitRef),
+}

filename: compiler/rustc_public/src/unstable/convert/stable/ty.rs
patch: @@ -1139,3 +1139,25 @@ impl<'tcx> Stable<'tcx> for rustc_middle::ty::util::Discr<'tcx> {
         crate::ty::Discr { val: self.val, ty: self.ty.stable(tables, cx) }
     }
 }
+
+impl<'tcx> Stable<'tcx> for rustc_middle::ty::VtblEntry<'tcx> {
+    type T = crate::ty::VtblEntry;
+
+    fn stable<'cx>(
+        &self,
+        tables: &mut Tables<'cx, BridgeTys>,
+        cx: &CompilerCtxt<'cx, BridgeTys>,
+    ) -> Self::T {
+        use crate::ty::VtblEntry;
+        match self {
+            ty::VtblEntry::MetadataDropInPlace => VtblEntry::MetadataDropInPlace,
+            ty::VtblEntry::MetadataSize => VtblEntry::MetadataSize,
+            ty::VtblEntry::MetadataAlign => VtblEntry::MetadataAlign,
+            ty::VtblEntry::Vacant => VtblEntry::Vacant,
+            ty::VtblEntry::Method(instance) => VtblEntry::Method(instance.stable(tables, cx)),
+            ty::VtblEntry::TraitVPtr(trait_ref) => {
+                VtblEntry::TraitVPtr(trait_ref.stable(tables, cx))
+            }
+        }
+    }
+}

filename: compiler/rustc_public_bridge/src/context/impls.rs
patch: @@ -18,7 +18,7 @@ use rustc_middle::ty::{
     AdtDef, AdtKind, AssocItem, Binder, ClosureKind, CoroutineArgsExt, EarlyBinder,
     ExistentialTraitRef, FnSig, GenericArgsRef, Instance, InstanceKind, IntrinsicDef, List,
     PolyFnSig, ScalarInt, TraitDef, TraitRef, Ty, TyCtxt, TyKind, TypeVisitableExt, UintTy,
-    ValTree, VariantDef,
+    ValTree, VariantDef, VtblEntry,
 };
 use rustc_middle::{mir, ty};
 use rustc_session::cstore::ForeignModule;
@@ -757,4 +757,16 @@ impl<'tcx, B: Bridge> CompilerCtxt<'tcx, B> {
         };
         assoc_items
     }
+
+    /// Get all vtable entries of a trait.
+    pub fn vtable_entries(&self, trait_ref: TraitRef<'tcx>) -> Vec<VtblEntry<'tcx>> {
+        self.tcx.vtable_entries(trait_ref).to_vec()
+    }
+
+    /// Returns the vtable entry at the given index.
+    ///
+    /// Returns `None` if the index is out of bounds.
+    pub fn vtable_entry(&self, trait_ref: TraitRef<'tcx>, idx: usize) -> Option<VtblEntry<'tcx>> {
+        self.vtable_entries(trait_ref).get(idx).copied()
+    }
 }

filename: tests/ui-fulldeps/rustc_public/check_vtable.rs
patch: @@ -0,0 +1,149 @@
+//@ run-pass
+// Test that users are able to use rustc_public to retrieve vtable info.
+
+//@ ignore-stage1
+//@ ignore-cross-compile
+//@ ignore-remote
+
+#![feature(rustc_private)]
+
+extern crate rustc_middle;
+extern crate rustc_driver;
+extern crate rustc_interface;
+#[macro_use]
+extern crate rustc_public;
+
+use rustc_public::ty::VtblEntry;
+use rustc_public::CrateDef;
+use std::io::Write;
+use std::ops::ControlFlow;
+
+const CRATE_NAME: &str = "vtable_test";
+
+/// This function uses the rustc_public APIs to test the `vtable_entries()`.
+fn test_vtable_entries() -> ControlFlow<()> {
+    let local_crate = rustc_public::local_crate();
+    let local_impls = local_crate.trait_impls();
+    let child_impl = local_impls
+        .iter()
+        .find(|i| i.trimmed_name() == "<Concrete as Child>")
+        .expect("Could not find <Concrete as Child>");
+
+    let child_trait_ref = child_impl.trait_impl().value;
+    let entries = child_trait_ref.vtable_entries();
+    match &entries[..] {
+        [
+            VtblEntry::MetadataDropInPlace,
+            VtblEntry::MetadataSize,
+            VtblEntry::MetadataAlign,
+            VtblEntry::Method(primary),
+            VtblEntry::Method(secondary),
+            VtblEntry::TraitVPtr(secondary_vptr),
+            VtblEntry::Method(child),
+        ] => {
+            assert!(
+                primary.name().contains("primary"),
+                "Expected primary method at index 3"
+            );
+            assert!(
+                secondary.name().contains("secondary"),
+                "Expected secondary method at index 4"
+            );
+            let vptr_str = secondary_vptr.def_id.name();
+            assert!(
+                vptr_str.contains("Secondary"),
+                "Expected Secondary VPtr at index 5"
+            );
+            assert!(
+                child.name().contains("child"),
+                "Expected child method at index 6"
+            );
+        }
+        _ => panic!(
+            "Unexpected vtable layout for <Concrete as Child>. Found: {:#?}",
+            entries
+        ),
+    }
+    let vacant_impl = local_impls
+        .iter()
+        .find(|i| i.trimmed_name() == "<Concrete as WithVacant>")
+        .expect("Could not find <Concrete as WithVacant>");
+    let vacant_trait_ref = vacant_impl.trait_impl().value;
+    let vacant_entries = vacant_trait_ref.vtable_entries();
+    match &vacant_entries[..] {
+        [
+            VtblEntry::MetadataDropInPlace,
+            VtblEntry::MetadataSize,
+            VtblEntry::MetadataAlign,
+            VtblEntry::Method(valid),
+        ] => {
+            assert!(valid.name().contains("valid"), "Expected valid method");
+        }
+        _ => panic!(
+            "Unexpected vtable layout for <Concrete as WithVacant>. Found: {:#?}",
+            vacant_entries
+        ),
+    }
+    ControlFlow::Continue(())
+}
+
+fn main() {
+    let path = "vtable_input.rs";
+    generate_input(&path).unwrap();
+    let args = &[
+        "rustc".to_string(),
+        "--crate-type=lib".to_string(),
+        "--crate-name".to_string(),
+        CRATE_NAME.to_string(),
+        path.to_string(),
+    ];
+    run!(args, test_vtable_entries).unwrap();
+}
+
+fn generate_input(path: &str) -> std::io::Result<()> {
+    let mut file = std::fs::File::create(path)?;
+    write!(
+        file,
+        r#"
+        pub struct Concrete;
+
+        pub trait Primary {{
+            fn primary(&self);
+        }}
+
+        pub trait Secondary {{
+            fn secondary(&self);
+        }}
+
+        pub trait Child: Primary + Secondary {{
+            fn child(&self);
+        }}
+
+        impl Primary for Concrete {{
+            fn primary(&self) {{}}
+        }}
+
+        impl Secondary for Concrete {{
+            fn secondary(&self) {{}}
+        }}
+
+        impl Child for Concrete {{
+            fn child(&self) {{}}
+        }}
+
+        pub trait WithVacant {{
+            fn valid(&self);
+
+            fn excluded<T>(&self, meow: T) where Self: Sized;
+        }}
+
+        impl WithVacant for Concrete {{
+            fn valid(&self) {{}}
+            fn excluded<T>(&self, meow: T) {{}}
+        }}
+
+        fn main() {{}}
+    "#
+    )?;
+    Ok(())
+}