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

@@ -93,9 +93,11 @@ bitflags! {
         // Other flags can still inhibit reordering and thus randomization.
         // The seed stored in `ReprOptions.field_shuffle_seed`.
         const RANDOMIZE_LAYOUT   = 1 << 4;
+        const IS_SCALABLE        = 1 << 5;
         // 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();
     }
@@ -132,6 +134,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(u128),
+    /// `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(
@@ -143,6 +158,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
@@ -159,6 +176,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/src/attributes/rustc_internal.rs

@@ -51,3 +51,26 @@ impl<S: Stage> SingleAttributeParser<S> for RustcObjectLifetimeDefaultParser {
         Some(AttributeKind::RustcObjectLifetimeDefault)
     }
 }
+
+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 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,
+            });
+        }
+
+        parse_single_integer(cx, args).map(|n| AttributeKind::RustcScalableVector {
+            element_count: Some(n),
+            span: cx.attr_span,
+        })
+    }
+}

compiler/rustc_attr_parsing/src/context.rs

@@ -44,7 +44,7 @@ use crate::attributes::proc_macro_attrs::{
 use crate::attributes::repr::{AlignParser, ReprParser};
 use crate::attributes::rustc_internal::{
     RustcLayoutScalarValidRangeEnd, RustcLayoutScalarValidRangeStart,
-    RustcObjectLifetimeDefaultParser,
+    RustcObjectLifetimeDefaultParser, RustcScalableVectorParser,
 };
 use crate::attributes::semantics::MayDangleParser;
 use crate::attributes::stability::{
@@ -174,6 +174,7 @@ attribute_parsers!(
         Single<RustcLayoutScalarValidRangeEnd>,
         Single<RustcLayoutScalarValidRangeStart>,
         Single<RustcObjectLifetimeDefaultParser>,
+        Single<RustcScalableVectorParser>,
         Single<ShouldPanicParser>,
         Single<SkipDuringMethodDispatchParser>,
         Single<TransparencyParser>,

compiler/rustc_feature/src/builtin_attrs.rs

@@ -1127,6 +1127,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

@@ -436,6 +436,14 @@ pub enum AttributeKind {
     /// Represents `#[rustc_object_lifetime_default]`.
     RustcObjectLifetimeDefault,
 
+    /// 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<u128>,
+        span: Span,
+    },
+
     /// Represents `#[should_panic]`
     ShouldPanic { reason: Option<Symbol>, span: Span },
 

compiler/rustc_hir/src/attrs/encode_cross_crate.rs

@@ -70,6 +70,7 @@ impl AttributeKind {
             RustcLayoutScalarValidRangeEnd(..) => Yes,
             RustcLayoutScalarValidRangeStart(..) => Yes,
             RustcObjectLifetimeDefault => No,
+            RustcScalableVector { .. } => Yes,
             ShouldPanic { .. } => No,
             SkipDuringMethodDispatch { .. } => No,
             SpecializationTrait(..) => No,

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::node_id::NodeMap;
 pub use rustc_ast_ir::{Movability, Mutability, try_visit};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexMap, FxIndexSet};
@@ -1469,9 +1471,20 @@ impl<'tcx> TyCtxt<'tcx> {
             field_shuffle_seed ^= user_seed;
         }
 
-        if let Some(reprs) =
-            find_attr!(self.get_all_attrs(did), AttributeKind::Repr { reprs, .. } => reprs)
-        {
+        let attrs = self.get_all_attrs(did);
+        let elt = find_attr!(
+            attrs,
+            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!(attrs, AttributeKind::Repr { reprs, .. } => reprs) {
             for (r, _) in reprs {
                 flags.insert(match *r {
                     attr::ReprRust => ReprFlags::empty(),
@@ -1530,7 +1543,14 @@ impl<'tcx> TyCtxt<'tcx> {
             flags.insert(ReprFlags::IS_LINEAR);
         }
 
-        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.

compiler/rustc_middle/src/ty/sty.rs

@@ -1196,6 +1196,14 @@ impl<'tcx> Ty<'tcx> {
         }
     }
 
+    #[inline]
+    pub fn is_scalable_vector(self) -> bool {
+        match self.kind() {
+            Adt(def, _) => def.repr().scalable(),
+            _ => false,
+        }
+    }
+
     pub fn sequence_element_type(self, tcx: TyCtxt<'tcx>) -> Ty<'tcx> {
         match self.kind() {
             Array(ty, _) | Slice(ty) => *ty,

compiler/rustc_passes/messages.ftl

@@ -631,6 +631,10 @@ passes_rustc_layout_scalar_valid_range_not_struct =
     attribute should be applied to a struct
     .label = not a struct
 
+passes_rustc_scalable_vector =
+    attribute should be applied to a struct
+    .label = not a struct
+
 passes_rustc_legacy_const_generics_index =
     #[rustc_legacy_const_generics] must have one index for each generic parameter
     .label = generic parameters

compiler/rustc_passes/src/check_attr.rs

@@ -271,6 +271,9 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                     AttributeKind::RustcLayoutScalarValidRangeStart(_num, attr_span)
                     | AttributeKind::RustcLayoutScalarValidRangeEnd(_num, attr_span),
                 ) => self.check_rustc_layout_scalar_valid_range(*attr_span, span, target),
+                Attribute::Parsed(AttributeKind::RustcScalableVector {
+                    span: attr_span, ..
+                }) => self.check_rustc_scalable_vector(*attr_span, span, target),
                 Attribute::Parsed(AttributeKind::ExportStable) => {
                     // handled in `check_export`
                 }
@@ -1776,6 +1779,13 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
         }
     }
 
+    fn check_rustc_scalable_vector(&self, attr_span: Span, span: Span, target: Target) {
+        if target != Target::Struct {
+            self.dcx().emit_err(errors::RustcScalableVector { attr_span, span });
+            return;
+        }
+    }
+
     /// Checks if `#[rustc_legacy_const_generics]` is applied to a function and has a valid argument.
     fn check_rustc_legacy_const_generics(
         &self,