attr: parse rustc_scalable_vector(N)

Extend parsing of `ReprOptions` with `rustc_scalable_vector(N)` which
optionally accepts a single literal integral value - the base multiple of
lanes that are in a scalable vector. Can only be applied to structs.

Co-authored-by: Jamie Cunliffe <Jamie.Cunliffe@arm.com>

Author: davidtwco

compiler/rustc_abi/src/lib.rs

@@ -96,9 +96,11 @@ bitflags! {
         /// If true, the type is always passed indirectly by non-Rustic ABIs.
         /// See [`TyAndLayout::pass_indirectly_in_non_rustic_abis`] for details.
         const PASS_INDIRECTLY_IN_NON_RUSTIC_ABIS = 1 << 5;
-        /// Any of these flags being set prevent field reordering optimisation.
-        const FIELD_ORDER_UNOPTIMIZABLE   = ReprFlags::IS_C.bits()
+        const IS_SCALABLE        = 1 << 6;
+         // Any of these flags being set prevent field reordering optimisation.
+        const FIELD_ORDER_UNOPTIMIZABLE = ReprFlags::IS_C.bits()
                                  | ReprFlags::IS_SIMD.bits()
+                                 | ReprFlags::IS_SCALABLE.bits()
                                  | ReprFlags::IS_LINEAR.bits();
         const ABI_UNOPTIMIZABLE = ReprFlags::IS_C.bits() | ReprFlags::IS_SIMD.bits();
     }
@@ -135,6 +137,19 @@ impl IntegerType {
     }
 }
 
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[cfg_attr(
+    feature = "nightly",
+    derive(Encodable_NoContext, Decodable_NoContext, HashStable_Generic)
+)]
+pub enum ScalableElt {
+    /// `N` in `rustc_scalable_vector(N)` - the element count of the scalable vector
+    ElementCount(u16),
+    /// `rustc_scalable_vector` w/out `N`, used for tuple types of scalable vectors that only
+    /// contain other scalable vectors
+    Container,
+}
+
 /// Represents the repr options provided by the user.
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
 #[cfg_attr(
@@ -146,6 +161,8 @@ pub struct ReprOptions {
     pub align: Option<Align>,
     pub pack: Option<Align>,
     pub flags: ReprFlags,
+    /// `#[rustc_scalable_vector]`
+    pub scalable: Option<ScalableElt>,
     /// The seed to be used for randomizing a type's layout
     ///
     /// Note: This could technically be a `u128` which would
@@ -162,6 +179,11 @@ impl ReprOptions {
         self.flags.contains(ReprFlags::IS_SIMD)
     }
 
+    #[inline]
+    pub fn scalable(&self) -> bool {
+        self.flags.contains(ReprFlags::IS_SCALABLE)
+    }
+
     #[inline]
     pub fn c(&self) -> bool {
         self.flags.contains(ReprFlags::IS_C)

compiler/rustc_attr_parsing/messages.ftl

@@ -189,6 +189,9 @@ attr_parsing_rustc_allowed_unstable_pairing =
 attr_parsing_rustc_promotable_pairing =
     `rustc_promotable` attribute must be paired with either a `rustc_const_unstable` or a `rustc_const_stable` attribute
 
+attr_parsing_rustc_scalable_vector_count_out_of_range = element count in `rustc_scalable_vector` is too large: `{$n}`
+    .note = the value may not exceed `u16::MAX`
+
 attr_parsing_soft_no_args =
     `soft` should not have any arguments
 

compiler/rustc_attr_parsing/src/attributes/rustc_internal.rs

@@ -1,5 +1,6 @@
 use super::prelude::*;
 use super::util::parse_single_integer;
+use crate::session_diagnostics::RustcScalableVectorCountOutOfRange;
 
 pub(crate) struct RustcMainParser;
 
@@ -76,3 +77,29 @@ impl<S: Stage> SingleAttributeParser<S> for RustcSimdMonomorphizeLaneLimitParser
         Some(AttributeKind::RustcSimdMonomorphizeLaneLimit(cx.parse_limit_int(nv)?))
     }
 }
+
+pub(crate) struct RustcScalableVectorParser;
+
+impl<S: Stage> SingleAttributeParser<S> for RustcScalableVectorParser {
+    const PATH: &[rustc_span::Symbol] = &[sym::rustc_scalable_vector];
+    const ATTRIBUTE_ORDER: AttributeOrder = AttributeOrder::KeepInnermost;
+    const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::Error;
+    const ALLOWED_TARGETS: AllowedTargets = AllowedTargets::AllowList(&[Allow(Target::Struct)]);
+    const TEMPLATE: AttributeTemplate = template!(Word, List: &["count"]);
+
+    fn convert(cx: &mut AcceptContext<'_, '_, S>, args: &ArgParser<'_>) -> Option<AttributeKind> {
+        if args.no_args().is_ok() {
+            return Some(AttributeKind::RustcScalableVector {
+                element_count: None,
+                span: cx.attr_span,
+            });
+        }
+
+        let n = parse_single_integer(cx, args)?;
+        let Ok(n) = n.try_into() else {
+            cx.emit_err(RustcScalableVectorCountOutOfRange { span: cx.attr_span, n });
+            return None;
+        };
+        Some(AttributeKind::RustcScalableVector { element_count: Some(n), span: cx.attr_span })
+    }
+}

compiler/rustc_attr_parsing/src/context.rs

@@ -59,7 +59,8 @@ use crate::attributes::prototype::CustomMirParser;
 use crate::attributes::repr::{AlignParser, AlignStaticParser, ReprParser};
 use crate::attributes::rustc_internal::{
     RustcLayoutScalarValidRangeEndParser, RustcLayoutScalarValidRangeStartParser, RustcMainParser,
-    RustcObjectLifetimeDefaultParser, RustcSimdMonomorphizeLaneLimitParser,
+    RustcObjectLifetimeDefaultParser, RustcScalableVectorParser,
+    RustcSimdMonomorphizeLaneLimitParser,
 };
 use crate::attributes::semantics::MayDangleParser;
 use crate::attributes::stability::{
@@ -207,6 +208,7 @@ attribute_parsers!(
         Single<RustcLayoutScalarValidRangeEndParser>,
         Single<RustcLayoutScalarValidRangeStartParser>,
         Single<RustcObjectLifetimeDefaultParser>,
+        Single<RustcScalableVectorParser>,
         Single<RustcSimdMonomorphizeLaneLimitParser>,
         Single<SanitizeParser>,
         Single<ShouldPanicParser>,

compiler/rustc_attr_parsing/src/session_diagnostics.rs

@@ -547,6 +547,15 @@ pub(crate) struct LinkOrdinalOutOfRange {
     pub ordinal: u128,
 }
 
+#[derive(Diagnostic)]
+#[diag(attr_parsing_rustc_scalable_vector_count_out_of_range)]
+#[note]
+pub(crate) struct RustcScalableVectorCountOutOfRange {
+    #[primary_span]
+    pub span: Span,
+    pub n: u128,
+}
+
 pub(crate) enum AttributeParseErrorReason<'a> {
     ExpectedNoArgs,
     ExpectedStringLiteral {

compiler/rustc_feature/src/builtin_attrs.rs

@@ -1417,6 +1417,10 @@ pub static BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         rustc_force_inline, Normal, template!(Word, NameValueStr: "reason"), WarnFollowing, EncodeCrossCrate::Yes,
         "`#[rustc_force_inline]` forces a free function to be inlined"
     ),
+    rustc_attr!(
+        rustc_scalable_vector, Normal, template!(List: &["count"]), WarnFollowing, EncodeCrossCrate::Yes,
+        "`#[rustc_scalable_vector]` defines a scalable vector type"
+    ),
 
     // ==========================================================================
     // Internal attributes, Testing:

compiler/rustc_hir/src/attrs/data_structures.rs

@@ -717,6 +717,14 @@ pub enum AttributeKind {
     /// Represents `#[rustc_pass_indirectly_in_non_rustic_abis]`
     RustcPassIndirectlyInNonRusticAbis(Span),
 
+    /// Represents `#[rustc_scalable_vector(N)]`
+    RustcScalableVector {
+        /// The base multiple of lanes that are in a scalable vector, if provided. `element_count`
+        /// is not provided for representing tuple types.
+        element_count: Option<u16>,
+        span: Span,
+    },
+
     /// Represents `#[rustc_should_not_be_called_on_const_items]`
     RustcShouldNotBeCalledOnConstItems(Span),
 

compiler/rustc_hir/src/attrs/encode_cross_crate.rs

@@ -91,6 +91,7 @@ impl AttributeKind {
             RustcMain => No,
             RustcObjectLifetimeDefault => No,
             RustcPassIndirectlyInNonRusticAbis(..) => No,
+            RustcScalableVector { .. } => Yes,
             RustcShouldNotBeCalledOnConstItems(..) => Yes,
             RustcSimdMonomorphizeLaneLimit(..) => Yes, // Affects layout computation, which needs to work cross-crate
             Sanitize { .. } => No,

compiler/rustc_hir/src/attrs/pretty_printing.rs

@@ -1,4 +1,5 @@
 use std::num::NonZero;
+use std::ops::Deref;
 
 use rustc_abi::Align;
 use rustc_ast::token::CommentKind;
@@ -35,6 +36,15 @@ impl<T: PrintAttribute> PrintAttribute for &T {
         T::print_attribute(self, p)
     }
 }
+impl<T: PrintAttribute> PrintAttribute for Box<T> {
+    fn should_render(&self) -> bool {
+        self.deref().should_render()
+    }
+
+    fn print_attribute(&self, p: &mut Printer) {
+        T::print_attribute(self.deref(), p)
+    }
+}
 impl<T: PrintAttribute> PrintAttribute for Option<T> {
     fn should_render(&self) -> bool {
         self.as_ref().is_some_and(|x| x.should_render())

compiler/rustc_middle/src/ty/mod.rs

@@ -24,7 +24,9 @@ pub use assoc::*;
 pub use generic_args::{GenericArgKind, TermKind, *};
 pub use generics::*;
 pub use intrinsic::IntrinsicDef;
-use rustc_abi::{Align, FieldIdx, Integer, IntegerType, ReprFlags, ReprOptions, VariantIdx};
+use rustc_abi::{
+    Align, FieldIdx, Integer, IntegerType, ReprFlags, ReprOptions, ScalableElt, VariantIdx,
+};
 use rustc_ast::expand::typetree::{FncTree, Kind, Type, TypeTree};
 use rustc_ast::node_id::NodeMap;
 pub use rustc_ast_ir::{Movability, Mutability, try_visit};
@@ -1513,6 +1515,17 @@ impl<'tcx> TyCtxt<'tcx> {
         }
 
         let attributes = self.get_all_attrs(did);
+        let elt = find_attr!(
+            attributes,
+            AttributeKind::RustcScalableVector { element_count, .. } => element_count
+        )
+        .map(|elt| match elt {
+            Some(n) => ScalableElt::ElementCount(*n),
+            None => ScalableElt::Container,
+        });
+        if elt.is_some() {
+            flags.insert(ReprFlags::IS_SCALABLE);
+        }
         if let Some(reprs) = find_attr!(attributes, AttributeKind::Repr { reprs, .. } => reprs) {
             for (r, _) in reprs {
                 flags.insert(match *r {
@@ -1577,7 +1590,14 @@ impl<'tcx> TyCtxt<'tcx> {
             flags.insert(ReprFlags::PASS_INDIRECTLY_IN_NON_RUSTIC_ABIS);
         }
 
-        ReprOptions { int: size, align: max_align, pack: min_pack, flags, field_shuffle_seed }
+        ReprOptions {
+            int: size,
+            align: max_align,
+            pack: min_pack,
+            flags,
+            field_shuffle_seed,
+            scalable: elt,
+        }
     }
 
     /// Look up the name of a definition across crates. This does not look at HIR.