Auto merge of #126505 - compiler-errors:no-vtable, r=lcnr

Only compute vtable information during codegen

This PR removes vtable information from the `Object` and `TraitUpcasting` candidate sources in the trait solvers, and defers the computation of relevant information to `Instance::resolve`. This is because vtables really aren't a thing in the trait world -- they're an implementation detail in codegen.

Previously it was just easiest to tangle this information together since we were already doing the work of looking at all the supertraits in the trait solver, and specifically because we use traits to represent when it's possible to call a method via a vtable (`Object` candidate) and do upcasting (`Unsize` candidate). but I am somewhat suspicious we're doing a *lot* of extra work, especially in polymorphic contexts, so let's see what perf says.

8 files changed, 100 insertions, 170 deletions

diff --git a/compiler/rustc_trait_selection/src/solve/assembly/mod.rs b/compiler/rustc_trait_selection/src/solve/assembly/mod.rs
index 72cdfb5e57b..1cdbf082078 100644
--- a/compiler/rustc_trait_selection/src/solve/assembly/mod.rs
+++ b/compiler/rustc_trait_selection/src/solve/assembly/mod.rs
@@ -5,6 +5,7 @@ use rustc_hir::def_id::DefId;
 use rustc_hir::LangItem;
 use rustc_infer::infer::InferCtxt;
 use rustc_infer::traits::query::NoSolution;
+use rustc_infer::traits::util::supertraits;
 use rustc_middle::bug;
 use rustc_middle::traits::solve::inspect::ProbeKind;
 use rustc_middle::traits::solve::{Certainty, Goal, MaybeCause, QueryResult};
@@ -744,14 +745,14 @@ impl<'tcx> EvalCtxt<'_, InferCtxt<'tcx>> {
         // a projection goal.
         if let Some(principal) = bounds.principal() {
             let principal_trait_ref = principal.with_self_ty(tcx, self_ty);
-            self.walk_vtable(principal_trait_ref, |ecx, assumption, vtable_base, _| {
+            for (idx, assumption) in supertraits(self.interner(), principal_trait_ref).enumerate() {
                 candidates.extend(G::probe_and_consider_object_bound_candidate(
-                    ecx,
-                    CandidateSource::BuiltinImpl(BuiltinImplSource::Object { vtable_base }),
+                    self,
+                    CandidateSource::BuiltinImpl(BuiltinImplSource::Object(idx)),
                     goal,
                     assumption.upcast(tcx),
                 ));
-            });
+            }
         }
     }
 
diff --git a/compiler/rustc_trait_selection/src/solve/eval_ctxt/mod.rs b/compiler/rustc_trait_selection/src/solve/eval_ctxt/mod.rs
index 860c580374d..6b8375b53e8 100644
--- a/compiler/rustc_trait_selection/src/solve/eval_ctxt/mod.rs
+++ b/compiler/rustc_trait_selection/src/solve/eval_ctxt/mod.rs
@@ -24,7 +24,6 @@ use rustc_type_ir_macros::{Lift_Generic, TypeFoldable_Generic, TypeVisitable_Gen
 use std::ops::ControlFlow;
 
 use crate::traits::coherence;
-use crate::traits::vtable::{count_own_vtable_entries, prepare_vtable_segments, VtblSegment};
 
 use super::inspect::ProofTreeBuilder;
 use super::{search_graph, GoalEvaluationKind, FIXPOINT_STEP_LIMIT};
@@ -1022,41 +1021,6 @@ impl<'tcx> EvalCtxt<'_, InferCtxt<'tcx>> {
             }
         }
     }
-
-    /// Walk through the vtable of a principal trait ref, executing a `supertrait_visitor`
-    /// for every trait ref encountered (including the principal). Passes both the vtable
-    /// base and the (optional) vptr slot.
-    pub(super) fn walk_vtable(
-        &mut self,
-        principal: ty::PolyTraitRef<'tcx>,
-        mut supertrait_visitor: impl FnMut(&mut Self, ty::PolyTraitRef<'tcx>, usize, Option<usize>),
-    ) {
-        let tcx = self.interner();
-        let mut offset = 0;
-        prepare_vtable_segments::<()>(tcx, principal, |segment| {
-            match segment {
-                VtblSegment::MetadataDSA => {
-                    offset += TyCtxt::COMMON_VTABLE_ENTRIES.len();
-                }
-                VtblSegment::TraitOwnEntries { trait_ref, emit_vptr } => {
-                    let own_vtable_entries = count_own_vtable_entries(tcx, trait_ref);
-
-                    supertrait_visitor(
-                        self,
-                        trait_ref,
-                        offset,
-                        emit_vptr.then(|| offset + own_vtable_entries),
-                    );
-
-                    offset += own_vtable_entries;
-                    if emit_vptr {
-                        offset += 1;
-                    }
-                }
-            }
-            ControlFlow::Continue(())
-        });
-    }
 }
 
 /// Eagerly replace aliases with inference variables, emitting `AliasRelate`
diff --git a/compiler/rustc_trait_selection/src/solve/eval_ctxt/select.rs b/compiler/rustc_trait_selection/src/solve/eval_ctxt/select.rs
index 68c0c8bf09e..257fd263b94 100644
--- a/compiler/rustc_trait_selection/src/solve/eval_ctxt/select.rs
+++ b/compiler/rustc_trait_selection/src/solve/eval_ctxt/select.rs
@@ -114,8 +114,8 @@ fn candidate_should_be_dropped_in_favor_of<'tcx>(
         // In the old trait solver, we arbitrarily choose lower vtable candidates
         // over higher ones.
         (
-            CandidateSource::BuiltinImpl(BuiltinImplSource::Object { vtable_base: a }),
-            CandidateSource::BuiltinImpl(BuiltinImplSource::Object { vtable_base: b }),
+            CandidateSource::BuiltinImpl(BuiltinImplSource::Object(a)),
+            CandidateSource::BuiltinImpl(BuiltinImplSource::Object(b)),
         ) => a >= b,
         // Prefer dyn candidates over non-dyn candidates. This is necessary to
         // handle the unsoundness between `impl<T: ?Sized> Any for T` and `dyn Any: Any`.
diff --git a/compiler/rustc_trait_selection/src/solve/trait_goals.rs b/compiler/rustc_trait_selection/src/solve/trait_goals.rs
index 3fdc11ade5f..bcb2ea18f78 100644
--- a/compiler/rustc_trait_selection/src/solve/trait_goals.rs
+++ b/compiler/rustc_trait_selection/src/solve/trait_goals.rs
@@ -9,6 +9,7 @@ use rustc_hir::{LangItem, Movability};
 use rustc_infer::infer::InferCtxt;
 use rustc_infer::traits::query::NoSolution;
 use rustc_infer::traits::solve::MaybeCause;
+use rustc_infer::traits::util::supertraits;
 use rustc_middle::bug;
 use rustc_middle::traits::solve::inspect::ProbeKind;
 use rustc_middle::traits::solve::{CandidateSource, Certainty, Goal, QueryResult};
@@ -756,24 +757,19 @@ impl<'tcx> EvalCtxt<'_, InferCtxt<'tcx>> {
                 a_data.principal(),
             ));
         } else if let Some(a_principal) = a_data.principal() {
-            self.walk_vtable(
-                a_principal.with_self_ty(tcx, a_ty),
-                |ecx, new_a_principal, _, vtable_vptr_slot| {
-                    responses.extend(ecx.consider_builtin_upcast_to_principal(
-                        goal,
-                        CandidateSource::BuiltinImpl(BuiltinImplSource::TraitUpcasting {
-                            vtable_vptr_slot,
-                        }),
-                        a_data,
-                        a_region,
-                        b_data,
-                        b_region,
-                        Some(new_a_principal.map_bound(|trait_ref| {
-                            ty::ExistentialTraitRef::erase_self_ty(tcx, trait_ref)
-                        })),
-                    ));
-                },
-            );
+            for new_a_principal in supertraits(tcx, a_principal.with_self_ty(tcx, a_ty)).skip(1) {
+                responses.extend(self.consider_builtin_upcast_to_principal(
+                    goal,
+                    CandidateSource::BuiltinImpl(BuiltinImplSource::TraitUpcasting),
+                    a_data,
+                    a_region,
+                    b_data,
+                    b_region,
+                    Some(new_a_principal.map_bound(|trait_ref| {
+                        ty::ExistentialTraitRef::erase_self_ty(tcx, trait_ref)
+                    })),
+                ));
+            }
         }
 
         responses
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 074fd487888..af599108c49 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
     };
 }