Auto merge of #157954 - teor2345:fn-arg-splat-perf-fndef-bitpack, r=<try>

perf: Store 1 bit for splatting in FnDeclFlags

Author: bors

compiler/rustc_ast/src/ast.rs

@@ -3058,23 +3058,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()
     }
 }
 

compiler/rustc_ast_lowering/src/delegation.rs

@@ -93,8 +93,9 @@ struct ParamInfo {
     /// Whether the function arguments end in a C variadic `...` parameter.
     pub c_variadic: bool,
 
-    /// The index of the splatted parameter, if any.
-    pub splatted: Option<u16>,
+    /// Does the function have a splatted parameter?
+    /// The index is available from the attributes.
+    pub has_splatted_arg: bool,
 }
 
 const PARENT_ID: hir::ItemLocalId = hir::ItemLocalId::ZERO;
@@ -384,15 +385,15 @@ impl<'hir> LoweringContext<'_, 'hir> {
         self.get_partial_res(node_id).and_then(|r| r.expect_full_res().opt_def_id())
     }
 
-    /// Returns function parameter info, including C variadic `...` and `#[splat]` if present.
+    /// Returns function parameter info, including C variadic `...` and `#[splat]` flag if present.
     fn param_info(&self, def_id: DefId) -> ParamInfo {
         let sig = self.tcx.fn_sig(def_id).skip_binder().skip_binder();
 
-        // FIXME(splat): use `sig.splatted()` once FnSig has it
+        // FIXME(splat): use `sig.splatted().is_some()` once FnSig has it
         ParamInfo {
             param_count: sig.inputs().len() + usize::from(sig.c_variadic()),
             c_variadic: sig.c_variadic(),
-            splatted: None,
+            has_splatted_arg: false,
         }
     }
 
@@ -407,7 +408,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         call_expr_id: HirId,
         unused_target_expr: bool,
     ) -> &'hir hir::FnDecl<'hir> {
-        let ParamInfo { param_count, c_variadic, splatted } = param_info;
+        let ParamInfo { param_count, c_variadic, has_splatted_arg } = param_info;
 
         // The last parameter in C variadic functions is skipped in the signature,
         // like during regular lowering.
@@ -454,8 +455,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             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),
         })
     }
 

compiler/rustc_ast_lowering/src/lib.rs

@@ -1776,15 +1776,15 @@ impl<'hir> LoweringContext<'_, 'hir> {
         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| {
@@ -1875,8 +1875,17 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 self.owner.id == fn_node_id && self.owner.lifetime_elision_allowed,
             )
             .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 })
     }

compiler/rustc_hir/src/hir.rs

@@ -36,6 +36,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;
@@ -4034,7 +4035,8 @@ pub enum SplattedArgIndexError {
 #[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`.
+    /// 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?
@@ -4058,8 +4060,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()
@@ -4076,20 +4078,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.
@@ -4132,38 +4133,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
     }
 
@@ -4189,9 +4167,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");
     }
 }
 
@@ -4240,8 +4247,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 {

compiler/rustc_hir_analysis/src/hir_ty_lowering/mod.rs

@@ -3501,7 +3501,9 @@ 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
+        // FIXME(splat): use ```
+        // set_splatted(decl.splatted_arg_index(&move |id| HasAttrs::get_attrs(id, &tcx)), input_tys.len())
+        // ``` once FnSig has it
         let fn_sig_kind = FnSigKind::default()
             .set_abi(abi)
             .set_safety(safety)

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};
@@ -2265,7 +2264,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() {