Auto merge of #155852 - teor2345:fn-arg-splat-fndecl-bit, r=<try>

[perf experiment]: Use 1 bit to track splatting in FnDecl

Author: bors

.gitignore

@@ -56,6 +56,9 @@ no_llvm_build
 /src/bootstrap/target
 /src/ci/citool/target
 /src/tools/x/target
+# Created by `x rust-analyzer[-clippy]` tests
+/library/core/src/iter/traits/target
+/library/core/src/slice/target
 # Created by `x vendor`
 /vendor
 # Created by default with `src/ci/docker/run.sh`

compiler/rustc_ast/src/ast.rs

@@ -3055,9 +3055,40 @@ impl FnDecl {
     pub fn has_self(&self) -> bool {
         self.inputs.get(0).is_some_and(Param::is_self)
     }
+
     pub fn c_variadic(&self) -> bool {
         self.inputs.last().is_some_and(|arg| matches!(arg.ty.kind, TyKind::CVarArgs))
     }
+
+    /// The marker index for "no splatted arguments".
+    /// Must have the same value as `FnSigKind::NO_SPLATTED_ARG_INDEX`.
+    // FIXME(splat): if we remove this limit from hir::FnDecl and FnSig, all instances of this
+    // constant can go away entirely.
+    pub const NO_SPLATTED_ARG_INDEX: u16 = u16::MAX;
+
+    /// Returns a splatted argument index and its span, if any splatted arguments are present.
+    #[inline]
+    pub fn splatted(&self) -> Option<(u16 /* arg_index */, Span)> {
+        let (index, span) = self.inputs.iter().enumerate().find_map(|(index, arg)| {
+            arg.attrs.iter().find_map(|attr| {
+                attr.has_name(sym::splat).then_some((u16::try_from(index).unwrap(), attr.span))
+            })
+        })?;
+
+        if index == Self::NO_SPLATTED_ARG_INDEX {
+            // AST validation has already checked the splatted argument index is valid, so just
+            // ignore invalid indexes here.
+            None
+        } else {
+            Some((index, span))
+        }
+    }
+
+    /// Returns `true` if the function has a splatted argument.
+    #[inline(always)]
+    pub fn has_splatted_arg(&self) -> bool {
+        self.splatted().is_some()
+    }
 }
 
 /// Is the trait definition an auto trait?

compiler/rustc_ast_lowering/src/delegation.rs

@@ -141,7 +141,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
                 let is_method = self.is_method(sig_id, span);
 
-                let (param_count, c_variadic) = self.param_count(sig_id);
+                let (param_count, c_variadic, has_splatted_arg) = self.param_count(sig_id);
 
                 let mut generics =
                     self.uplift_delegation_generics(delegation, sig_id, item_id, is_method);
@@ -154,8 +154,14 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     span,
                 );
 
-                let decl =
-                    self.lower_delegation_decl(sig_id, param_count, c_variadic, span, &generics);
+                let decl = self.lower_delegation_decl(
+                    sig_id,
+                    param_count,
+                    c_variadic,
+                    has_splatted_arg,
+                    span,
+                    &generics,
+                );
 
                 let sig = self.lower_delegation_sig(sig_id, decl, span);
                 let ident = self.lower_ident(delegation.ident);
@@ -269,17 +275,25 @@ impl<'hir> LoweringContext<'_, 'hir> {
         self.get_partial_res(node_id).and_then(|r| r.expect_full_res().opt_def_id())
     }
 
-    // Function parameter count, including C variadic `...` if present.
-    fn param_count(&self, def_id: DefId) -> (usize, bool /*c_variadic*/) {
+    // Function parameter count, including C variadic `...` and `#[splat]` if present.
+    fn param_count(
+        &self,
+        def_id: DefId,
+    ) -> (usize, bool /*c_variadic*/, bool /*has_splatted_arg*/) {
         let sig = self.tcx.fn_sig(def_id).skip_binder().skip_binder();
-        (sig.inputs().len() + usize::from(sig.c_variadic()), sig.c_variadic())
+        (
+            sig.inputs().len() + usize::from(sig.c_variadic()),
+            sig.c_variadic(),
+            sig.splatted().is_some(),
+        )
     }
 
     fn lower_delegation_decl(
         &mut self,
         sig_id: DefId,
         param_count: usize,
         c_variadic: bool,
+        has_splatted_arg: bool,
         span: Span,
         generics: &GenericsGenerationResults<'hir>,
     ) -> &'hir hir::FnDecl<'hir> {
@@ -314,7 +328,8 @@ impl<'hir> LoweringContext<'_, 'hir> {
             output: hir::FnRetTy::Return(output),
             fn_decl_kind: FnDeclFlags::default()
                 .set_lifetime_elision_allowed(true)
-                .set_c_variadic(c_variadic),
+                .set_c_variadic(c_variadic)
+                .set_has_splatted_arg(has_splatted_arg),
         })
     }
 

compiler/rustc_ast_lowering/src/lib.rs

@@ -1767,12 +1767,15 @@ impl<'hir> LoweringContext<'_, 'hir> {
         coro: Option<CoroutineKind>,
     ) -> &'hir hir::FnDecl<'hir> {
         let c_variadic = decl.c_variadic();
+        let mut splatted_arg_index = decl.splatted();
 
         // Skip the `...` (`CVarArgs`) trailing arguments from the AST,
         // as they are not explicit in HIR/Ty function signatures.
         // (instead, the `c_variadic` flag is set to `true`)
         let mut inputs = &decl.inputs[..];
         if decl.c_variadic() {
+            // Splat + variadic errors in AST validation, so just ignore one of them here.
+            splatted_arg_index = None;
             inputs = &inputs[..inputs.len() - 1];
         }
         let inputs = self.arena.alloc_from_iter(inputs.iter().map(|param| {
@@ -1860,7 +1863,18 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 }
             }))
             .set_lifetime_elision_allowed(self.resolver.lifetime_elision_allowed(fn_node_id))
-            .set_c_variadic(c_variadic);
+            .set_c_variadic(c_variadic)
+            .set_has_splatted_arg(splatted_arg_index.is_some());
+
+        if let Some((index, span)) = splatted_arg_index {
+            // For performance, just lower the one attribute fn args care about to HIR.
+            let local_id = inputs[usize::from(index)].hir_id.local_id;
+            assert!(!self.attrs.contains_key(&local_id));
+            self.attrs.insert(
+                local_id,
+                arena_vec![self; hir::Attribute::Parsed(hir::attrs::AttributeKind::Splat(span))],
+            );
+        }
 
         self.arena.alloc(hir::FnDecl { inputs, output, fn_decl_kind })
     }

compiler/rustc_ast_passes/src/ast_validation.rs

@@ -350,7 +350,8 @@ impl<'a> AstValidator<'a> {
 
     fn check_fn_decl(&self, fn_decl: &FnDecl, self_semantic: SelfSemantic) {
         self.check_decl_num_args(fn_decl);
-        self.check_decl_cvariadic_pos(fn_decl);
+        let c_variadic_span = self.check_decl_cvariadic_pos(fn_decl);
+        self.check_decl_splatting(fn_decl, c_variadic_span);
         self.check_decl_attrs(fn_decl);
         self.check_decl_self_param(fn_decl, self_semantic);
     }
@@ -368,17 +369,68 @@ impl<'a> AstValidator<'a> {
     /// Emits an error if a function declaration has a variadic parameter in the
     /// beginning or middle of parameter list.
     /// Example: `fn foo(..., x: i32)` will emit an error.
-    fn check_decl_cvariadic_pos(&self, fn_decl: &FnDecl) {
+    /// Returns true if a C-variadic parameter is found.
+    fn check_decl_cvariadic_pos(&self, fn_decl: &FnDecl) -> Option<Span> {
+        let mut c_variadic_span = None;
+
         match &*fn_decl.inputs {
             [ps @ .., _] => {
                 for Param { ty, span, .. } in ps {
                     if let TyKind::CVarArgs = ty.kind {
+                        c_variadic_span = Some(*span);
                         self.dcx().emit_err(errors::FnParamCVarArgsNotLast { span: *span });
                     }
                 }
             }
             _ => {}
         }
+
+        if let Some(Param { ty, span, .. }) = &fn_decl.inputs.last() {
+            if let TyKind::CVarArgs = ty.kind {
+                c_variadic_span = Some(*span);
+            }
+        }
+
+        c_variadic_span
+    }
+
+    /// Emits an error if a function declaration has more than one splatted argument, with a
+    /// C-variadic parameter, or a splat at an unsupported index (for performance).
+    /// Example: `fn foo(#[splat] x: (), #[splat] y: ())` will emit an error.
+    fn check_decl_splatting(&self, fn_decl: &FnDecl, c_variadic_span: Option<Span>) {
+        let (splatted_arg_indexes, mut splatted_spans): (Vec<u16>, Vec<Span>) = fn_decl
+            .inputs
+            .iter()
+            .enumerate()
+            .filter_map(|(index, arg)| {
+                arg.attrs
+                    .iter()
+                    .any(|attr| attr.has_name(sym::splat))
+                    .then_some((u16::try_from(index).unwrap(), arg.span))
+            })
+            .unzip();
+
+        // A splatted argument at the "no splatted" marker index is not supported (this is an
+        // unlikely edge case).
+        if let (Some(&splatted_arg_index), Some(&splatted_span)) =
+            (splatted_arg_indexes.last(), splatted_spans.last())
+            && splatted_arg_index == FnDecl::NO_SPLATTED_ARG_INDEX
+        {
+            self.dcx()
+                .emit_err(errors::InvalidSplattedArg { splatted_arg_index, span: splatted_span });
+        }
+
+        // Multiple splatted arguments are invalid: we can't know which arguments go in each splat.
+        if splatted_arg_indexes.len() > 1 {
+            self.dcx().emit_err(errors::DuplicateSplattedArgs { spans: splatted_spans.clone() });
+        }
+
+        if let Some(c_variadic_span) = c_variadic_span
+            && !splatted_spans.is_empty()
+        {
+            splatted_spans.push(c_variadic_span);
+            self.dcx().emit_err(errors::CVarArgsAndSplat { spans: splatted_spans });
+        }
     }
 
     fn check_decl_attrs(&self, fn_decl: &FnDecl) {
@@ -394,6 +446,7 @@ impl<'a> AstValidator<'a> {
                     sym::deny,
                     sym::expect,
                     sym::forbid,
+                    sym::splat,
                     sym::warn,
                 ];
                 !attr.has_any_name(&arr) && rustc_attr_parsing::is_builtin_attr(*attr)

compiler/rustc_ast_passes/src/errors.rs

@@ -124,6 +124,33 @@ pub(crate) struct FnParamCVarArgsNotLast {
     pub span: Span,
 }
 
+#[derive(Diagnostic)]
+#[diag("`#[splat]` is not supported on argument index {$splatted_arg_index}")]
+#[help("remove `#[splat]`, or use it on an argument closer to the start of the argument list")]
+pub(crate) struct InvalidSplattedArg {
+    pub splatted_arg_index: u16,
+
+    #[primary_span]
+    #[label("`#[splat]` is not supported here")]
+    pub span: Span,
+}
+
+#[derive(Diagnostic)]
+#[diag("multiple `#[splat]`s are not allowed in the same function")]
+#[help("remove `#[splat]` from all but one argument")]
+pub(crate) struct DuplicateSplattedArgs {
+    #[primary_span]
+    pub spans: Vec<Span>,
+}
+
+#[derive(Diagnostic)]
+#[diag("`...` and `#[splat]` are not allowed in the same function")]
+#[help("remove `#[splat]` or remove `...`")]
+pub(crate) struct CVarArgsAndSplat {
+    #[primary_span]
+    pub spans: Vec<Span>,
+}
+
 #[derive(Diagnostic)]
 #[diag("documentation comments cannot be applied to function parameters")]
 pub(crate) struct FnParamDocComment {
@@ -132,6 +159,7 @@ pub(crate) struct FnParamDocComment {
     pub span: Span,
 }
 
+// FIXME(splat): add splat to the allowed built-in attributes when it is complete/stabilized
 #[derive(Diagnostic)]
 #[diag(
     "allow, cfg, cfg_attr, deny, expect, forbid, and warn are the only allowed built-in attributes in function parameters"

compiler/rustc_ast_passes/src/feature_gate.rs

@@ -506,6 +506,7 @@ pub fn check_crate(krate: &ast::Crate, sess: &Session, features: &Features) {
     gate_all!(pin_ergonomics, "pinned reference syntax is experimental");
     gate_all!(postfix_match, "postfix match is experimental");
     gate_all!(return_type_notation, "return type notation is experimental");
+    gate_all!(splat, "`fn(#[splat] (a, ...))` is incomplete", "call as func((a, ...)) instead");
     gate_all!(super_let, "`super let` is experimental");
     gate_all!(try_blocks_heterogeneous, "`try bikeshed` expression is experimental");
     gate_all!(unsafe_binders, "unsafe binder types are experimental");

compiler/rustc_attr_parsing/src/attributes/mod.rs

@@ -64,6 +64,7 @@ pub(crate) mod rustc_allocator;
 pub(crate) mod rustc_dump;
 pub(crate) mod rustc_internal;
 pub(crate) mod semantics;
+pub(crate) mod splat;
 pub(crate) mod stability;
 pub(crate) mod test_attrs;
 pub(crate) mod traits;

compiler/rustc_attr_parsing/src/attributes/splat.rs

@@ -0,0 +1,17 @@
+//! Attribute parsing for the `#[splat]` function argument overloading attribute.
+//! This attribute modifies typecheck to support overload resolution, then modifies codegen for performance.
+
+use super::prelude::*;
+
+pub(crate) struct SplatParser;
+
+impl<S: Stage> NoArgsAttributeParser<S> for SplatParser {
+    const PATH: &[Symbol] = &[sym::splat];
+    const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::Warn;
+    const ALLOWED_TARGETS: AllowedTargets = AllowedTargets::AllowList(&[
+        Allow(Target::Param),
+        // FIXME(splat): only allow MacroCall if the macro creates an argument
+        Allow(Target::MacroCall),
+    ]);
+    const CREATE: fn(Span) -> AttributeKind = AttributeKind::Splat;
+}

compiler/rustc_attr_parsing/src/context.rs

@@ -56,6 +56,7 @@ use crate::attributes::rustc_allocator::*;
 use crate::attributes::rustc_dump::*;
 use crate::attributes::rustc_internal::*;
 use crate::attributes::semantics::*;
+use crate::attributes::splat::*;
 use crate::attributes::stability::*;
 use crate::attributes::test_attrs::*;
 use crate::attributes::traits::*;
@@ -338,6 +339,7 @@ attribute_parsers!(
         Single<WithoutArgs<RustcStrictCoherenceParser>>,
         Single<WithoutArgs<RustcTrivialFieldReadsParser>>,
         Single<WithoutArgs<RustcUnsafeSpecializationMarkerParser>>,
+        Single<WithoutArgs<SplatParser>>,
         Single<WithoutArgs<ThreadLocalParser>>,
         Single<WithoutArgs<TrackCallerParser>>,
         // tidy-alphabetical-end