Rollup merge of rust-lang#152003 - 9SonSteroids:trait_info_of, r=oli-obk

Reflection TypeId::trait_info_of

*[View all comments](https://triagebot.infra.rust-lang.org/gh-comments/rust-lang/rust/pull/152003)*

This is for rust-lang#146922.

As rust-lang#151236 was requested to be remade by someone I implemented the functionality as `TypeId::trait_info_of` which additionally allows getting the vtable pointer to build `dyn` Objects in recursive reflection.

It allows checking if a TypeId implements a trait. Since this is my first PR feel free to tell me if there are any formal issues.

Author: jhpratt

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;
@@ -24,7 +24,7 @@ use crate::interpret::{
     self, AllocId, AllocInit, AllocRange, ConstAllocation, CtfeProvenance, FnArg, Frame,
     GlobalAlloc, ImmTy, InterpCx, InterpResult, OpTy, PlaceTy, Pointer, RangeSet, Scalar,
     compile_time_machine, err_inval, interp_ok, throw_exhaust, throw_inval, throw_ub,
-    throw_ub_custom, throw_unsup, throw_unsup_format,
+    throw_ub_custom, throw_unsup, throw_unsup_format, type_implements_dyn_trait,
 };
 
 /// When hitting this many interpreted terminators we emit a deny by default lint
@@ -598,6 +598,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/intrinsics.rs

@@ -4,19 +4,16 @@
 
 mod simd;
 
-use rustc_abi::{FIRST_VARIANT, FieldIdx, HasDataLayout, Size, VariantIdx};
+use rustc_abi::{FieldIdx, HasDataLayout, Size, VariantIdx};
 use rustc_apfloat::ieee::{Double, Half, Quad, Single};
 use rustc_data_structures::assert_matches;
 use rustc_errors::msg;
-use rustc_hir::def_id::CRATE_DEF_ID;
-use rustc_infer::infer::TyCtxtInferExt;
 use rustc_middle::mir::interpret::{CTFE_ALLOC_SALT, read_target_uint, write_target_uint};
 use rustc_middle::mir::{self, BinOp, ConstValue, NonDivergingIntrinsic};
 use rustc_middle::ty::layout::TyAndLayout;
-use rustc_middle::ty::{FloatTy, PolyExistentialPredicate, Ty, TyCtxt, TypeVisitableExt};
+use rustc_middle::ty::{FloatTy, Ty, TyCtxt, TypeVisitableExt};
 use rustc_middle::{bug, span_bug, ty};
 use rustc_span::{Symbol, sym};
-use rustc_trait_selection::traits::{Obligation, ObligationCause, ObligationCtxt};
 use tracing::trace;
 
 use super::memory::MemoryKind;
@@ -227,44 +224,6 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
 
                 self.write_scalar(Scalar::from_target_usize(offset, self), dest)?;
             }
-            sym::vtable_for => {
-                let tp_ty = instance.args.type_at(0);
-                let result_ty = instance.args.type_at(1);
-
-                ensure_monomorphic_enough(tcx, tp_ty)?;
-                ensure_monomorphic_enough(tcx, result_ty)?;
-                let ty::Dynamic(preds, _) = result_ty.kind() else {
-                    span_bug!(
-                        self.find_closest_untracked_caller_location(),
-                        "Invalid type provided to vtable_for::<T, U>. U must be dyn Trait, got {result_ty}."
-                    );
-                };
-
-                let (infcx, param_env) =
-                    self.tcx.infer_ctxt().build_with_typing_env(self.typing_env);
-
-                let ocx = ObligationCtxt::new(&infcx);
-                ocx.register_obligations(preds.iter().map(|pred: PolyExistentialPredicate<'_>| {
-                    let pred = pred.with_self_ty(tcx, tp_ty);
-                    // Lifetimes can only be 'static because of the bound on T
-                    let pred = ty::fold_regions(tcx, pred, |r, _| {
-                        if r == tcx.lifetimes.re_erased { tcx.lifetimes.re_static } else { r }
-                    });
-                    Obligation::new(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();
-
-                if regions_are_valid && type_impls_trait {
-                    let vtable_ptr = self.get_vtable_ptr(tp_ty, preds)?;
-                    // Writing a non-null pointer into an `Option<NonNull>` will automatically make it `Some`.
-                    self.write_pointer(vtable_ptr, dest)?;
-                } else {
-                    // Write `None`
-                    self.write_discriminant(FIRST_VARIANT, dest)?;
-                }
-            }
             sym::variant_count => {
                 let tp_ty = instance.args.type_at(0);
                 let ty = match tp_ty.kind() {

compiler/rustc_const_eval/src/interpret/mod.rs

@@ -38,6 +38,6 @@ use self::place::{MemPlace, Place};
 pub use self::projection::{OffsetMode, Projectable};
 pub use self::stack::{Frame, FrameInfo, LocalState, ReturnContinuation};
 pub use self::util::EnteredTraceSpan;
-pub(crate) use self::util::create_static_alloc;
+pub(crate) use self::util::{create_static_alloc, type_implements_dyn_trait};
 pub use self::validity::{CtfeValidationMode, RangeSet, RefTracking};
 pub use self::visitor::ValueVisitor;

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,12 +213,12 @@ 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
         | sym::va_copy
         | sym::variant_count
-        | sym::vtable_for
         | sym::wrapping_add
         | sym::wrapping_mul
         | sym::wrapping_sub
@@ -323,6 +323,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,
@@ -675,20 +694,6 @@ pub(crate) fn check_intrinsic_type(
             (0, 0, vec![Ty::new_imm_ptr(tcx, tcx.types.unit)], tcx.types.usize)
         }
 
-        sym::vtable_for => {
-            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(&[param(1).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);
-
-            (2, 0, vec![], ret_ty)
-        }
-
         // This type check is not particularly useful, but the `where` bounds
         // on the definition in `core` do the heavy lifting for checking it.
         sym::aggregate_raw_ptr => (3, 0, vec![param(1), param(2)], param(0)),

compiler/rustc_span/src/symbol.rs

@@ -2054,6 +2054,7 @@ symbols! {
         type_changing_struct_update,
         type_id,
         type_id_eq,
+        type_id_vtable,
         type_info,
         type_ir,
         type_ir_infer_ctxt_like,
@@ -2206,7 +2207,6 @@ symbols! {
         vsreg,
         vsx,
         vtable_align,
-        vtable_for,
         vtable_size,
         warn,
         wasip2,

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];
 
@@ -948,7 +1012,8 @@ pub const fn try_as_dyn<
 >(
     t: &T,
 ) -> Option<&U> {
-    let vtable: Option<ptr::DynMetadata<U>> = const { intrinsics::vtable_for::<T, U>() };
+    let vtable: Option<ptr::DynMetadata<U>> =
+        const { TypeId::of::<T>().trait_info_of::<U>().as_ref().map(TraitImpl::get_vtable) };
     match vtable {
         Some(dyn_metadata) => {
             let pointer = ptr::from_raw_parts(t, dyn_metadata);
@@ -1001,7 +1066,8 @@ pub const fn try_as_dyn_mut<
 >(
     t: &mut T,
 ) -> Option<&mut U> {
-    let vtable: Option<ptr::DynMetadata<U>> = const { intrinsics::vtable_for::<T, U>() };
+    let vtable: Option<ptr::DynMetadata<U>> =
+        const { TypeId::of::<T>().trait_info_of::<U>().as_ref().map(TraitImpl::get_vtable) };
     match vtable {
         Some(dyn_metadata) => {
             let pointer = ptr::from_raw_parts_mut(t, dyn_metadata);

library/core/src/intrinsics/mod.rs

@@ -2751,18 +2751,6 @@ pub unsafe fn vtable_size(ptr: *const ()) -> usize;
 #[rustc_intrinsic]
 pub unsafe fn vtable_align(ptr: *const ()) -> usize;
 
-/// The intrinsic returns the `U` vtable for `T` if `T` can be coerced to the trait object type `U`.
-///
-/// # Compile-time failures
-/// Determining whether `T` can be coerced to the trait object type `U` requires trait resolution by the compiler.
-/// In some cases, that resolution can exceed the recursion limit,
-/// and compilation will fail instead of this function returning `None`.
-#[rustc_nounwind]
-#[unstable(feature = "core_intrinsics", issue = "none")]
-#[rustc_intrinsic]
-pub const fn vtable_for<T, U: ptr::Pointee<Metadata = ptr::DynMetadata<U>> + ?Sized>()
--> Option<ptr::DynMetadata<U>>;
-
 /// The size of a type in bytes.
 ///
 /// Note that, unlike most intrinsics, this is safe to call;
@@ -2864,6 +2852,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.