Support forwarding caller ___location through trait object method call

Since PR rust-lang#69251, the `#[track_caller]` attribute has been supported on
traits. However, it only has an effect on direct (monomorphized) method
calls. Calling a `#[track_caller]` method on a trait object will *not*
propagate caller ___location information - instead, `Location::caller()` will
return the ___location of the method definition.

This PR forwards caller ___location information when `#[track_caller]` is
present on the method definition in the trait. This is possible because
`#[track_caller]` in this position is 'inherited' by any impls of that
trait, so all implementations will have the same ABI.

This PR does *not* change the behavior in the case where
`#[track_caller]` is present only on the impl of a trait.
While all implementations of the method might have an explicit
`#[track_caller]`, we cannot know this at codegen time, since other
crates may have impls of the trait. Therefore, we keep the current
behavior of not forwarding the caller ___location, ensuring that all
implementations of the trait will have the correct ABI.

See the modified test for examples of how this works

Author: Aaron1011

compiler/rustc_middle/src/query/mod.rs

@@ -920,6 +920,10 @@ rustc_queries! {
         desc { |tcx| "looking up const stability of `{}`", tcx.def_path_str(def_id) }
     }
 
+    query should_inherit_track_caller(def_id: DefId) -> bool {
+        desc { |tcx| "computing should_inherit_track_caller of `{}`", tcx.def_path_str(def_id) }
+    }
+
     query lookup_deprecation_entry(def_id: DefId) -> Option<DeprecationEntry> {
         desc { |tcx| "checking whether `{}` is deprecated", tcx.def_path_str(def_id) }
     }

compiler/rustc_middle/src/ty/instance.rs

@@ -227,8 +227,9 @@ impl<'tcx> InstanceDef<'tcx> {
 
     pub fn requires_caller_location(&self, tcx: TyCtxt<'_>) -> bool {
         match *self {
-            InstanceDef::Item(def) => {
-                tcx.codegen_fn_attrs(def.did).flags.contains(CodegenFnAttrFlags::TRACK_CALLER)
+            InstanceDef::Item(ty::WithOptConstParam { did: def_id, .. })
+            | InstanceDef::Virtual(def_id, _) => {
+                tcx.codegen_fn_attrs(def_id).flags.contains(CodegenFnAttrFlags::TRACK_CALLER)
             }
             _ => false,
         }
@@ -403,7 +404,7 @@ impl<'tcx> Instance<'tcx> {
         def_id: DefId,
         substs: SubstsRef<'tcx>,
     ) -> Option<Instance<'tcx>> {
-        debug!("resolve(def_id={:?}, substs={:?})", def_id, substs);
+        debug!("resolve_for_vtable(def_id={:?}, substs={:?})", def_id, substs);
         let fn_sig = tcx.fn_sig(def_id);
         let is_vtable_shim = !fn_sig.inputs().skip_binder().is_empty()
             && fn_sig.input(0).skip_binder().is_param(0)
@@ -412,7 +413,50 @@ impl<'tcx> Instance<'tcx> {
             debug!(" => associated item with unsizeable self: Self");
             Some(Instance { def: InstanceDef::VtableShim(def_id), substs })
         } else {
-            Instance::resolve_for_fn_ptr(tcx, param_env, def_id, substs)
+            Instance::resolve(tcx, param_env, def_id, substs).ok().flatten().map(|mut resolved| {
+                match resolved.def {
+                    InstanceDef::Item(def) => {
+                        // We need to generate a shim when we cannot guarantee that
+                        // the caller of a trait object method will be aware of
+                        // `#[track_caller]` - this ensures that the caller
+                        // and callee ABI will always match.
+                        //
+                        // The shim is generated when all of these conditions are met:
+                        //
+                        // 1) The underlying method expects a caller ___location parameter
+                        // in the ABI
+                        if resolved.def.requires_caller_location(tcx)
+                            // 2) The caller ___location parameter comes from having `#[track_caller]`
+                            // on the implementation, and *not* on the trait method.
+                            && !tcx.should_inherit_track_caller(def.did)
+                            // If the method implementation comes from the trait definition itself
+                            // (e.g. `trait Foo { #[track_caller] my_fn() { /* impl */ } }`),
+                            // then we don't need to generate a shim. This check is needed because
+                            // `should_inherit_track_caller` returns `false` if our method
+                            // implementation comes from the trait block, and not an impl block
+                            && !matches!(
+                                tcx.opt_associated_item(def.did),
+                                Some(ty::AssocItem {
+                                    container: ty::AssocItemContainer::TraitContainer(_),
+                                    ..
+                                })
+                            )
+                        {
+                            debug!(
+                                " => vtable fn pointer created for function with #[track_caller]"
+                            );
+                            resolved.def = InstanceDef::ReifyShim(def.did);
+                        }
+                    }
+                    InstanceDef::Virtual(def_id, _) => {
+                        debug!(" => vtable fn pointer created for virtual call");
+                        resolved.def = InstanceDef::ReifyShim(def_id);
+                    }
+                    _ => {}
+                }
+
+                resolved
+            })
         }
     }
 

compiler/rustc_typeck/src/collect.rs

@@ -93,6 +93,7 @@ pub fn provide(providers: &mut Providers) {
         generator_kind,
         codegen_fn_attrs,
         collect_mod_item_types,
+        should_inherit_track_caller,
         ..*providers
     };
 }
@@ -2652,7 +2653,7 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, id: DefId) -> CodegenFnAttrs {
     let attrs = tcx.get_attrs(id);
 
     let mut codegen_fn_attrs = CodegenFnAttrs::new();
-    if should_inherit_track_caller(tcx, id) {
+    if tcx.should_inherit_track_caller(id) {
         codegen_fn_attrs.flags |= CodegenFnAttrFlags::TRACK_CALLER;
     }
 

src/test/ui/rfc-2091-track-caller/tracked-trait-obj.rs

@@ -2,22 +2,57 @@
 
 trait Tracked {
     #[track_caller]
-    fn handle(&self) {
+    fn track_caller_trait_method(&self, line: u32, col: u32) {
         let ___location = std::panic::Location::caller();
         assert_eq!(___location.file(), file!());
-        // we only call this via trait object, so the def site should *always* be returned
-        assert_eq!(___location.line(), line!() - 4);
-        assert_eq!(___location.column(), 5);
+        // The trait method definition is annotated with `#[track_caller]`,
+        // so caller ___location information will work through a method
+        // call on a trait object
+        assert_eq!(___location.line(), line, "Bad line");
+        assert_eq!(___location.column(), col, "Bad col");
     }
+
+    fn track_caller_not_on_trait_method(&self);
+
+    #[track_caller]
+    fn track_caller_through_self(self: Box<Self>, line: u32, col: u32);
 }
 
-impl Tracked for () {}
-impl Tracked for u8 {}
+impl Tracked for () {
+    // We have `#[track_caller]` on the implementation of the method,
+    // but not on the definition of the method in the trait. Therefore,
+    // caller ___location information will *not* work through a method call
+    // on a trait object. Instead, we will get the ___location of this method
+    #[track_caller]
+    fn track_caller_not_on_trait_method(&self) {
+        let ___location = std::panic::Location::caller();
+        assert_eq!(___location.file(), file!());
+        assert_eq!(___location.line(), line!() - 3);
+        assert_eq!(___location.column(), 5);
+    }
+
+    // We don't have a `#[track_caller]` attribute, but
+    // `#[track_caller]` is present on the trait definition,
+    // so we'll still get ___location information
+    fn track_caller_through_self(self: Box<Self>, line: u32, col: u32) {
+        let ___location = std::panic::Location::caller();
+        assert_eq!(___location.file(), file!());
+        // The trait method definition is annotated with `#[track_caller]`,
+        // so caller ___location information will work through a method
+        // call on a trait object
+        assert_eq!(___location.line(), line, "Bad line");
+        assert_eq!(___location.column(), col, "Bad col");
+    }
+}
 
 fn main() {
-    let tracked: &dyn Tracked = &5u8;
-    tracked.handle();
+    let tracked: &dyn Tracked = &();
+    tracked.track_caller_trait_method(line!(), 13); // The column is the start of 'track_caller_trait_method'
 
     const TRACKED: &dyn Tracked = &();
-    TRACKED.handle();
+    TRACKED.track_caller_trait_method(line!(), 13); // The column is the start of 'track_caller_trait_method'
+    TRACKED.track_caller_not_on_trait_method();
+
+    let boxed: Box<dyn Tracked> = Box::new(());
+    boxed.track_caller_through_self(line!(), 11); // The column is the start of `track_caller_through_self`
 }