Auto merge of #89619 - michaelwoerister:incr-vtables, r=nagisa

Turn vtable_allocation() into a query

This PR removes the untracked vtable-const-allocation cache from the `tcx` and turns the `vtable_allocation()` method into a query.

The change is pretty straightforward and should be backportable without too much effort.

Fixes https://github.com/rust-lang/rust/issues/89598.

compiler/rustc_codegen_cranelift/src/vtable.rs

@@ -68,7 +68,7 @@ pub(crate) fn get_vtable<'tcx>(
     ty: Ty<'tcx>,
     trait_ref: Option<ty::PolyExistentialTraitRef<'tcx>>,
 ) -> Value {
-    let alloc_id = fx.tcx.vtable_allocation(ty, trait_ref);
+    let alloc_id = fx.tcx.vtable_allocation((ty, trait_ref));
     let data_id =
         data_id_for_alloc_id(&mut fx.constants_cx, &mut *fx.module, alloc_id, Mutability::Not);
     let local_data_id = fx.module.declare_data_in_func(data_id, &mut fx.bcx.func);

compiler/rustc_codegen_ssa/src/meth.rs

@@ -72,7 +72,7 @@ pub fn get_vtable<'tcx, Cx: CodegenMethods<'tcx>>(
         return val;
     }
 
-    let vtable_alloc_id = tcx.vtable_allocation(ty, trait_ref);
+    let vtable_alloc_id = tcx.vtable_allocation((ty, trait_ref));
     let vtable_allocation = tcx.global_alloc(vtable_alloc_id).unwrap_memory();
     let vtable_const = cx.const_data_from_alloc(vtable_allocation);
     let align = cx.data_layout().pointer_align.abi;

compiler/rustc_const_eval/src/interpret/traits.rs

@@ -30,7 +30,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         ensure_monomorphic_enough(*self.tcx, ty)?;
         ensure_monomorphic_enough(*self.tcx, poly_trait_ref)?;
 
-        let vtable_allocation = self.tcx.vtable_allocation(ty, poly_trait_ref);
+        let vtable_allocation = self.tcx.vtable_allocation((ty, poly_trait_ref));
 
         let vtable_ptr = self.memory.global_base_pointer(Pointer::from(vtable_allocation))?;
 

compiler/rustc_middle/src/query/mod.rs

@@ -1012,6 +1012,13 @@ rustc_queries! {
             key.1, key.0 }
     }
 
+    query vtable_allocation(key: (Ty<'tcx>, Option<ty::PolyExistentialTraitRef<'tcx>>)) -> mir::interpret::AllocId {
+        desc { |tcx| "vtable const allocation for <{} as {}>",
+            key.0,
+            key.1.map(|trait_ref| format!("{}", trait_ref)).unwrap_or("_".to_owned())
+        }
+    }
+
     query codegen_fulfill_obligation(
         key: (ty::ParamEnv<'tcx>, ty::PolyTraitRef<'tcx>)
     ) -> Result<ImplSource<'tcx, ()>, ErrorReported> {

compiler/rustc_middle/src/ty/context.rs

@@ -8,7 +8,7 @@ use crate::lint::{struct_lint_level, LintDiagnosticBuilder, LintLevelSource};
 use crate::middle;
 use crate::middle::resolve_lifetime::{self, LifetimeScopeForPath, ObjectLifetimeDefault};
 use crate::middle::stability;
-use crate::mir::interpret::{self, AllocId, Allocation, ConstValue, Scalar};
+use crate::mir::interpret::{self, Allocation, ConstValue, Scalar};
 use crate::mir::{Body, Field, Local, Place, PlaceElem, ProjectionKind, Promoted};
 use crate::thir::Thir;
 use crate::traits;
@@ -1047,10 +1047,6 @@ pub struct GlobalCtxt<'tcx> {
     pub(crate) alloc_map: Lock<interpret::AllocMap<'tcx>>,
 
     output_filenames: Arc<OutputFilenames>,
-
-    // FIXME(eddyb) this doesn't belong here and should be using a query.
-    pub(super) vtables_cache:
-        Lock<FxHashMap<(Ty<'tcx>, Option<ty::PolyExistentialTraitRef<'tcx>>), AllocId>>,
 }
 
 impl<'tcx> TyCtxt<'tcx> {
@@ -1189,7 +1185,6 @@ impl<'tcx> TyCtxt<'tcx> {
             const_stability_interner: Default::default(),
             alloc_map: Lock::new(interpret::AllocMap::new()),
             output_filenames: Arc::new(output_filenames),
-            vtables_cache: Default::default(),
         }
     }
 

compiler/rustc_middle/src/ty/mod.rs

@@ -2051,6 +2051,7 @@ pub fn provide(providers: &mut ty::query::Providers) {
         trait_impls_of: trait_def::trait_impls_of_provider,
         type_uninhabited_from: inhabitedness::type_uninhabited_from,
         const_param_default: consts::const_param_default,
+        vtable_allocation: vtable::vtable_allocation_provider,
         ..*providers
     };
 }

compiler/rustc_middle/src/ty/print/pretty.rs

@@ -2208,6 +2208,7 @@ forward_display_to_print! {
     // because `for<'tcx>` isn't possible yet.
     ty::Binder<'tcx, ty::ExistentialPredicate<'tcx>>,
     ty::Binder<'tcx, ty::TraitRef<'tcx>>,
+    ty::Binder<'tcx, ty::ExistentialTraitRef<'tcx>>,
     ty::Binder<'tcx, TraitRefPrintOnlyTraitPath<'tcx>>,
     ty::Binder<'tcx, TraitRefPrintOnlyTraitName<'tcx>>,
     ty::Binder<'tcx, ty::FnSig<'tcx>>,

compiler/rustc_middle/src/ty/vtable.rs

@@ -43,20 +43,13 @@ pub const COMMON_VTABLE_ENTRIES_DROPINPLACE: usize = 0;
 pub const COMMON_VTABLE_ENTRIES_SIZE: usize = 1;
 pub const COMMON_VTABLE_ENTRIES_ALIGN: usize = 2;
 
-impl<'tcx> TyCtxt<'tcx> {
-    /// Retrieves an allocation that represents the contents of a vtable.
-    /// There's a cache within `TyCtxt` so it will be deduplicated.
-    pub fn vtable_allocation(
-        self,
-        ty: Ty<'tcx>,
-        poly_trait_ref: Option<ty::PolyExistentialTraitRef<'tcx>>,
-    ) -> AllocId {
-        let tcx = self;
-        let vtables_cache = tcx.vtables_cache.lock();
-        if let Some(alloc_id) = vtables_cache.get(&(ty, poly_trait_ref)).cloned() {
-            return alloc_id;
-        }
-        drop(vtables_cache);
+/// Retrieves an allocation that represents the contents of a vtable.
+/// Since this is a query, allocations are cached and not duplicated.
+pub(super) fn vtable_allocation_provider<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    key: (Ty<'tcx>, Option<ty::PolyExistentialTraitRef<'tcx>>),
+) -> AllocId {
+    let (ty, poly_trait_ref) = key;
 
     let vtable_entries = if let Some(poly_trait_ref) = poly_trait_ref {
         let trait_ref = poly_trait_ref.with_self_ty(tcx, ty);
@@ -78,8 +71,7 @@ impl<'tcx> TyCtxt<'tcx> {
     let ptr_align = tcx.data_layout.pointer_align.abi;
 
     let vtable_size = ptr_size * u64::try_from(vtable_entries.len()).unwrap();
-        let mut vtable =
-            Allocation::uninit(vtable_size, ptr_align, /* panic_on_fail */ true).unwrap();
+    let mut vtable = Allocation::uninit(vtable_size, ptr_align, /* panic_on_fail */ true).unwrap();
 
     // No need to do any alignment checks on the memory accesses below, because we know the
     // allocation is correctly aligned as we created it above. Also we're only offsetting by
@@ -105,10 +97,9 @@ impl<'tcx> TyCtxt<'tcx> {
                 ScalarMaybeUninit::from_pointer(fn_ptr, &tcx)
             }
             VtblEntry::TraitVPtr(trait_ref) => {
-                    let super_trait_ref = trait_ref.map_bound(|trait_ref| {
-                        ty::ExistentialTraitRef::erase_self_ty(tcx, trait_ref)
-                    });
-                    let supertrait_alloc_id = self.vtable_allocation(ty, Some(super_trait_ref));
+                let super_trait_ref = trait_ref
+                    .map_bound(|trait_ref| ty::ExistentialTraitRef::erase_self_ty(tcx, trait_ref));
+                let supertrait_alloc_id = tcx.vtable_allocation((ty, Some(super_trait_ref)));
                 let vptr = Pointer::from(supertrait_alloc_id);
                 ScalarMaybeUninit::from_pointer(vptr, &tcx)
             }
@@ -119,9 +110,5 @@ impl<'tcx> TyCtxt<'tcx> {
     }
 
     vtable.mutability = Mutability::Not;
-        let alloc_id = tcx.create_memory_alloc(tcx.intern_const_alloc(vtable));
-        let mut vtables_cache = self.vtables_cache.lock();
-        vtables_cache.insert((ty, poly_trait_ref), alloc_id);
-        alloc_id
-    }
+    tcx.create_memory_alloc(tcx.intern_const_alloc(vtable))
 }

compiler/rustc_query_impl/src/keys.rs

@@ -72,6 +72,17 @@ impl<'tcx> Key for mir::interpret::GlobalId<'tcx> {
     }
 }
 
+impl<'tcx> Key for (Ty<'tcx>, Option<ty::PolyExistentialTraitRef<'tcx>>) {
+    #[inline(always)]
+    fn query_crate_is_local(&self) -> bool {
+        true
+    }
+
+    fn default_span(&self, _: TyCtxt<'_>) -> Span {
+        DUMMY_SP
+    }
+}
+
 impl<'tcx> Key for mir::interpret::LitToConstInput<'tcx> {
     #[inline(always)]
     fn query_crate_is_local(&self) -> bool {

src/test/incremental/reorder_vtable.rs

@@ -0,0 +1,41 @@
+// revisions:rpass1 rpass2
+
+// This test case makes sure re-order the methods in a vtable will
+// trigger recompilation of codegen units that instantiate it.
+//
+// See https://github.com/rust-lang/rust/issues/89598
+
+trait Foo {
+    #[cfg(rpass1)]
+    fn method1(&self) -> u32;
+
+    fn method2(&self) -> u32;
+
+    #[cfg(rpass2)]
+    fn method1(&self) -> u32;
+}
+
+impl Foo for u32 {
+    fn method1(&self) -> u32 { 17 }
+    fn method2(&self) -> u32 { 42 }
+}
+
+fn main() {
+    // Before #89598 was fixed, the vtable allocation would be cached during
+    // a MIR optimization pass and then the codegen pass for the main object
+    // file would not register a dependency on it (because of the missing
+    // dep-tracking).
+    //
+    // In the rpass2 session, the main object file would not be re-compiled,
+    // thus the mod1::foo(x) call would pass in an outdated vtable, while the
+    // mod1 object would expect the new, re-ordered vtable, resulting in a
+    // call to the wrong method.
+    let x: &dyn Foo = &0u32;
+    assert_eq!(mod1::foo(x), 17);
+}
+
+mod mod1 {
+    pub(super) fn foo(x: &dyn super::Foo) -> u32 {
+        x.method1()
+    }
+}