inherit `#[align]` from trait method prototypes

4 files changed, 93 insertions, 26 deletions

diff --git a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
index acdda32d58a..47c1e9085e3 100644
--- a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
+++ b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
@@ -1,6 +1,6 @@
 use std::str::FromStr;
 
-use rustc_abi::ExternAbi;
+use rustc_abi::{Align, ExternAbi};
 use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode};
 use rustc_ast::{LitKind, MetaItem, MetaItemInner, attr};
 use rustc_attr_data_structures::{
@@ -402,6 +402,9 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
     codegen_fn_attrs.alignment =
         Ord::max(codegen_fn_attrs.alignment, tcx.sess.opts.unstable_opts.min_function_alignment);
 
+    // On trait methods, inherit the `#[align]` of the trait's method prototype.
+    codegen_fn_attrs.alignment = Ord::max(codegen_fn_attrs.alignment, tcx.inherited_align(did));
+
     let inline_span;
     (codegen_fn_attrs.inline, inline_span) = if let Some((inline_attr, span)) =
         find_attr!(attrs, AttributeKind::Inline(i, span) => (*i, *span))
@@ -556,17 +559,26 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
     codegen_fn_attrs
 }
 
+/// If the provided DefId is a method in a trait impl, return the DefId of the method prototype.
+fn opt_trait_item(tcx: TyCtxt<'_>, def_id: DefId) -> Option<DefId> {
+    let impl_item = tcx.opt_associated_item(def_id)?;
+    match impl_item.container {
+        ty::AssocItemContainer::Impl => impl_item.trait_item_def_id,
+        _ => None,
+    }
+}
+
 /// Checks if the provided DefId is a method in a trait impl for a trait which has track_caller
 /// applied to the method prototype.
 fn should_inherit_track_caller(tcx: TyCtxt<'_>, def_id: DefId) -> bool {
-    if let Some(impl_item) = tcx.opt_associated_item(def_id)
-        && let ty::AssocItemContainer::Impl = impl_item.container
-        && let Some(trait_item) = impl_item.trait_item_def_id
-    {
-        return tcx.codegen_fn_attrs(trait_item).flags.intersects(CodegenFnAttrFlags::TRACK_CALLER);
-    }
+    let Some(trait_item) = opt_trait_item(tcx, def_id) else { return false };
+    tcx.codegen_fn_attrs(trait_item).flags.intersects(CodegenFnAttrFlags::TRACK_CALLER)
+}
 
-    false
+/// If the provided DefId is a method in a trait impl, return the value of the `#[align]`
+/// attribute on the method prototype (if any).
+fn inherited_align<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<Align> {
+    tcx.codegen_fn_attrs(opt_trait_item(tcx, def_id)?).alignment
 }
 
 fn check_link_ordinal(tcx: TyCtxt<'_>, attr: &hir::Attribute) -> Option<u16> {
@@ -734,5 +746,6 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
 }
 
 pub(crate) fn provide(providers: &mut Providers) {
-    *providers = Providers { codegen_fn_attrs, should_inherit_track_caller, ..*providers };
+    *providers =
+        Providers { codegen_fn_attrs, should_inherit_track_caller, inherited_align, ..*providers };
 }
diff --git a/compiler/rustc_middle/src/query/erase.rs b/compiler/rustc_middle/src/query/erase.rs
index 26a31cb055e..f138c5ca039 100644
--- a/compiler/rustc_middle/src/query/erase.rs
+++ b/compiler/rustc_middle/src/query/erase.rs
@@ -266,6 +266,7 @@ trivial! {
     Option<rustc_target::spec::PanicStrategy>,
     Option<usize>,
     Option<rustc_middle::ty::IntrinsicDef>,
+    Option<rustc_abi::Align>,
     Result<(), rustc_errors::ErrorGuaranteed>,
     Result<(), rustc_middle::traits::query::NoSolution>,
     Result<rustc_middle::traits::EvaluationResult, rustc_middle::traits::OverflowError>,
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index ddedea32112..a57475f4fa8 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -67,6 +67,7 @@ use std::mem;
 use std::path::PathBuf;
 use std::sync::Arc;
 
+use rustc_abi::Align;
 use rustc_arena::TypedArena;
 use rustc_ast::expand::StrippedCfgItem;
 use rustc_ast::expand::allocator::AllocatorKind;
@@ -1481,6 +1482,10 @@ rustc_queries! {
         desc { |tcx| "computing should_inherit_track_caller of `{}`", tcx.def_path_str(def_id) }
     }
 
+    query inherited_align(def_id: DefId) -> Option<Align> {
+        desc { |tcx| "computing inherited_align 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) }
         cache_on_disk_if { def_id.is_local() }
diff --git a/tests/codegen/align-fn.rs b/tests/codegen/align-fn.rs
index 267da060240..90073ff3081 100644
--- a/tests/codegen/align-fn.rs
+++ b/tests/codegen/align-fn.rs
@@ -1,10 +1,10 @@
-//@ compile-flags: -C no-prepopulate-passes -Z mir-opt-level=0
+//@ compile-flags: -C no-prepopulate-passes -Z mir-opt-level=0 -Clink-dead-code
 
 #![crate_type = "lib"]
 #![feature(fn_align)]
 
 // CHECK: align 16
-#[no_mangle]
+#[unsafe(no_mangle)]
 #[align(16)]
 pub fn fn_align() {}
 
@@ -12,12 +12,12 @@ pub struct A;
 
 impl A {
     // CHECK: align 16
-    #[no_mangle]
+    #[unsafe(no_mangle)]
     #[align(16)]
     pub fn method_align(self) {}
 
     // CHECK: align 16
-    #[no_mangle]
+    #[unsafe(no_mangle)]
     #[align(16)]
     pub fn associated_fn() {}
 }
@@ -25,46 +25,94 @@ impl A {
 trait T: Sized {
     fn trait_fn() {}
 
-    // CHECK: align 32
-    #[align(32)]
     fn trait_method(self) {}
+
+    #[align(8)]
+    fn trait_method_inherit_low(self);
+
+    #[align(32)]
+    fn trait_method_inherit_high(self);
+
+    #[align(32)]
+    fn trait_method_inherit_default(self) {}
+
+    #[align(4)]
+    #[align(128)]
+    #[align(8)]
+    fn inherit_highest(self) {}
 }
 
 impl T for A {
-    // CHECK: align 16
-    #[no_mangle]
+    // CHECK-LABEL: trait_fn
+    // CHECK-SAME: align 16
+    #[unsafe(no_mangle)]
     #[align(16)]
     fn trait_fn() {}
 
-    // CHECK: align 16
-    #[no_mangle]
+    // CHECK-LABEL: trait_method
+    // CHECK-SAME: align 16
+    #[unsafe(no_mangle)]
     #[align(16)]
     fn trait_method(self) {}
-}
 
-impl T for () {}
+    // The prototype's align is ignored because the align here is higher.
+    // CHECK-LABEL: trait_method_inherit_low
+    // CHECK-SAME: align 16
+    #[unsafe(no_mangle)]
+    #[align(16)]
+    fn trait_method_inherit_low(self) {}
+
+    // The prototype's align is used because it is higher.
+    // CHECK-LABEL: trait_method_inherit_high
+    // CHECK-SAME: align 32
+    #[unsafe(no_mangle)]
+    #[align(16)]
+    fn trait_method_inherit_high(self) {}
+
+    // The prototype's align inherited.
+    // CHECK-LABEL: trait_method_inherit_default
+    // CHECK-SAME: align 32
+    #[unsafe(no_mangle)]
+    fn trait_method_inherit_default(self) {}
+
+    // The prototype's highest align inherited.
+    // CHECK-LABEL: inherit_highest
+    // CHECK-SAME: align 128
+    #[unsafe(no_mangle)]
+    #[align(32)]
+    #[align(64)]
+    fn inherit_highest(self) {}
+}
 
-pub fn foo() {
-    ().trait_method();
+trait HasDefaultImpl: Sized {
+    // CHECK-LABEL: inherit_from_default_method
+    // CHECK-LABEL: inherit_from_default_method
+    // CHECK-SAME: align 32
+    #[align(32)]
+    fn inherit_from_default_method(self) {}
 }
 
+pub struct InstantiateDefaultMethods;
+
+impl HasDefaultImpl for InstantiateDefaultMethods {}
+
 // CHECK-LABEL: align_specified_twice_1
 // CHECK-SAME: align 64
-#[no_mangle]
+#[unsafe(no_mangle)]
 #[align(32)]
 #[align(64)]
 pub fn align_specified_twice_1() {}
 
 // CHECK-LABEL: align_specified_twice_2
 // CHECK-SAME: align 128
-#[no_mangle]
+#[unsafe(no_mangle)]
 #[align(128)]
 #[align(32)]
 pub fn align_specified_twice_2() {}
 
 // CHECK-LABEL: align_specified_twice_3
 // CHECK-SAME: align 256
-#[no_mangle]
+#[unsafe(no_mangle)]
 #[align(32)]
 #[align(256)]
 pub fn align_specified_twice_3() {}