Implement struct_target_features for non-generic functions.

Author: veluca93

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -1,5 +1,6 @@
 use rustc_ast::{ast, attr, MetaItemKind, NestedMetaItem};
 use rustc_attr::{list_contains_name, InlineAttr, InstructionSetAttr, OptimizeAttr};
+use rustc_data_structures::fx::FxHashSet;
 use rustc_errors::codes::*;
 use rustc_errors::{struct_span_code_err, DiagMessage, SubdiagMessage};
 use rustc_hir as hir;
@@ -8,15 +9,16 @@ use rustc_hir::def_id::{DefId, LocalDefId, LOCAL_CRATE};
 use rustc_hir::weak_lang_items::WEAK_LANG_ITEMS;
 use rustc_hir::{lang_items, LangItem};
 use rustc_middle::middle::codegen_fn_attrs::{
-    CodegenFnAttrFlags, CodegenFnAttrs, PatchableFunctionEntry,
+    CodegenFnAttrFlags, CodegenFnAttrs, PatchableFunctionEntry, TargetFeature,
 };
 use rustc_middle::mir::mono::Linkage;
 use rustc_middle::query::Providers;
-use rustc_middle::ty::{self as ty, TyCtxt};
+use rustc_middle::ty::{self as ty, Ty, TyCtxt};
 use rustc_session::lint;
 use rustc_session::parse::feature_err;
 use rustc_span::symbol::Ident;
 use rustc_span::{sym, Span};
+use rustc_target::abi::VariantIdx;
 use rustc_target::spec::{abi, SanitizerSet};
 
 use crate::errors;
@@ -78,23 +80,26 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
     let mut link_ordinal_span = None;
     let mut no_sanitize_span = None;
 
+    let fn_sig_outer = || {
+        use DefKind::*;
+
+        let def_kind = tcx.def_kind(did);
+        if let Fn | AssocFn | Variant | Ctor(..) = def_kind { Some(tcx.fn_sig(did)) } else { None }
+    };
+
     for attr in attrs.iter() {
         // In some cases, attribute are only valid on functions, but it's the `check_attr`
         // pass that check that they aren't used anywhere else, rather this module.
         // In these cases, we bail from performing further checks that are only meaningful for
         // functions (such as calling `fn_sig`, which ICEs if given a non-function). We also
         // report a delayed bug, just in case `check_attr` isn't doing its job.
         let fn_sig = || {
-            use DefKind::*;
-
-            let def_kind = tcx.def_kind(did);
-            if let Fn | AssocFn | Variant | Ctor(..) = def_kind {
-                Some(tcx.fn_sig(did))
-            } else {
+            let sig = fn_sig_outer();
+            if sig.is_none() {
                 tcx.dcx()
                     .span_delayed_bug(attr.span, "this attribute can only be applied to functions");
-                None
             }
+            sig
         };
 
         let Some(Ident { name, .. }) = attr.ident() else {
@@ -613,7 +618,30 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
         }
     }
 
-    // If a function uses #[target_feature] it can't be inlined into general
+    if let Some(sig) = fn_sig_outer() {
+        for ty in sig.skip_binder().inputs().skip_binder() {
+            let additional_tf =
+                tcx.struct_reachable_target_features(tcx.param_env(did.to_def_id()).and(*ty));
+            // FIXME(struct_target_features): is this really necessary?
+            if !additional_tf.is_empty() && sig.skip_binder().abi() != abi::Abi::Rust {
+                tcx.dcx().span_err(
+                    tcx.hir().span(tcx.local_def_id_to_hir_id(did)),
+                    "cannot use a struct with target features in a function with non-Rust ABI",
+                );
+            }
+            if !additional_tf.is_empty() && codegen_fn_attrs.inline == InlineAttr::Always {
+                tcx.dcx().span_err(
+                    tcx.hir().span(tcx.local_def_id_to_hir_id(did)),
+                    "cannot use a struct with target features in a #[inline(always)] function",
+                );
+            }
+            codegen_fn_attrs
+                .target_features
+                .extend(additional_tf.iter().map(|tf| TargetFeature { implied: true, ..*tf }));
+        }
+    }
+
+    // If a function uses non-default target_features it can't be inlined into general
     // purpose functions as they wouldn't have the right target features
     // enabled. For that reason we also forbid #[inline(always)] as it can't be
     // respected.
@@ -758,6 +786,91 @@ fn check_link_name_xor_ordinal(
     }
 }
 
+fn struct_target_features(tcx: TyCtxt<'_>, def_id: LocalDefId) -> &[TargetFeature] {
+    let mut features = vec![];
+    let supported_features = tcx.supported_target_features(LOCAL_CRATE);
+    for attr in tcx.get_attrs(def_id, sym::target_feature) {
+        from_target_feature(tcx, attr, supported_features, &mut features);
+    }
+    tcx.arena.alloc_slice(&features)
+}
+
+fn struct_reachable_target_features<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>,
+) -> &'tcx [TargetFeature] {
+    // Collect target features from types reachable from `env.value`.
+    // We define a type as "reachable" if:
+    //  - it is a function argument
+    //  - it is a field of a reachable struct
+    //  - there is a reachable reference to it
+    let mut visited_types = FxHashSet::default();
+    let mut reachable_types = vec![env.value];
+    let mut reachable_tf = vec![];
+
+    while let Some(ty) = reachable_types.pop() {
+        if visited_types.contains(&ty) {
+            continue;
+        }
+        visited_types.insert(ty);
+        match ty.kind() {
+            ty::Alias(..) => {
+                if let Ok(t) = tcx.try_normalize_erasing_regions(env.param_env, ty) {
+                    reachable_types.push(t)
+                }
+            }
+
+            ty::Ref(_, inner, _) => reachable_types.push(*inner),
+            ty::Tuple(tys) => reachable_types.extend(tys.iter()),
+            ty::Adt(adt_def, args) => {
+                reachable_tf.extend_from_slice(tcx.struct_target_features(adt_def.did()));
+                // This only recurses into structs as i.e. an Option<TargetFeature> is an ADT
+                // that doesn't actually always contain a TargetFeature.
+                if adt_def.is_struct() {
+                    reachable_types.extend(
+                        adt_def
+                            .variant(VariantIdx::from_usize(0))
+                            .fields
+                            .iter()
+                            .map(|field| field.ty(tcx, args)),
+                    );
+                }
+            }
+            ty::Bool
+            | ty::Char
+            | ty::Int(..)
+            | ty::Uint(..)
+            | ty::Float(..)
+            | ty::Foreign(..)
+            | ty::Str
+            | ty::Array(..)
+            | ty::Pat(..)
+            | ty::Slice(..)
+            | ty::RawPtr(..)
+            | ty::FnDef(..)
+            | ty::FnPtr(..)
+            | ty::Dynamic(..)
+            | ty::Closure(..)
+            | ty::CoroutineClosure(..)
+            | ty::Coroutine(..)
+            | ty::CoroutineWitness(..)
+            | ty::Never
+            | ty::Param(..)
+            | ty::Bound(..)
+            | ty::Placeholder(..)
+            | ty::Infer(..)
+            | ty::Error(..) => (),
+        }
+    }
+    tcx.arena.alloc_slice(&reachable_tf)
+}
+
 pub fn provide(providers: &mut Providers) {
-    *providers = Providers { codegen_fn_attrs, should_inherit_track_caller, ..*providers };
+    *providers = Providers {
+        codegen_fn_attrs,
+        should_inherit_track_caller,
+        struct_target_features,
+        struct_reachable_target_features,
+        ..*providers
+    };
 }

compiler/rustc_feature/src/unstable.rs

@@ -597,6 +597,8 @@ declare_features! (
     (unstable, strict_provenance, "1.61.0", Some(95228)),
     /// Allows string patterns to dereference values to match them.
     (unstable, string_deref_patterns, "1.67.0", Some(87121)),
+    /// Allows structs to carry target_feature information.
+    (incomplete, struct_target_features, "CURRENT_RUSTC_VERSION", Some(129107)),
     /// Allows the use of `#[target_feature]` on safe functions.
     (unstable, target_feature_11, "1.45.0", Some(69098)),
     /// Allows using `#[thread_local]` on `static` items.

compiler/rustc_hir/src/def.rs

@@ -326,6 +326,43 @@ impl DefKind {
             | DefKind::ExternCrate => false,
         }
     }
+
+    /// Whether `query struct_target_features` should be used with this definition.
+    pub fn has_struct_target_features(self) -> bool {
+        match self {
+            DefKind::Struct => true,
+            DefKind::Fn
+            | DefKind::Union
+            | DefKind::Enum
+            | DefKind::AssocFn
+            | DefKind::Ctor(..)
+            | DefKind::Closure
+            | DefKind::Static { .. }
+            | DefKind::Mod
+            | DefKind::Variant
+            | DefKind::Trait
+            | DefKind::TyAlias
+            | DefKind::ForeignTy
+            | DefKind::TraitAlias
+            | DefKind::AssocTy
+            | DefKind::Const
+            | DefKind::AssocConst
+            | DefKind::Macro(..)
+            | DefKind::Use
+            | DefKind::ForeignMod
+            | DefKind::OpaqueTy
+            | DefKind::Impl { .. }
+            | DefKind::Field
+            | DefKind::TyParam
+            | DefKind::ConstParam
+            | DefKind::LifetimeParam
+            | DefKind::AnonConst
+            | DefKind::InlineConst
+            | DefKind::SyntheticCoroutineBody
+            | DefKind::GlobalAsm
+            | DefKind::ExternCrate => false,
+        }
+    }
 }
 
 /// The resolution of a path or export.

compiler/rustc_hir_typeck/src/coercion.rs

@@ -851,6 +851,8 @@ impl<'f, 'tcx> Coerce<'f, 'tcx> {
                     }
 
                     // Safe `#[target_feature]` functions are not assignable to safe fn pointers (RFC 2396).
+                    // FIXME(struct_target_features): should this be true also for functions that inherit
+                    // target features from structs?
 
                     if b_hdr.safety == hir::Safety::Safe
                         && !self.tcx.codegen_fn_attrs(def_id).target_features.is_empty()

compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs

@@ -254,6 +254,7 @@ provide! { tcx, def_id, other, cdata,
     variances_of => { table }
     fn_sig => { table }
     codegen_fn_attrs => { table }
+    struct_target_features => { table_defaulted_array }
     impl_trait_header => { table }
     const_param_default => { table }
     object_lifetime_default => { table }

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -1392,6 +1392,9 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             if def_kind.has_codegen_attrs() {
                 record!(self.tables.codegen_fn_attrs[def_id] <- self.tcx.codegen_fn_attrs(def_id));
             }
+            if def_kind.has_struct_target_features() {
+                record_defaulted_array!(self.tables.struct_target_features[def_id] <- self.tcx.struct_target_features(def_id));
+            }
             if should_encode_visibility(def_kind) {
                 let vis =
                     self.tcx.local_visibility(local_id).map_id(|def_id| def_id.local_def_index);

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -19,7 +19,7 @@ use rustc_macros::{
     Decodable, Encodable, MetadataDecodable, MetadataEncodable, TyDecodable, TyEncodable,
 };
 use rustc_middle::metadata::ModChild;
-use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
+use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrs, TargetFeature};
 use rustc_middle::middle::debugger_visualizer::DebuggerVisualizerFile;
 use rustc_middle::middle::exported_symbols::{ExportedSymbol, SymbolExportInfo};
 use rustc_middle::middle::lib_features::FeatureStability;
@@ -402,6 +402,7 @@ define_tables! {
     // individually instead of `DefId`s.
     module_children_reexports: Table<DefIndex, LazyArray<ModChild>>,
     cross_crate_inlinable: Table<DefIndex, bool>,
+    struct_target_features: Table<DefIndex, LazyArray<TargetFeature>>,
 
 - optional:
     attributes: Table<DefIndex, LazyArray<ast::Attribute>>,

compiler/rustc_middle/src/middle/codegen_fn_attrs.rs

@@ -26,8 +26,8 @@ pub struct CodegenFnAttrs {
     /// be set when `link_name` is set. This is for foreign items with the
     /// "raw-dylib" kind.
     pub link_ordinal: Option<u16>,
-    /// The `#[target_feature(enable = "...")]` attribute and the enabled
-    /// features (only enabled features are supported right now).
+    /// All the target features that are enabled for this function. Some features might be enabled
+    /// implicitly.
     pub target_features: Vec<TargetFeature>,
     /// The `#[linkage = "..."]` attribute on Rust-defined items and the value we found.
     pub linkage: Option<Linkage>,
@@ -55,8 +55,8 @@ pub struct CodegenFnAttrs {
 pub struct TargetFeature {
     /// The name of the target feature (e.g. "avx")
     pub name: Symbol,
-    /// The feature is implied by another feature, rather than explicitly added by the
-    /// `#[target_feature]` attribute
+    /// The feature is implied by another feature or by an argument, rather than explicitly
+    /// added by the `#[target_feature]` attribute
     pub implied: bool,
 }
 

compiler/rustc_middle/src/query/mod.rs

@@ -47,7 +47,7 @@ use {rustc_ast as ast, rustc_attr as attr, rustc_hir as hir};
 use crate::infer::canonical::{self, Canonical};
 use crate::lint::LintExpectation;
 use crate::metadata::ModChild;
-use crate::middle::codegen_fn_attrs::CodegenFnAttrs;
+use crate::middle::codegen_fn_attrs::{CodegenFnAttrs, TargetFeature};
 use crate::middle::debugger_visualizer::DebuggerVisualizerFile;
 use crate::middle::exported_symbols::{ExportedSymbol, SymbolExportInfo};
 use crate::middle::lib_features::LibFeatures;
@@ -1249,6 +1249,15 @@ rustc_queries! {
         feedable
     }
 
+    query struct_target_features(def_id: DefId) -> &'tcx [TargetFeature] {
+        separate_provide_extern
+        desc { |tcx| "computing target features for struct `{}`", tcx.def_path_str(def_id) }
+    }
+
+    query struct_reachable_target_features(env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> &'tcx [TargetFeature] {
+        desc { |tcx| "computing target features reachable from {}", env.value }
+    }
+
     query asm_target_features(def_id: DefId) -> &'tcx FxIndexSet<Symbol> {
         desc { |tcx| "computing target features for inline asm of `{}`", tcx.def_path_str(def_id) }
     }

compiler/rustc_middle/src/ty/parameterized.rs

@@ -59,6 +59,7 @@ trivially_parameterized_over_tcx! {
     std::string::String,
     crate::metadata::ModChild,
     crate::middle::codegen_fn_attrs::CodegenFnAttrs,
+    crate::middle::codegen_fn_attrs::TargetFeature,
     crate::middle::debugger_visualizer::DebuggerVisualizerFile,
     crate::middle::exported_symbols::SymbolExportInfo,
     crate::middle::lib_features::FeatureStability,