Add HIR FnDecl for #[splat]

Author: teor2345

compiler/rustc_ast/src/ast.rs

@@ -3052,9 +3052,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

@@ -199,7 +199,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, splatted) = self.param_count(sig_id);
 
         if !self.check_block_soundness(delegation, sig_id, is_method, param_count) {
             return self.generate_delegation_error(span, delegation);
@@ -215,6 +215,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             sig_id,
             param_count,
             c_variadic,
+            splatted,
             span,
             &generics,
             delegation.id,
@@ -367,10 +368,11 @@ 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*/, 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())
+        // FIXME(splat): use `sig.splatted()` once FnSig has it
+        (sig.inputs().len() + usize::from(sig.c_variadic()), sig.c_variadic(), None)
     }
 
     fn lower_delegation_decl(
@@ -379,6 +381,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         sig_id: DefId,
         param_count: usize,
         c_variadic: bool,
+        splatted: Option<u16>,
         span: Span,
         generics: &GenericsGenerationResults<'hir>,
         call_path_node_id: NodeId,
@@ -429,7 +432,9 @@ 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_splatted(splatted, inputs.len())
+                .unwrap(),
         })
     }
 

compiler/rustc_ast_lowering/src/lib.rs

@@ -1776,12 +1776,15 @@ impl<'hir> LoweringContext<'_, 'hir> {
         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| {
@@ -1871,7 +1874,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
             .set_lifetime_elision_allowed(
                 self.owner.id == fn_node_id && self.owner.lifetime_elision_allowed,
             )
-            .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

@@ -412,6 +412,18 @@ impl<'a> AstValidator<'a> {
             })
             .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(diagnostics::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()

compiler/rustc_ast_passes/src/diagnostics.rs

@@ -123,6 +123,17 @@ 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")]

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;
@@ -1094,9 +1093,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_safety(hir::Safety::Safe)
-            .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_safety(hir::Safety::Safe);
         let closure_sig_as_fn_ptr_ty = Ty::new_fn_ptr(
             tcx,
             ty::Binder::dummy(tcx.mk_fn_sig(

compiler/rustc_hir/src/hir.rs

@@ -4018,13 +4018,29 @@ pub struct Param<'hir> {
     pub span: Span,
 }
 
+/// Error type for splatted argument index errors.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum SplattedArgIndexError {
+    /// The splatted argument index is invalid.
+    /// Currently the only unsupported index is `u16::MAX`, which is used to indicate that no argument
+    /// is splatted.
+    InvalidIndex { splatted_arg_index: u16 },
+
+    /// The splatted argument index is outside the bounds of the function arguments.
+    OutOfBounds { splatted_arg_index: u16, args_len: u16 },
+}
+
 /// Contains the packed non-type fields of a function declaration.
-// FIXME(splat): add the splatted argument index as a u16
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 #[derive(Encodable, Decodable, StableHash)]
 pub struct FnDeclFlags {
     /// Holds the c_variadic and lifetime_elision_allowed bitflags, and 3 bits for the `ImplicitSelfKind`.
     flags: u8,
+
+    /// Which function argument is splatted into multiple arguments in callers, if any?
+    /// Splatting functions with `u16::MAX` arguments is not supported, see `FnSigKind` for
+    /// details.
+    splatted: u16,
 }
 
 impl fmt::Debug for FnDeclFlags {
@@ -4036,11 +4052,15 @@ impl fmt::Debug for FnDeclFlags {
             f.field(&"LifetimeElisionAllowed");
         } else {
             f.field(&"NoLifetimeElision");
-        };
+        }
 
         if self.c_variadic() {
             f.field(&"CVariadic");
-        };
+        }
+
+        if let Some(index) = self.splatted() {
+            f.field(&format!("Splatted({})", index));
+        }
 
         f.finish()
     }
@@ -4056,14 +4076,20 @@ impl FnDeclFlags {
     /// Bitflag for lifetime elision.
     const LIFETIME_ELISION_ALLOWED_FLAG: u8 = 1 << 4;
 
-    /// Create a new FnDeclKind with no implicit self, no lifetime elision, and no C-style variadic argument.
+    /// Marker index for "no splatted argument".
+    /// Must have the same value as `FnSigKind::NO_SPLATTED_ARG_INDEX` and `rustc_ast::FnDecl::NO_SPLATTED_ARG_INDEX`.
+    const NO_SPLATTED_ARG_INDEX: u16 = u16::MAX;
+
+    /// Create a new FnDeclKind with no implicit self, no lifetime elision, no C-style variadic
+    /// argument, and no splatting.
     /// To modify these flags, use the `set_*` methods, for readability.
     // FIXME: use Default instead when that trait is const stable.
     pub const fn default() -> Self {
-        Self { flags: 0 }
+        Self { flags: 0, splatted: 0 }
             .set_implicit_self(ImplicitSelfKind::None)
             .set_lifetime_elision_allowed(false)
             .set_c_variadic(false)
+            .set_no_splatted_args()
     }
 
     /// Set the implicit self kind.
@@ -4106,6 +4132,41 @@ impl FnDeclFlags {
         self
     }
 
+    /// Set the splatted argument index.
+    /// The number of function arguments is used for error checking.
+    #[must_use = "this method does not modify the receiver"]
+    pub const fn set_splatted(
+        mut self,
+        splatted: Option<u16>,
+        args_len: usize,
+    ) -> Result<Self, SplattedArgIndexError> {
+        if let Some(splatted_arg_index) = splatted {
+            if splatted_arg_index == Self::NO_SPLATTED_ARG_INDEX {
+                // This index value is used as a marker for "no splatting", so it is unsupported.
+                return Err(SplattedArgIndexError::InvalidIndex { splatted_arg_index });
+            } else if splatted_arg_index as usize >= args_len {
+                return Err(SplattedArgIndexError::OutOfBounds {
+                    splatted_arg_index,
+                    args_len: args_len as u16,
+                });
+            }
+
+            self.splatted = splatted_arg_index;
+        } else {
+            self.splatted = Self::NO_SPLATTED_ARG_INDEX;
+        }
+
+        Ok(self)
+    }
+
+    /// Set "no splatted arguments" for the function declaration.
+    #[must_use = "this method does not modify the receiver"]
+    pub const fn set_no_splatted_args(mut self) -> Self {
+        self.splatted = Self::NO_SPLATTED_ARG_INDEX;
+
+        self
+    }
+
     /// Get the implicit self kind.
     pub const fn implicit_self(self) -> ImplicitSelfKind {
         match self.flags & Self::IMPLICIT_SELF_MASK {
@@ -4127,6 +4188,11 @@ impl FnDeclFlags {
     pub const fn lifetime_elision_allowed(self) -> bool {
         self.flags & Self::LIFETIME_ELISION_ALLOWED_FLAG != 0
     }
+
+    /// Get the splatted argument index, if any.
+    pub const fn splatted(self) -> Option<u16> {
+        if self.splatted == Self::NO_SPLATTED_ARG_INDEX { None } else { Some(self.splatted) }
+    }
 }
 
 /// Represents the header (not the body) of a function declaration.
@@ -4174,6 +4240,10 @@ impl<'hir> FnDecl<'hir> {
         self.fn_decl_kind.lifetime_elision_allowed()
     }
 
+    pub fn splatted(&self) -> Option<u16> {
+        self.fn_decl_kind.splatted()
+    }
+
     pub fn dummy(span: Span) -> Self {
         Self {
             inputs: &[],

compiler/rustc_hir_analysis/src/check/mod.rs

@@ -450,7 +450,7 @@ fn fn_sig_suggestion<'tcx>(
         .iter()
         .enumerate()
         .map(|(i, ty)| {
-            Some(match ty.kind() {
+            let arg_ty = match ty.kind() {
                 ty::Param(_) if assoc.is_method() && i == 0 => "self".to_string(),
                 ty::Ref(reg, ref_ty, mutability) if i == 0 => {
                     let reg = format!("{reg} ");
@@ -477,7 +477,8 @@ fn fn_sig_suggestion<'tcx>(
                         format!("_: {ty}")
                     }
                 }
-            })
+            };
+            Some(format!("{arg_ty}"))
         })
         .chain(std::iter::once(if sig.c_variadic() { Some("...".to_string()) } else { None }))
         .flatten()

compiler/rustc_hir_analysis/src/check/wfcheck.rs

@@ -1711,8 +1711,9 @@ fn check_fn_or_method<'tcx>(
 
     if sig.abi() == ExternAbi::RustCall {
         let span = tcx.def_span(def_id);
-        let has_implicit_self = hir_decl.implicit_self() != hir::ImplicitSelfKind::None;
+        let has_implicit_self = hir_decl.implicit_self().has_implicit_self();
         let mut inputs = sig.inputs().iter().skip(if has_implicit_self { 1 } else { 0 });
+        // FIXME(splat): use `sig.splatted()` once FnSig has it
         // Check that the argument is a tuple and is sized
         if let Some(ty) = inputs.next() {
             wfcx.register_bound(

compiler/rustc_hir_analysis/src/hir_ty_lowering/mod.rs

@@ -3497,6 +3497,8 @@ impl<'tcx> dyn HirTyLowerer<'tcx> + '_ {
 
         debug!(?output_ty);
 
+        debug!(?abi, ?safety, ?decl.fn_decl_kind, input_tys_len = ?input_tys.len());
+        // FIXME(splat): use `set_splatted()` once FnSig has it
         let fn_sig_kind = FnSigKind::default()
             .set_abi(abi)
             .set_safety(safety)