Only compute vtable information during codegen

4 files changed, 79 insertions, 110 deletions

diff --git a/compiler/rustc_trait_selection/src/traits/mod.rs b/compiler/rustc_trait_selection/src/traits/mod.rs
index eae2f9d1792..af6bfdae440 100644
--- a/compiler/rustc_trait_selection/src/traits/mod.rs
+++ b/compiler/rustc_trait_selection/src/traits/mod.rs
@@ -65,7 +65,7 @@ pub use self::structural_match::search_for_structural_match_violation;
 pub use self::structural_normalize::StructurallyNormalizeExt;
 pub use self::util::elaborate;
 pub use self::util::{expand_trait_aliases, TraitAliasExpander, TraitAliasExpansionInfo};
-pub use self::util::{get_vtable_index_of_object_method, impl_item_is_final, upcast_choices};
+pub use self::util::{impl_item_is_final, upcast_choices};
 pub use self::util::{supertraits, transitive_bounds, transitive_bounds_that_define_assoc_item};
 pub use self::util::{with_replaced_escaping_bound_vars, BoundVarReplacer, PlaceholderReplacer};
 
diff --git a/compiler/rustc_trait_selection/src/traits/select/confirmation.rs b/compiler/rustc_trait_selection/src/traits/select/confirmation.rs
index 749081006f3..6cf87a9ac78 100644
--- a/compiler/rustc_trait_selection/src/traits/select/confirmation.rs
+++ b/compiler/rustc_trait_selection/src/traits/select/confirmation.rs
@@ -22,10 +22,6 @@ use rustc_span::def_id::DefId;
 
 use crate::traits::normalize::{normalize_with_depth, normalize_with_depth_to};
 use crate::traits::util::{self, closure_trait_ref_and_return_type};
-use crate::traits::vtable::{
-    count_own_vtable_entries, prepare_vtable_segments, vtable_trait_first_method_offset,
-    VtblSegment,
-};
 use crate::traits::{
     ImplDerivedCause, ImplSource, ImplSourceUserDefinedData, Normalized, Obligation,
     ObligationCause, PolyTraitObligation, PredicateObligation, Selection, SelectionError,
@@ -689,13 +685,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
 
         debug!(?nested, "object nested obligations");
 
-        let vtable_base = vtable_trait_first_method_offset(
-            tcx,
-            unnormalized_upcast_trait_ref,
-            ty::Binder::dummy(object_trait_ref),
-        );
-
-        Ok(ImplSource::Builtin(BuiltinImplSource::Object { vtable_base: vtable_base }, nested))
+        Ok(ImplSource::Builtin(BuiltinImplSource::Object(index), nested))
     }
 
     fn confirm_fn_pointer_candidate(
@@ -1125,36 +1115,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
             )?
             .expect("did not expect ambiguity during confirmation");
 
-        let vtable_segment_callback = {
-            let mut vptr_offset = 0;
-            move |segment| {
-                match segment {
-                    VtblSegment::MetadataDSA => {
-                        vptr_offset += TyCtxt::COMMON_VTABLE_ENTRIES.len();
-                    }
-                    VtblSegment::TraitOwnEntries { trait_ref, emit_vptr } => {
-                        vptr_offset += count_own_vtable_entries(tcx, trait_ref);
-                        if trait_ref == unnormalized_upcast_principal {
-                            if emit_vptr {
-                                return ControlFlow::Break(Some(vptr_offset));
-                            } else {
-                                return ControlFlow::Break(None);
-                            }
-                        }
-
-                        if emit_vptr {
-                            vptr_offset += 1;
-                        }
-                    }
-                }
-                ControlFlow::Continue(())
-            }
-        };
-
-        let vtable_vptr_slot =
-            prepare_vtable_segments(tcx, source_principal, vtable_segment_callback).unwrap();
-
-        Ok(ImplSource::Builtin(BuiltinImplSource::TraitUpcasting { vtable_vptr_slot }, nested))
+        Ok(ImplSource::Builtin(BuiltinImplSource::TraitUpcasting, nested))
     }
 
     fn confirm_builtin_unsize_candidate(
diff --git a/compiler/rustc_trait_selection/src/traits/util.rs b/compiler/rustc_trait_selection/src/traits/util.rs
index ce7245d93a4..c3fe816028e 100644
--- a/compiler/rustc_trait_selection/src/traits/util.rs
+++ b/compiler/rustc_trait_selection/src/traits/util.rs
@@ -208,23 +208,6 @@ pub fn upcast_choices<'tcx>(
     supertraits(tcx, source_trait_ref).filter(|r| r.def_id() == target_trait_def_id).collect()
 }
 
-/// Given an upcast trait object described by `object`, returns the
-/// index of the method `method_def_id` (which should be part of
-/// `object.upcast_trait_ref`) within the vtable for `object`.
-pub fn get_vtable_index_of_object_method<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    vtable_base: usize,
-    method_def_id: DefId,
-) -> Option<usize> {
-    // Count number of methods preceding the one we are selecting and
-    // add them to the total offset.
-    tcx.own_existential_vtable_entries(tcx.parent(method_def_id))
-        .iter()
-        .copied()
-        .position(|def_id| def_id == method_def_id)
-        .map(|index| vtable_base + index)
-}
-
 pub fn closure_trait_ref_and_return_type<'tcx>(
     tcx: TyCtxt<'tcx>,
     fn_trait_def_id: DefId,
diff --git a/compiler/rustc_trait_selection/src/traits/vtable.rs b/compiler/rustc_trait_selection/src/traits/vtable.rs
index 9bd4a9aab0a..017b0a45d1f 100644
--- a/compiler/rustc_trait_selection/src/traits/vtable.rs
+++ b/compiler/rustc_trait_selection/src/traits/vtable.rs
@@ -1,15 +1,11 @@
 use crate::errors::DumpVTableEntries;
 use crate::traits::{impossible_predicates, is_vtable_safe_method};
 use rustc_hir::def_id::DefId;
-use rustc_hir::lang_items::LangItem;
 use rustc_infer::traits::util::PredicateSet;
-use rustc_infer::traits::ImplSource;
 use rustc_middle::bug;
 use rustc_middle::query::Providers;
-use rustc_middle::traits::BuiltinImplSource;
-use rustc_middle::ty::visit::TypeVisitableExt;
-use rustc_middle::ty::GenericArgs;
 use rustc_middle::ty::{self, GenericParamDefKind, Ty, TyCtxt, Upcast, VtblEntry};
+use rustc_middle::ty::{GenericArgs, TypeVisitableExt};
 use rustc_span::{sym, Span};
 use smallvec::{smallvec, SmallVec};
 
@@ -320,30 +316,42 @@ fn vtable_entries<'tcx>(
     tcx.arena.alloc_from_iter(entries)
 }
 
-/// Find slot base for trait methods within vtable entries of another trait
-pub(super) fn vtable_trait_first_method_offset<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    trait_to_be_found: ty::PolyTraitRef<'tcx>,
-    trait_owning_vtable: ty::PolyTraitRef<'tcx>,
-) -> usize {
-    // #90177
-    let trait_to_be_found_erased = tcx.erase_regions(trait_to_be_found);
+// Given a `dyn Subtrait: Supertrait` trait ref, find corresponding first slot
+// for `Supertrait`'s methods in the vtable of `Subtrait`.
+pub(crate) fn first_method_vtable_slot<'tcx>(tcx: TyCtxt<'tcx>, key: ty::TraitRef<'tcx>) -> usize {
+    debug_assert!(!key.has_non_region_infer() && !key.has_non_region_param());
 
-    let vtable_segment_callback = {
-        let mut vtable_base = 0;
+    let ty::Dynamic(source, _, _) = *key.self_ty().kind() else {
+        bug!();
+    };
+    let source_principal = tcx
+        .normalize_erasing_regions(ty::ParamEnv::reveal_all(), source.principal().unwrap())
+        .with_self_ty(tcx, tcx.types.trait_object_dummy_self);
+
+    let target_principal = tcx
+        .normalize_erasing_regions(ty::ParamEnv::reveal_all(), key)
+        // We don't care about the self type, since it will always be the same thing.
+        .with_self_ty(tcx, tcx.types.trait_object_dummy_self);
 
+    let vtable_segment_callback = {
+        let mut vptr_offset = 0;
         move |segment| {
             match segment {
                 VtblSegment::MetadataDSA => {
-                    vtable_base += TyCtxt::COMMON_VTABLE_ENTRIES.len();
+                    vptr_offset += TyCtxt::COMMON_VTABLE_ENTRIES.len();
                 }
                 VtblSegment::TraitOwnEntries { trait_ref, emit_vptr } => {
-                    if tcx.erase_regions(trait_ref) == trait_to_be_found_erased {
-                        return ControlFlow::Break(vtable_base);
+                    if tcx
+                        .normalize_erasing_late_bound_regions(ty::ParamEnv::reveal_all(), trait_ref)
+                        == target_principal
+                    {
+                        return ControlFlow::Break(vptr_offset);
                     }
-                    vtable_base += count_own_vtable_entries(tcx, trait_ref);
+
+                    vptr_offset += tcx.own_existential_vtable_entries(trait_ref.def_id()).len();
+
                     if emit_vptr {
-                        vtable_base += 1;
+                        vptr_offset += 1;
                     }
                 }
             }
@@ -351,55 +359,72 @@ pub(super) fn vtable_trait_first_method_offset<'tcx>(
         }
     };
 
-    if let Some(vtable_base) =
-        prepare_vtable_segments(tcx, trait_owning_vtable, vtable_segment_callback)
-    {
-        vtable_base
-    } else {
-        bug!("Failed to find info for expected trait in vtable");
-    }
+    prepare_vtable_segments(tcx, source_principal, vtable_segment_callback).unwrap()
 }
 
-/// Find slot offset for trait vptr within vtable entries of another trait
-pub(crate) fn vtable_trait_upcasting_coercion_new_vptr_slot<'tcx>(
+/// Given a `dyn Subtrait` and `dyn Supertrait` trait object, find the slot of
+/// // the trait vptr in the subtrait's vtable.
+pub(crate) fn supertrait_vtable_slot<'tcx>(
     tcx: TyCtxt<'tcx>,
     key: (
-        Ty<'tcx>, // trait object type whose trait owning vtable
-        Ty<'tcx>, // trait object for supertrait
+        Ty<'tcx>, // Source -- `dyn Subtrait`.
+        Ty<'tcx>, // Target -- `dyn Supertrait` being coerced to.
     ),
 ) -> Option<usize> {
+    debug_assert!(!key.has_non_region_infer() && !key.has_non_region_param());
+
     let (source, target) = key;
-    assert!(matches!(&source.kind(), &ty::Dynamic(..)) && !source.has_infer());
-    assert!(matches!(&target.kind(), &ty::Dynamic(..)) && !target.has_infer());
+    let ty::Dynamic(source, _, _) = *source.kind() else {
+        bug!();
+    };
+    let source_principal = tcx
+        .normalize_erasing_regions(ty::ParamEnv::reveal_all(), source.principal().unwrap())
+        .with_self_ty(tcx, tcx.types.trait_object_dummy_self);
 
-    // this has been typecked-before, so diagnostics is not really needed.
-    let unsize_trait_did = tcx.require_lang_item(LangItem::Unsize, None);
+    let ty::Dynamic(target, _, _) = *target.kind() else {
+        bug!();
+    };
+    let target_principal = tcx
+        .normalize_erasing_regions(ty::ParamEnv::reveal_all(), target.principal().unwrap())
+        .with_self_ty(tcx, tcx.types.trait_object_dummy_self);
 
-    let trait_ref = ty::TraitRef::new(tcx, unsize_trait_did, [source, target]);
+    let vtable_segment_callback = {
+        let mut vptr_offset = 0;
+        move |segment| {
+            match segment {
+                VtblSegment::MetadataDSA => {
+                    vptr_offset += TyCtxt::COMMON_VTABLE_ENTRIES.len();
+                }
+                VtblSegment::TraitOwnEntries { trait_ref, emit_vptr } => {
+                    vptr_offset += tcx.own_existential_vtable_entries(trait_ref.def_id()).len();
+                    if tcx.normalize_erasing_regions(ty::ParamEnv::reveal_all(), trait_ref)
+                        == target_principal
+                    {
+                        if emit_vptr {
+                            return ControlFlow::Break(Some(vptr_offset));
+                        } else {
+                            return ControlFlow::Break(None);
+                        }
+                    }
 
-    match tcx.codegen_select_candidate((ty::ParamEnv::reveal_all(), trait_ref)) {
-        Ok(ImplSource::Builtin(BuiltinImplSource::TraitUpcasting { vtable_vptr_slot }, _)) => {
-            *vtable_vptr_slot
+                    if emit_vptr {
+                        vptr_offset += 1;
+                    }
+                }
+            }
+            ControlFlow::Continue(())
         }
-        otherwise => bug!("expected TraitUpcasting candidate, got {otherwise:?}"),
-    }
-}
+    };
 
-/// Given a trait `trait_ref`, returns the number of vtable entries
-/// that come from `trait_ref`, excluding its supertraits. Used in
-/// computing the vtable base for an upcast trait of a trait object.
-pub(crate) fn count_own_vtable_entries<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    trait_ref: ty::PolyTraitRef<'tcx>,
-) -> usize {
-    tcx.own_existential_vtable_entries(trait_ref.def_id()).len()
+    prepare_vtable_segments(tcx, source_principal, vtable_segment_callback).unwrap()
 }
 
 pub(super) fn provide(providers: &mut Providers) {
     *providers = Providers {
         own_existential_vtable_entries,
         vtable_entries,
-        vtable_trait_upcasting_coercion_new_vptr_slot,
+        first_method_vtable_slot,
+        supertrait_vtable_slot,
         ..*providers
     };
 }