Allow parsing of #[spirv(...)] attribute with Iterators

Author: Firestar99

crates/rustc_codegen_spirv/src/abi.rs

@@ -546,11 +546,8 @@ impl<'tcx> ConvSpirvType<'tcx> for TyAndLayout<'tcx> {
                 span = cx.tcx.def_span(adt.did());
             }
 
-            let attrs = AggregatedSpirvAttributes::parse(
-                cx.tcx,
-                &cx.sym,
-                cx.tcx.get_attrs_unchecked(adt.did()),
-            );
+            let attrs =
+                AggregatedSpirvAttributes::parse(cx.tcx, &cx.sym, cx.tcx.get_all_attrs(adt.did()));
 
             if let Some(intrinsic_type_attr) = attrs.intrinsic_type.map(|attr| attr.value)
                 && let Ok(spirv_type) =

crates/rustc_codegen_spirv/src/attr.rs

@@ -2,7 +2,7 @@
 //!
 //! The attribute-checking parts of this try to follow `rustc_passes::check_attr`.
 
-use crate::symbols::Symbols;
+use crate::symbols::{Symbols, parse_attrs_for_checking};
 use rspirv::spirv::{BuiltIn, ExecutionMode, ExecutionModel, StorageClass};
 use rustc_hir as hir;
 use rustc_hir::def_id::LocalModDefId;
@@ -146,12 +146,16 @@ impl AggregatedSpirvAttributes {
     ///
     /// Any errors for malformed/duplicate attributes will have been reported
     /// prior to codegen, by the `attr` check pass.
-    pub fn parse<'tcx>(tcx: TyCtxt<'tcx>, sym: &Symbols, attrs: &'tcx [Attribute]) -> Self {
+    pub fn parse<'tcx>(
+        tcx: TyCtxt<'tcx>,
+        sym: &Symbols,
+        attrs: impl Iterator<Item = &'tcx Attribute>,
+    ) -> Self {
         let mut aggregated_attrs = Self::default();
 
         // NOTE(eddyb) `span_delayed_bug` ensures that if attribute checking fails
         // to see an attribute error, it will cause an ICE instead.
-        for parse_attr_result in crate::symbols::parse_attrs_for_checking(sym, attrs) {
+        for parse_attr_result in parse_attrs_for_checking(sym, attrs) {
             let (span, parsed_attr) = match parse_attr_result {
                 Ok(span_and_parsed_attr) => span_and_parsed_attr,
                 Err((span, msg)) => {
@@ -276,10 +280,8 @@ impl CheckSpirvAttrVisitor<'_> {
     fn check_spirv_attributes(&self, hir_id: HirId, target: Target) {
         let mut aggregated_attrs = AggregatedSpirvAttributes::default();
 
-        let parse_attrs = |attrs| crate::symbols::parse_attrs_for_checking(&self.sym, attrs);
-
         let attrs = self.tcx.hir_attrs(hir_id);
-        for parse_attr_result in parse_attrs(attrs) {
+        for parse_attr_result in parse_attrs_for_checking(&self.sym, attrs.into_iter()) {
             let (span, parsed_attr) = match parse_attr_result {
                 Ok(span_and_parsed_attr) => span_and_parsed_attr,
                 Err((span, msg)) => {
@@ -324,9 +326,12 @@ impl CheckSpirvAttrVisitor<'_> {
                 | SpirvAttribute::SpecConstant(_) => match target {
                     Target::Param => {
                         let parent_hir_id = self.tcx.parent_hir_id(hir_id);
-                        let parent_is_entry_point = parse_attrs(self.tcx.hir_attrs(parent_hir_id))
-                            .filter_map(|r| r.ok())
-                            .any(|(_, attr)| matches!(attr, SpirvAttribute::Entry(_)));
+                        let parent_is_entry_point = parse_attrs_for_checking(
+                            &self.sym,
+                            self.tcx.hir_attrs(parent_hir_id).iter(),
+                        )
+                        .filter_map(|r| r.ok())
+                        .any(|(_, attr)| matches!(attr, SpirvAttribute::Entry(_)));
                         if !parent_is_entry_point {
                             self.tcx.dcx().span_err(
                                 span,

crates/rustc_codegen_spirv/src/codegen_cx/declare.rs

@@ -133,11 +133,8 @@ impl<'tcx> CodegenCx<'tcx> {
             self.set_linkage(fn_id, symbol_name.to_owned(), linkage);
         }
 
-        let attrs = AggregatedSpirvAttributes::parse(
-            self.tcx,
-            &self.sym,
-            self.tcx.get_attrs_unchecked(def_id),
-        );
+        let attrs =
+            AggregatedSpirvAttributes::parse(self.tcx, &self.sym, self.tcx.get_all_attrs(def_id));
         if let Some(entry) = attrs.entry.map(|attr| attr.value) {
             // HACK(eddyb) early insert to let `shader_entry_stub` call this
             // very function via `get_fn_addr`.

crates/rustc_codegen_spirv/src/codegen_cx/entry.rs

@@ -436,7 +436,7 @@ impl<'tcx> CodegenCx<'tcx> {
         let attrs = AggregatedSpirvAttributes::parse(
             self.tcx,
             &self.sym,
-            self.tcx.hir_attrs(hir_param.hir_id),
+            self.tcx.hir_attrs(hir_param.hir_id).iter(),
         );
 
         let EntryParamDeducedFromRustRefOrValue {

crates/rustc_codegen_spirv/src/symbols.rs

@@ -441,11 +441,11 @@ impl Symbols {
 type ParseAttrError = (Span, String);
 
 // FIXME(eddyb) maybe move this to `attr`?
-pub(crate) fn parse_attrs_for_checking<'a>(
-    sym: &'a Symbols,
-    attrs: &'a [Attribute],
-) -> impl Iterator<Item = Result<(Span, SpirvAttribute), ParseAttrError>> + 'a {
-    attrs.iter().flat_map(move |attr| {
+pub(crate) fn parse_attrs_for_checking<'a, 'b>(
+    sym: &'b Symbols,
+    attrs: impl Iterator<Item = &'a Attribute> + 'b,
+) -> impl Iterator<Item = Result<(Span, SpirvAttribute), ParseAttrError>> + 'b {
+    attrs.flat_map(move |attr| {
         let (whole_attr_error, args) = match attr {
             Attribute::Unparsed(item) => {
                 // #[...]