Support getting TypeId's Trait and vtable

Author: jasper3108

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -2,7 +2,7 @@ use std::borrow::{Borrow, Cow};
 use std::fmt;
 use std::hash::Hash;
 
-use rustc_abi::{Align, Size};
+use rustc_abi::{Align, FIRST_VARIANT, Size};
 use rustc_ast::Mutability;
 use rustc_data_structures::fx::{FxHashMap, FxIndexMap, IndexEntry};
 use rustc_errors::msg;
@@ -21,6 +21,7 @@ use tracing::debug;
 
 use super::error::*;
 use crate::errors::{LongRunning, LongRunningWarn};
+use crate::interpret::util::type_implements_dyn_trait;
 use crate::interpret::{
     self, AllocId, AllocInit, AllocRange, ConstAllocation, CtfeProvenance, FnArg, Frame,
     GlobalAlloc, ImmTy, InterpCx, InterpResult, OpTy, PlaceTy, Pointer, RangeSet, Scalar,
@@ -601,6 +602,22 @@ impl<'tcx> interpret::Machine<'tcx> for CompileTimeMachine<'tcx> {
                 }
             }
 
+            sym::type_id_vtable => {
+                let tp_ty = ecx.read_type_id(&args[0])?;
+                let result_ty = ecx.read_type_id(&args[1])?;
+
+                let (implements_trait, preds) = type_implements_dyn_trait(ecx, tp_ty, result_ty)?;
+
+                if implements_trait {
+                    let vtable_ptr = ecx.get_vtable_ptr(tp_ty, preds)?;
+                    // Writing a non-null pointer into an `Option<NonNull>` will automatically make it `Some`.
+                    ecx.write_pointer(vtable_ptr, dest)?;
+                } else {
+                    // Write `None`
+                    ecx.write_discriminant(FIRST_VARIANT, dest)?;
+                }
+            }
+
             sym::type_of => {
                 let ty = ecx.read_type_id(&args[0])?;
                 ecx.write_type_info(ty, dest)?;

compiler/rustc_const_eval/src/interpret/mod.rs

@@ -15,7 +15,7 @@ mod projection;
 mod stack;
 mod step;
 mod traits;
-mod util;
+pub(crate) mod util;
 mod validity;
 mod visitor;
 

compiler/rustc_const_eval/src/interpret/util.rs

@@ -1,12 +1,51 @@
-use rustc_hir::def_id::LocalDefId;
-use rustc_middle::mir;
+use rustc_hir::def_id::{CRATE_DEF_ID, LocalDefId};
+use rustc_infer::infer::TyCtxtInferExt;
+use rustc_infer::traits::{Obligation, ObligationCause};
 use rustc_middle::mir::interpret::{AllocInit, Allocation, GlobalAlloc, InterpResult, Pointer};
 use rustc_middle::ty::layout::TyAndLayout;
-use rustc_middle::ty::{TyCtxt, TypeVisitable, TypeVisitableExt};
+use rustc_middle::ty::{PolyExistentialPredicate, Ty, TyCtxt, TypeVisitable, TypeVisitableExt};
+use rustc_middle::{mir, span_bug, ty};
+use rustc_trait_selection::traits::ObligationCtxt;
 use tracing::debug;
 
 use super::{InterpCx, MPlaceTy, MemoryKind, interp_ok, throw_inval};
 use crate::const_eval::{CompileTimeInterpCx, CompileTimeMachine, InterpretationResult};
+use crate::interpret::Machine;
+
+/// Checks if a type implements predicates.
+/// Calls `ensure_monomorphic_enough` on `ty` and `trait_ty` for you.
+pub(crate) fn type_implements_dyn_trait<'tcx, M: Machine<'tcx>>(
+    ecx: &mut InterpCx<'tcx, M>,
+    ty: Ty<'tcx>,
+    trait_ty: Ty<'tcx>,
+) -> InterpResult<'tcx, (bool, &'tcx ty::List<ty::PolyExistentialPredicate<'tcx>>)> {
+    ensure_monomorphic_enough(ecx.tcx.tcx, ty)?;
+    ensure_monomorphic_enough(ecx.tcx.tcx, trait_ty)?;
+
+    let ty::Dynamic(preds, _) = trait_ty.kind() else {
+        span_bug!(
+            ecx.find_closest_untracked_caller_location(),
+            "Invalid type provided to type_implements_predicates. U must be dyn Trait, got {trait_ty}."
+        );
+    };
+
+    let (infcx, param_env) = ecx.tcx.infer_ctxt().build_with_typing_env(ecx.typing_env);
+
+    let ocx = ObligationCtxt::new(&infcx);
+    ocx.register_obligations(preds.iter().map(|pred: PolyExistentialPredicate<'_>| {
+        let pred = pred.with_self_ty(ecx.tcx.tcx, ty);
+        // Lifetimes can only be 'static because of the bound on T
+        let pred = rustc_middle::ty::fold_regions(ecx.tcx.tcx, pred, |r, _| {
+            if r == ecx.tcx.tcx.lifetimes.re_erased { ecx.tcx.tcx.lifetimes.re_static } else { r }
+        });
+        Obligation::new(ecx.tcx.tcx, ObligationCause::dummy(), param_env, pred)
+    }));
+    let type_impls_trait = ocx.evaluate_obligations_error_on_ambiguity().is_empty();
+    // Since `assumed_wf_tys=[]` the choice of LocalDefId is irrelevant, so using the "default"
+    let regions_are_valid = ocx.resolve_regions(CRATE_DEF_ID, param_env, []).is_empty();
+
+    interp_ok((regions_are_valid && type_impls_trait, preds))
+}
 
 /// Checks whether a type contains generic parameters which must be instantiated.
 ///

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -213,6 +213,7 @@ fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: LocalDefId) -> hi
         | sym::truncf128
         | sym::type_id
         | sym::type_id_eq
+        | sym::type_id_vtable
         | sym::type_name
         | sym::type_of
         | sym::ub_checks
@@ -323,6 +324,25 @@ pub(crate) fn check_intrinsic_type(
             let type_id = tcx.type_of(tcx.lang_items().type_id().unwrap()).no_bound_vars().unwrap();
             (0, 0, vec![type_id, type_id], tcx.types.bool)
         }
+        sym::type_id_vtable => {
+            let dyn_metadata = tcx.require_lang_item(LangItem::DynMetadata, span);
+            let dyn_metadata_adt_ref = tcx.adt_def(dyn_metadata);
+            let dyn_metadata_args =
+                tcx.mk_args(&[Ty::new_ptr(tcx, tcx.types.unit, ty::Mutability::Not).into()]);
+            let dyn_ty = Ty::new_adt(tcx, dyn_metadata_adt_ref, dyn_metadata_args);
+
+            let option_did = tcx.require_lang_item(LangItem::Option, span);
+            let option_adt_ref = tcx.adt_def(option_did);
+            let option_args = tcx.mk_args(&[dyn_ty.into()]);
+            let ret_ty = Ty::new_adt(tcx, option_adt_ref, option_args);
+
+            (
+                0,
+                0,
+                vec![tcx.type_of(tcx.lang_items().type_id().unwrap()).no_bound_vars().unwrap(); 2],
+                ret_ty,
+            )
+        }
         sym::type_of => (
             0,
             0,

compiler/rustc_span/src/symbol.rs

@@ -2358,6 +2358,7 @@ symbols! {
         type_changing_struct_update,
         type_id,
         type_id_eq,
+        type_id_vtable,
         type_info,
         type_ir,
         type_ir_infer_ctxt_like,

library/core/src/any.rs

@@ -86,7 +86,10 @@
 
 #![stable(feature = "rust1", since = "1.0.0")]
 
-use crate::{fmt, hash, intrinsics, ptr};
+use crate::intrinsics::{self, type_id_vtable};
+use crate::mem::transmute;
+use crate::mem::type_info::{TraitImpl, TypeKind};
+use crate::{fmt, hash, ptr};
 
 ///////////////////////////////////////////////////////////////////////////////
 // Any trait
@@ -788,6 +791,67 @@ impl TypeId {
         const { intrinsics::type_id::<T>() }
     }
 
+    /// Checks if the [TypeId] implements the trait. If it does it returns [TraitImpl] which can be used to build a fat pointer.
+    /// It can only be called at compile time. `self` must be the [TypeId] of a sized type or None will be returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(type_info)]
+    /// use std::any::{TypeId};
+    ///
+    /// pub trait Blah {}
+    /// impl Blah for u8 {}
+    ///
+    /// assert!(const { TypeId::of::<u8>().trait_info_of::<dyn Blah>() }.is_some());
+    /// assert!(const { TypeId::of::<u16>().trait_info_of::<dyn Blah>() }.is_none());
+    /// ```
+    #[unstable(feature = "type_info", issue = "146922")]
+    #[rustc_const_unstable(feature = "type_info", issue = "146922")]
+    pub const fn trait_info_of<
+        T: ptr::Pointee<Metadata = ptr::DynMetadata<T>> + ?Sized + 'static,
+    >(
+        self,
+    ) -> Option<TraitImpl<T>> {
+        // SAFETY: The vtable was obtained for `T`, so it is guaranteed to be `DynMetadata<T>`.
+        // The intrinsic can't infer this because it is designed to work with arbitrary TypeIds.
+        unsafe { transmute(self.trait_info_of_trait_type_id(const { TypeId::of::<T>() })) }
+    }
+
+    /// Checks if the [TypeId] implements the trait of `trait_represented_by_type_id`. If it does it returns [TraitImpl] which can be used to build a fat pointer.
+    /// It can only be called at compile time. `self` must be the [TypeId] of a sized type or None will be returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(type_info)]
+    /// use std::any::{TypeId};
+    ///
+    /// pub trait Blah {}
+    /// impl Blah for u8 {}
+    ///
+    /// assert!(const { TypeId::of::<u8>().trait_info_of_trait_type_id(TypeId::of::<dyn Blah>()) }.is_some());
+    /// assert!(const { TypeId::of::<u16>().trait_info_of_trait_type_id(TypeId::of::<dyn Blah>()) }.is_none());
+    /// ```
+    #[unstable(feature = "type_info", issue = "146922")]
+    #[rustc_const_unstable(feature = "type_info", issue = "146922")]
+    pub const fn trait_info_of_trait_type_id(
+        self,
+        trait_represented_by_type_id: TypeId,
+    ) -> Option<TraitImpl<*const ()>> {
+        if self.info().size.is_none() {
+            return None;
+        }
+
+        if matches!(trait_represented_by_type_id.info().kind, TypeKind::DynTrait(_))
+            && let Some(vtable) = type_id_vtable(self, trait_represented_by_type_id)
+        {
+            Some(TraitImpl { vtable })
+        } else {
+            None
+        }
+    }
+
     fn as_u128(self) -> u128 {
         let mut bytes = [0; 16];
 

library/core/src/intrinsics/mod.rs

@@ -2864,6 +2864,20 @@ pub const unsafe fn size_of_val<T: ?Sized>(ptr: *const T) -> usize;
 #[rustc_intrinsic_const_stable_indirect]
 pub const unsafe fn align_of_val<T: ?Sized>(ptr: *const T) -> usize;
 
+#[rustc_intrinsic]
+#[unstable(feature = "core_intrinsics", issue = "none")]
+/// Check if a type represented by a `TypeId` implements a trait represented by a `TypeId`.
+/// It can only be called at compile time, the backends do
+/// not implement it. If it implements the trait the dyn metadata gets returned for vtable access.
+pub const fn type_id_vtable(
+    _id: crate::any::TypeId,
+    _trait: crate::any::TypeId,
+) -> Option<ptr::DynMetadata<*const ()>> {
+    panic!(
+        "`TypeId::trait_info_of` and `trait_info_of_trait_type_id` can only be called at compile-time"
+    )
+}
+
 /// Compute the type information of a concrete type.
 /// It can only be called at compile time, the backends do
 /// not implement it.

library/core/src/mem/type_info.rs

@@ -3,6 +3,8 @@
 
 use crate::any::TypeId;
 use crate::intrinsics::{type_id, type_of};
+use crate::marker::PointeeSized;
+use crate::ptr::DynMetadata;
 
 /// Compile-time type information.
 #[derive(Debug)]
@@ -16,6 +18,19 @@ pub struct Type {
     pub size: Option<usize>,
 }
 
+/// Info of a trait implementation, you can retrieve the vtable with [Self::get_vtable]
+#[derive(Debug, PartialEq, Eq)]
+pub struct TraitImpl<T: PointeeSized> {
+    pub(crate) vtable: DynMetadata<T>,
+}
+
+impl<T: PointeeSized> TraitImpl<T> {
+    /// Gets the raw vtable for type reflection mapping
+    pub fn get_vtable(&self) -> DynMetadata<T> {
+        self.vtable
+    }
+}
+
 impl TypeId {
     /// Compute the type information of a concrete type.
     /// It can only be called at compile time.

library/coretests/tests/mem.rs

@@ -1,4 +1,5 @@
 mod fn_ptr;
+mod trait_info_of;
 mod type_info;
 
 use core::mem::*;

library/coretests/tests/mem/trait_info_of.rs

@@ -0,0 +1,70 @@
+use std::any::TypeId;
+use std::ptr::DynMetadata;
+
+struct Garlic(i32);
+trait Blah {
+    fn get_truth(&self) -> i32;
+}
+impl Blah for Garlic {
+    fn get_truth(&self) -> i32 {
+        self.0 * 21
+    }
+}
+
+#[test]
+fn test_implements_trait() {
+    const {
+        assert!(TypeId::of::<Garlic>().trait_info_of::<dyn Blah>().is_some());
+        assert!(TypeId::of::<Garlic>().trait_info_of::<dyn Blah + Send>().is_some());
+        assert!(TypeId::of::<*const Box<Garlic>>().trait_info_of::<dyn Sync>().is_none());
+        assert!(TypeId::of::<u8>().trait_info_of_trait_type_id(TypeId::of::<dyn Blah>()).is_none());
+    }
+}
+
+#[test]
+fn test_dyn_creation() {
+    let garlic = Garlic(2);
+    unsafe {
+        assert_eq!(
+            std::ptr::from_raw_parts::<dyn Blah>(
+                &raw const garlic,
+                const { TypeId::of::<Garlic>().trait_info_of::<dyn Blah>() }.unwrap().get_vtable()
+            )
+            .as_ref()
+            .unwrap()
+            .get_truth(),
+            42
+        );
+    }
+
+    assert_eq!(
+        const {
+            TypeId::of::<Garlic>()
+            .trait_info_of_trait_type_id(TypeId::of::<dyn Blah>())
+            .unwrap()
+        }.get_vtable(),
+        unsafe {
+            crate::mem::transmute::<_, DynMetadata<*const ()>>(
+                const {
+                    TypeId::of::<Garlic>().trait_info_of::<dyn Blah>()
+                }.unwrap().get_vtable(),
+            )
+        }
+    );
+}
+
+#[test]
+fn test_incorrect_use() {
+    assert_eq!(
+        const { TypeId::of::<i32>().trait_info_of_trait_type_id(TypeId::of::<u32>()) },
+        None
+    );
+}
+
+trait DstTrait {}
+impl DstTrait for [i32] {}
+
+#[test]
+fn dst_ice() {
+    assert!(const { TypeId::of::<[i32]>().trait_info_of::<dyn DstTrait>() }.is_none());
+}