Impl typeck & MIR lowering for argument splatting

Author: teor2345

compiler/rustc_ast/src/ast.rs

@@ -3055,9 +3055,30 @@ 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` and `FnDeclFlags::NO_SPLATTED_ARG_INDEX`.
+    pub const NO_SPLATTED_ARG_INDEX: u16 = u16::MAX;
+
+    /// Returns a splatted argument index, if any are present.
+    pub fn splatted(&self) -> Option<u16> {
+        self.inputs.iter().enumerate().find_map(|(index, arg)| {
+            if index == Self::NO_SPLATTED_ARG_INDEX as usize {
+                // AST validation has already checked the splatted argument index is valid, so just
+                // ignore invalid indexes here.
+                None
+            } else {
+                arg.attrs
+                    .iter()
+                    .any(|attr| attr.has_name(sym::splat))
+                    .then_some(u16::try_from(index).unwrap())
+            }
+        })
+    }
 }
 
 /// Is the trait definition an auto trait?

compiler/rustc_ast_lowering/src/delegation.rs

@@ -141,7 +141,7 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
 
                 let is_method = self.is_method(sig_id, span);
 
-                let (param_count, c_variadic) = self.param_count(sig_id);
+                let (param_count, c_variadic, splatted) = self.param_count(sig_id);
 
                 let mut generics = self.uplift_delegation_generics(delegation, sig_id, item_id);
 
@@ -153,8 +153,14 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
                     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,
+                    splatted,
+                    span,
+                    &generics,
+                );
 
                 let sig = self.lower_delegation_sig(sig_id, decl, span);
                 let ident = self.lower_ident(delegation.ident);
@@ -268,17 +274,18 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
         self.resolver.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*/, Option<u16> /*splatted*/) {
         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())
     }
 
     fn lower_delegation_decl(
         &mut self,
         sig_id: DefId,
         param_count: usize,
         c_variadic: bool,
+        splatted: Option<u16>,
         span: Span,
         generics: &GenericsGenerationResults<'hir>,
     ) -> &'hir hir::FnDecl<'hir> {
@@ -311,7 +318,9 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
             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_splatted(splatted, inputs.len())
+                .unwrap(),
         })
     }
 

compiler/rustc_ast_lowering/src/lib.rs

@@ -1844,12 +1844,15 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
         coro: Option<CoroutineKind>,
     ) -> &'hir hir::FnDecl<'hir> {
         let c_variadic = decl.c_variadic();
+        let mut splatted = 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 = None;
             inputs = &inputs[..inputs.len() - 1];
         }
         let inputs = self.arena.alloc_from_iter(inputs.iter().map(|param| {
@@ -1937,7 +1940,9 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
                 }
             }))
             .set_lifetime_elision_allowed(self.resolver.lifetime_elision_allowed(fn_node_id))
-            .set_c_variadic(c_variadic);
+            .set_c_variadic(c_variadic)
+            .set_splatted(splatted, inputs.len())
+            .unwrap();
 
         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) {

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 {

compiler/rustc_borrowck/src/diagnostics/region_errors.rs

@@ -15,8 +15,7 @@ use rustc_middle::bug;
 use rustc_middle::hir::place::PlaceBase;
 use rustc_middle::mir::{AnnotationSource, ConstraintCategory, ReturnConstraint};
 use rustc_middle::ty::{
-    self, FnSigKind, GenericArgs, Region, RegionVid, Ty, TyCtxt, TypeFoldable, TypeVisitor,
-    fold_regions,
+    self, GenericArgs, Region, RegionVid, Ty, TyCtxt, TypeFoldable, TypeVisitor, fold_regions,
 };
 use rustc_span::{Ident, Span, kw};
 use rustc_trait_selection::error_reporting::InferCtxtErrorExt;
@@ -1085,8 +1084,8 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
         }
 
         // Build a new closure where the return type is an owned value, instead of a ref.
-        let fn_sig_kind =
-            FnSigKind::default().set_safe(true).set_c_variadic(liberated_sig.c_variadic());
+        // The new closure is safe, but otherwise has the same ABI, splat, and c-variadic.
+        let fn_sig_kind = liberated_sig.fn_sig_kind.set_safe(true);
         let closure_sig_as_fn_ptr_ty = Ty::new_fn_ptr(
             tcx,
             ty::Binder::dummy(tcx.mk_fn_sig(

compiler/rustc_codegen_cranelift/src/abi/mod.rs

@@ -265,6 +265,7 @@ pub(crate) fn codegen_fn_prelude<'tcx>(fx: &mut FunctionCx<'_, '_, 'tcx>, start_
         .map(|local| {
             let arg_ty = fx.monomorphize(fx.mir.local_decls[local].ty);
 
+            // FIXME(splat): un-tuple splatted arguments in codegen, for performance
             // Adapted from https://github.com/rust-lang/rust/blob/145155dc96757002c7b2e9de8489416e2fdbbd57/src/librustc_codegen_llvm/mir/mod.rs#L442-L482
             if Some(local) == fx.mir.spread_arg {
                 // This argument (e.g. the last argument in the "rust-call" ABI)

compiler/rustc_codegen_ssa/src/debuginfo/type_names.rs

@@ -375,6 +375,7 @@ fn push_debuginfo_type_name<'tcx>(
                 output.push_str("fn(");
             }
 
+            // FIXME(splat): should debuginfo be de-tupled in the callee (and caller)?
             if !sig.inputs().is_empty() {
                 for &parameter_type in sig.inputs() {
                     push_debuginfo_type_name(tcx, parameter_type, true, output, visited);

compiler/rustc_const_eval/src/const_eval/type_info.rs

@@ -7,7 +7,7 @@ use rustc_ast::Mutability;
 use rustc_hir::LangItem;
 use rustc_middle::span_bug;
 use rustc_middle::ty::layout::TyAndLayout;
-use rustc_middle::ty::{self, Const, FnHeader, FnSigTys, ScalarInt, Ty, TyCtxt};
+use rustc_middle::ty::{self, Const, FnHeader, FnSigKind, FnSigTys, ScalarInt, Ty, TyCtxt};
 use rustc_span::{Symbol, sym};
 
 use crate::const_eval::CompileTimeMachine;
@@ -465,6 +465,15 @@ impl<'tcx> InterpCx<'tcx, CompileTimeMachine<'tcx>> {
                 sym::variadic => {
                     self.write_scalar(Scalar::from_bool(fn_sig_kind.c_variadic()), &field_place)?;
                 }
+                sym::splat => {
+                    self.write_scalar(
+                        // Use the same encoding as FnSigKind.splatted
+                        Scalar::from_u16(
+                            fn_sig_kind.splatted().unwrap_or(FnSigKind::NO_SPLATTED_ARG_INDEX),
+                        ),
+                        &field_place,
+                    )?;
+                }
                 other => span_bug!(self.tcx.def_span(field.did), "unimplemented field {other}"),
             }
         }