Store 1 bit for splatting in FnDeclFlags

teor2345 · teor2345 · commit 7bb820bca1c3 · 2026-04-28T11:21:22.000+10:00
diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
@@ -3061,23 +3061,33 @@ impl FnDecl {
     }
 
     /// The marker index for "no splatted arguments".
-    /// Must have the same value as `FnSigKind::NO_SPLATTED_ARG_INDEX` and `FnDeclFlags::NO_SPLATTED_ARG_INDEX`.
+    /// 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, 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())
-            }
-        })
+    /// 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()
     }
 }
 
diff --git a/compiler/rustc_ast_lowering/src/delegation.rs b/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, splatted) = 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);
 
@@ -157,7 +157,7 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
                     sig_id,
                     param_count,
                     c_variadic,
-                    splatted,
+                    has_splatted_arg,
                     span,
                     &generics,
                 );
@@ -275,17 +275,24 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
     }
 
     // Function parameter count, including C variadic `...` and `#[splat]` if present.
-    fn param_count(&self, def_id: DefId) -> (usize, bool /*c_variadic*/, Option<u16> /*splatted*/) {
+    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.splatted())
+        (
+            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,
-        splatted: Option<u16>,
+        has_splatted_arg: bool,
         span: Span,
         generics: &GenericsGenerationResults<'hir>,
     ) -> &'hir hir::FnDecl<'hir> {
@@ -319,8 +326,7 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
             fn_decl_kind: FnDeclFlags::default()
                 .set_lifetime_elision_allowed(true)
                 .set_c_variadic(c_variadic)
-                .set_splatted(splatted, inputs.len())
-                .unwrap(),
+                .set_has_splatted_arg(has_splatted_arg),
         })
     }
 
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
@@ -1844,15 +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();
+        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 = None;
+            splatted_arg_index = None;
             inputs = &inputs[..inputs.len() - 1];
         }
         let inputs = self.arena.alloc_from_iter(inputs.iter().map(|param| {
@@ -1941,8 +1941,17 @@ 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_splatted(splatted, inputs.len())
-            .unwrap();
+            .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 })
     }
diff --git a/compiler/rustc_hir/src/hir.rs b/compiler/rustc_hir/src/hir.rs
@@ -34,6 +34,7 @@ use tracing::debug;
 use crate::attrs::AttributeKind;
 use crate::def::{CtorKind, DefKind, MacroKinds, PerNS, Res};
 use crate::def_id::{DefId, LocalDefIdMap};
+use crate::find_attr;
 pub(crate) use crate::hir_id::{HirId, ItemLocalId, ItemLocalMap, OwnerId};
 use crate::intravisit::{FnKind, VisitorExt};
 use crate::lints::DelayedLints;
@@ -3937,7 +3938,8 @@ pub enum SplattedArgIndexError {
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 #[derive(Encodable, Decodable, HashStable_Generic)]
 pub struct FnDeclFlags {
-    /// Holds the c_variadic and lifetime_elision_allowed bitflags, and 3 bits for the `ImplicitSelfKind`.
+    /// Holds the c_variadic, lifetime_elision_allowed, and has_splatted_arg bitflags, and 3 bits
+    /// for the `ImplicitSelfKind`.
     flags: u8,
 
     /// Which function argument is splatted into multiple arguments in callers, if any?
@@ -3961,8 +3963,8 @@ impl fmt::Debug for FnDeclFlags {
             f.field(&"CVariadic");
         }
 
-        if let Some(index) = self.splatted() {
-            f.field(&format!("Splatted({})", index));
+        if self.has_splatted_arg() {
+            f.field(&"HasSplattedArg");
         }
 
         f.finish()
@@ -3979,20 +3981,19 @@ impl FnDeclFlags {
     /// Bitflag for lifetime elision.
     const LIFETIME_ELISION_ALLOWED_FLAG: u8 = 1 << 4;
 
-    /// 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;
+    /// Bitflag set if any argument is splatted (for performance).
+    const HAS_SPLATTED_ARG_FLAG: u8 = 1 << 5;
 
     /// Create a new FnDeclKind with no implicit self, no lifetime elision, no C-style variadic
-    /// argument, and no splatting.
+    /// argument, and no splatted argument.
     /// 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, splatted: 0 }
             .set_implicit_self(ImplicitSelfKind::None)
             .set_lifetime_elision_allowed(false)
             .set_c_variadic(false)
-            .set_no_splatted_args()
+            .set_has_splatted_arg(false)
     }
 
     /// Set the implicit self kind.
@@ -4035,38 +4036,15 @@ impl FnDeclFlags {
         self
     }
 
-    /// Set the splatted argument index.
-    /// The number of function arguments is used for error checking.
+    /// Set the splatted argument flag.
     #[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;
+    pub const fn set_has_splatted_arg(mut self, has_splatted_arg: bool) -> Self {
+        if has_splatted_arg {
+            self.flags |= Self::HAS_SPLATTED_ARG_FLAG;
         } else {
-            self.splatted = Self::NO_SPLATTED_ARG_INDEX;
+            self.flags &= !Self::HAS_SPLATTED_ARG_FLAG;
         }
 
-        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
     }
 
@@ -4092,9 +4070,38 @@ impl FnDeclFlags {
         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) }
+    /// Does this function have a splatted argument?
+    pub const fn has_splatted_arg(self) -> bool {
+        self.flags & Self::HAS_SPLATTED_ARG_FLAG != 0
+    }
+
+    /// Returns `true` if the given input contains a `#[splat]` attribute in `attrs`.
+    pub fn is_splatted_arg<'hir>(
+        &self,
+        input: &'hir Ty<'hir>,
+        attrs: &'hir dyn Fn(HirId) -> &'hir [Attribute],
+    ) -> bool {
+        self.has_splatted_arg() && find_attr!(attrs(input.hir_id), Splat(_))
+    }
+
+    /// Searches `inputs` and `attrs` for the index of the splatted argument. Returns `None` if
+    /// there is no splatted argument.
+    pub fn splatted_arg_index<'hir>(
+        &self,
+        inputs: &'hir [Ty<'hir>],
+        attrs: &'hir dyn Fn(HirId) -> &'hir [Attribute],
+    ) -> Option<u16> {
+        if !self.has_splatted_arg() {
+            return None;
+        }
+
+        for (index, input) in inputs.iter().enumerate() {
+            if self.is_splatted_arg(input, attrs) {
+                return Some(u16::try_from(index).unwrap());
+            }
+        }
+
+        unreachable!("no splatted argument found");
     }
 }
 
@@ -4143,8 +4150,27 @@ impl<'hir> FnDecl<'hir> {
         self.fn_decl_kind.lifetime_elision_allowed()
     }
 
-    pub fn splatted(&self) -> Option<u16> {
-        self.fn_decl_kind.splatted()
+    /// Returns `true` if the function has a splatted argument.
+    pub fn has_splatted_arg(&self) -> bool {
+        self.fn_decl_kind.has_splatted_arg()
+    }
+
+    /// Returns `true` if the given argument `index` contains a `#[splat]` attribute in `attrs`.
+    pub fn is_splatted_arg(
+        &self,
+        index: usize,
+        attrs: &'hir dyn Fn(HirId) -> &'hir [Attribute],
+    ) -> bool {
+        self.fn_decl_kind.is_splatted_arg(&self.inputs[index], attrs)
+    }
+
+    /// Searches `self.inputs` and `attrs` for the index of the splatted argument. Returns `None`
+    /// if there is no splatted argument.
+    pub fn splatted_arg_index(
+        &self,
+        attrs: &'hir dyn Fn(HirId) -> &'hir [Attribute],
+    ) -> Option<u16> {
+        self.fn_decl_kind.splatted_arg_index(self.inputs, attrs)
     }
 
     pub fn dummy(span: Span) -> Self {
diff --git a/compiler/rustc_hir_analysis/src/hir_ty_lowering/mod.rs b/compiler/rustc_hir_analysis/src/hir_ty_lowering/mod.rs
@@ -29,6 +29,7 @@ use rustc_errors::{
     Applicability, Diag, DiagCtxtHandle, Diagnostic, ErrorGuaranteed, FatalError, Level, StashKey,
     struct_span_code_err,
 };
+use rustc_hir::attrs::HasAttrs;
 use rustc_hir::def::{CtorKind, CtorOf, DefKind, Res};
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_hir::{self as hir, AnonConst, GenericArg, GenericArgs, HirId};
@@ -3593,7 +3594,10 @@ impl<'tcx> dyn HirTyLowerer<'tcx> + '_ {
             .set_abi(abi)
             .set_safe(safety.is_safe())
             .set_c_variadic(decl.fn_decl_kind.c_variadic())
-            .set_splatted(decl.fn_decl_kind.splatted(), input_tys.len())
+            .set_splatted(
+                decl.splatted_arg_index(&move |id| HasAttrs::get_attrs(id, &tcx)),
+                input_tys.len(),
+            )
             .unwrap();
         let fn_ty = tcx.mk_fn_sig(input_tys, output_ty, fn_sig_kind);
         let fn_ptr_ty = ty::Binder::bind_with_vars(fn_ty, bound_vars);
diff --git a/compiler/rustc_hir_pretty/src/lib.rs b/compiler/rustc_hir_pretty/src/lib.rs
@@ -16,12 +16,11 @@ use rustc_ast_pretty::pp::Breaks::{Consistent, Inconsistent};
 use rustc_ast_pretty::pp::{self, BoxMarker, Breaks};
 use rustc_ast_pretty::pprust::state::MacHeader;
 use rustc_ast_pretty::pprust::{Comments, PrintState};
-use rustc_hir as hir;
 use rustc_hir::attrs::{AttributeKind, PrintAttribute};
 use rustc_hir::{
-    BindingMode, ByRef, ConstArg, ConstArgExprField, ConstArgKind, GenericArg, GenericBound,
-    GenericParam, GenericParamKind, HirId, ImplicitSelfKind, LifetimeParamKind, Node, PatKind,
-    PreciseCapturingArg, RangeEnd, Term, TyFieldPath, TyPatKind,
+    self as hir, BindingMode, ByRef, ConstArg, ConstArgExprField, ConstArgKind, GenericArg,
+    GenericBound, GenericParam, GenericParamKind, HirId, ImplicitSelfKind, LifetimeParamKind, Node,
+    PatKind, PreciseCapturingArg, RangeEnd, Term, TyFieldPath, TyPatKind, find_attr,
 };
 use rustc_span::source_map::SourceMap;
 use rustc_span::{DUMMY_SP, FileName, Ident, Span, Spanned, Symbol, kw, sym};
@@ -2264,7 +2263,7 @@ impl<'a> State<'a> {
         assert!(arg_idents.is_empty() || body_id.is_none());
         let mut i = 0;
         let mut print_arg = |s: &mut Self, ty: Option<&hir::Ty<'_>>| {
-            if Some(i) == decl.splatted().map(usize::from) {
+            if decl.has_splatted_arg() && find_attr!(s.attrs(decl.inputs[i].hir_id), Splat(_)) {
                 s.word("#[splat]");
             }
             if i == 0 && decl.implicit_self().has_implicit_self() {
diff --git a/compiler/rustc_type_ir/src/ty_kind.rs b/compiler/rustc_type_ir/src/ty_kind.rs
@@ -858,9 +858,11 @@ impl FnSigKind {
     const C_VARIADIC_FLAG: u8 = 1 << 7;
 
     /// The marker index for "no splatted arguments".
-    /// Must have the same value as `FnDeclFlags::NO_SPLATTED_ARG_INDEX` and `rustc_ast::FnDecl::NO_SPLATTED_ARG_INDEX`.
+    /// Must have the same value as `rustc_ast::FnDecl::NO_SPLATTED_ARG_INDEX`.
     ///
     /// This is an implementation detail, which should only be used in low-level encoding.
+    // 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;
 
     /// Create a new FnSigKind with the "Rust" ABI, "Unsafe" safety, and no C-style variadic or splatted arguments.
