Port rustc_expected_cgu_reuse to the new attribute parser

Cargo.lock

@@ -3687,7 +3687,6 @@ dependencies = [
  "serde_json",
  "smallvec",
  "tempfile",
- "thin-vec",
  "thorin-dwp",
  "tracing",
  "wasm-encoder 0.219.2",

compiler/rustc_attr_parsing/src/attributes/rustc_internal.rs

@@ -2,7 +2,8 @@ use std::path::PathBuf;
 
 use rustc_ast::{LitIntType, LitKind, MetaItemLit};
 use rustc_hir::attrs::{
-    BorrowckGraphvizFormatKind, RustcCleanAttribute, RustcCleanQueries, RustcLayoutType,
+    BorrowckGraphvizFormatKind, CguFields, CguKind, DivergingBlockBehavior,
+    DivergingFallbackBehavior, RustcCleanAttribute, RustcCleanQueries, RustcLayoutType,
     RustcMirKind,
 };
 use rustc_session::errors;
@@ -10,7 +11,9 @@ use rustc_span::Symbol;
 
 use super::prelude::*;
 use super::util::parse_single_integer;
-use crate::session_diagnostics::{AttributeRequiresOpt, RustcScalableVectorCountOutOfRange};
+use crate::session_diagnostics::{
+    AttributeRequiresOpt, CguFieldsMissing, RustcScalableVectorCountOutOfRange,
+};
 
 pub(crate) struct RustcMainParser;
 
@@ -204,6 +207,144 @@ impl<S: Stage> NoArgsAttributeParser<S> for RustcLintOptTyParser {
     const CREATE: fn(Span) -> AttributeKind = |_| AttributeKind::RustcLintOptTy;
 }
 
+fn parse_cgu_fields<S: Stage>(
+    cx: &mut AcceptContext<'_, '_, S>,
+    args: &ArgParser,
+    accepts_kind: bool,
+) -> Option<(Symbol, Symbol, Option<CguKind>)> {
+    let Some(args) = args.list() else {
+        cx.expected_list(cx.attr_span, args);
+        return None;
+    };
+
+    let mut cfg = None::<(Symbol, Span)>;
+    let mut module = None::<(Symbol, Span)>;
+    let mut kind = None::<(Symbol, Span)>;
+
+    for arg in args.mixed() {
+        let Some(arg) = arg.meta_item() else {
+            cx.expected_name_value(args.span, None);
+            continue;
+        };
+
+        let res = match arg.ident().map(|i| i.name) {
+            Some(sym::cfg) => &mut cfg,
+            Some(sym::module) => &mut module,
+            Some(sym::kind) if accepts_kind => &mut kind,
+            _ => {
+                cx.expected_specific_argument(
+                    arg.path().span(),
+                    if accepts_kind {
+                        &[sym::cfg, sym::module, sym::kind]
+                    } else {
+                        &[sym::cfg, sym::module]
+                    },
+                );
+                continue;
+            }
+        };
+
+        let Some(i) = arg.args().name_value() else {
+            cx.expected_name_value(arg.span(), None);
+            continue;
+        };
+
+        let Some(str) = i.value_as_str() else {
+            cx.expected_string_literal(i.value_span, Some(i.value_as_lit()));
+            continue;
+        };
+
+        if res.is_some() {
+            cx.duplicate_key(arg.span(), arg.ident().unwrap().name);
+            continue;
+        }
+
+        *res = Some((str, i.value_span));
+    }
+
+    let Some((cfg, _)) = cfg else {
+        cx.emit_err(CguFieldsMissing { span: args.span, name: &cx.attr_path, field: sym::cfg });
+        return None;
+    };
+    let Some((module, _)) = module else {
+        cx.emit_err(CguFieldsMissing { span: args.span, name: &cx.attr_path, field: sym::module });
+        return None;
+    };
+    let kind = if let Some((kind, span)) = kind {
+        Some(match kind {
+            sym::no => CguKind::No,
+            sym::pre_dash_lto => CguKind::PreDashLto,
+            sym::post_dash_lto => CguKind::PostDashLto,
+            sym::any => CguKind::Any,
+            _ => {
+                cx.expected_specific_argument_strings(
+                    span,
+                    &[sym::no, sym::pre_dash_lto, sym::post_dash_lto, sym::any],
+                );
+                return None;
+            }
+        })
+    } else {
+        // return None so that an unwrap for the attributes that need it is ok.
+        if accepts_kind {
+            cx.emit_err(CguFieldsMissing {
+                span: args.span,
+                name: &cx.attr_path,
+                field: sym::kind,
+            });
+            return None;
+        };
+
+        None
+    };
+
+    Some((cfg, module, kind))
+}
+
+#[derive(Default)]
+pub(crate) struct RustcCguTestAttributeParser {
+    items: ThinVec<(Span, CguFields)>,
+}
+
+impl<S: Stage> AttributeParser<S> for RustcCguTestAttributeParser {
+    const ATTRIBUTES: AcceptMapping<Self, S> = &[
+        (
+            &[sym::rustc_partition_reused],
+            template!(List: &[r#"cfg = "...", module = "...""#]),
+            |this, cx, args| {
+                this.items.extend(parse_cgu_fields(cx, args, false).map(|(cfg, module, _)| {
+                    (cx.attr_span, CguFields::PartitionReused { cfg, module })
+                }));
+            },
+        ),
+        (
+            &[sym::rustc_partition_codegened],
+            template!(List: &[r#"cfg = "...", module = "...""#]),
+            |this, cx, args| {
+                this.items.extend(parse_cgu_fields(cx, args, false).map(|(cfg, module, _)| {
+                    (cx.attr_span, CguFields::PartitionCodegened { cfg, module })
+                }));
+            },
+        ),
+        (
+            &[sym::rustc_expected_cgu_reuse],
+            template!(List: &[r#"cfg = "...", module = "...", kind = "...""#]),
+            |this, cx, args| {
+                this.items.extend(parse_cgu_fields(cx, args, true).map(|(cfg, module, kind)| {
+                    // unwrap ok because if not given, we return None in `parse_cgu_fields`.
+                    (cx.attr_span, CguFields::ExpectedCguReuse { cfg, module, kind: kind.unwrap() })
+                }));
+            },
+        ),
+    ];
+
+    const ALLOWED_TARGETS: AllowedTargets =
+        AllowedTargets::AllowList(&[Allow(Target::Mod), Allow(Target::Crate)]);
+
+    fn finalize(self, _cx: &FinalizeContext<'_, '_, S>) -> Option<AttributeKind> {
+        Some(AttributeKind::RustcCguTestAttr(self.items))
+    }
+}
 pub(crate) struct RustcLintQueryInstabilityParser;
 
 impl<S: Stage> NoArgsAttributeParser<S> for RustcLintQueryInstabilityParser {

compiler/rustc_attr_parsing/src/context.rs

@@ -149,6 +149,7 @@ attribute_parsers!(
         NakedParser,
         StabilityParser,
         UsedParser,
+        RustcCguTestAttributeParser,
         // tidy-alphabetical-end
 
         // tidy-alphabetical-start

compiler/rustc_attr_parsing/src/session_diagnostics.rs

@@ -45,6 +45,15 @@ pub(crate) struct DocAliasStartEnd<'a> {
     pub attr_str: &'a str,
 }
 
+#[derive(Diagnostic)]
+#[diag("`#[{$name})]` is missing a `{$field}` argument")]
+pub(crate) struct CguFieldsMissing<'a> {
+    #[primary_span]
+    pub span: Span,
+    pub name: &'a AttrPath,
+    pub field: Symbol,
+}
+
 #[derive(Diagnostic)]
 #[diag("`#![doc({$attr_name} = \"...\")]` isn't allowed as a crate-level attribute")]
 pub(crate) struct DocAttrNotCrateLevel {

compiler/rustc_codegen_ssa/Cargo.toml

@@ -36,7 +36,6 @@ rustc_trait_selection = { path = "../rustc_trait_selection" }
 serde_json = "1.0.59"
 smallvec = { version = "1.8.1", features = ["union", "may_dangle"] }
 tempfile = "3.2"
-thin-vec = "0.2.12"
 thorin-dwp = "0.9"
 tracing = "0.1"
 wasm-encoder = "0.219"

compiler/rustc_codegen_ssa/src/assert_module_sources.rs

@@ -28,13 +28,13 @@ use std::fmt;
 
 use rustc_data_structures::unord::{UnordMap, UnordSet};
 use rustc_errors::{DiagArgValue, IntoDiagArg};
-use rustc_hir as hir;
+use rustc_hir::attrs::{AttributeKind, CguFields, CguKind};
 use rustc_hir::def_id::LOCAL_CRATE;
+use rustc_hir::{self as hir, find_attr};
 use rustc_middle::mir::mono::CodegenUnitNameBuilder;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::Session;
-use rustc_span::{Span, Symbol, sym};
-use thin_vec::ThinVec;
+use rustc_span::{Span, Symbol};
 use tracing::debug;
 
 use crate::errors;
@@ -63,9 +63,7 @@ pub fn assert_module_sources(tcx: TyCtxt<'_>, set_reuse: &dyn Fn(&mut CguReuseTr
             },
         };
 
-        for attr in tcx.hir_attrs(rustc_hir::CRATE_HIR_ID) {
-            ams.check_attr(attr);
-        }
+        ams.check_attrs(tcx.hir_attrs(rustc_hir::CRATE_HIR_ID));
 
         set_reuse(&mut ams.cgu_reuse_tracker);
 
@@ -89,109 +87,91 @@ struct AssertModuleSource<'tcx> {
 }
 
 impl<'tcx> AssertModuleSource<'tcx> {
-    fn check_attr(&mut self, attr: &hir::Attribute) {
-        let (expected_reuse, comp_kind) = if attr.has_name(sym::rustc_partition_reused) {
-            (CguReuse::PreLto, ComparisonKind::AtLeast)
-        } else if attr.has_name(sym::rustc_partition_codegened) {
-            (CguReuse::No, ComparisonKind::Exact)
-        } else if attr.has_name(sym::rustc_expected_cgu_reuse) {
-            match self.field(attr, sym::kind) {
-                sym::no => (CguReuse::No, ComparisonKind::Exact),
-                sym::pre_dash_lto => (CguReuse::PreLto, ComparisonKind::Exact),
-                sym::post_dash_lto => (CguReuse::PostLto, ComparisonKind::Exact),
-                sym::any => (CguReuse::PreLto, ComparisonKind::AtLeast),
-                other => {
-                    self.tcx
-                        .dcx()
-                        .emit_fatal(errors::UnknownReuseKind { span: attr.span(), kind: other });
-                }
+    fn check_attrs(&mut self, attrs: &[hir::Attribute]) {
+        for &(span, cgu_fields) in find_attr!(attrs,
+            AttributeKind::RustcCguTestAttr(e) => e)
+        .into_iter()
+        .flatten()
+        {
+            let (expected_reuse, comp_kind) = match cgu_fields {
+                CguFields::PartitionReused { .. } => (CguReuse::PreLto, ComparisonKind::AtLeast),
+                CguFields::PartitionCodegened { .. } => (CguReuse::No, ComparisonKind::Exact),
+                CguFields::ExpectedCguReuse { kind, .. } => match kind {
+                    CguKind::No => (CguReuse::No, ComparisonKind::Exact),
+                    CguKind::PreDashLto => (CguReuse::PreLto, ComparisonKind::Exact),
+                    CguKind::PostDashLto => (CguReuse::PostLto, ComparisonKind::Exact),
+                    CguKind::Any => (CguReuse::PreLto, ComparisonKind::AtLeast),
+                },
+            };
+            let (CguFields::ExpectedCguReuse { cfg, module, .. }
+            | CguFields::PartitionCodegened { cfg, module }
+            | CguFields::PartitionReused { cfg, module }) = cgu_fields;
+
+            if !self.tcx.sess.opts.unstable_opts.query_dep_graph {
+                self.tcx.dcx().emit_fatal(errors::MissingQueryDepGraph { span });
             }
-        } else {
-            return;
-        };
 
-        if !self.tcx.sess.opts.unstable_opts.query_dep_graph {
-            self.tcx.dcx().emit_fatal(errors::MissingQueryDepGraph { span: attr.span() });
-        }
+            if !self.check_config(cfg) {
+                debug!("check_attr: config does not match, ignoring attr");
+                return;
+            }
 
-        if !self.check_config(attr) {
-            debug!("check_attr: config does not match, ignoring attr");
-            return;
-        }
+            let user_path = module.as_str();
+            let crate_name = self.tcx.crate_name(LOCAL_CRATE);
+            let crate_name = crate_name.as_str();
 
-        let user_path = self.field(attr, sym::module).to_string();
-        let crate_name = self.tcx.crate_name(LOCAL_CRATE).to_string();
+            if !user_path.starts_with(&crate_name) {
+                self.tcx.dcx().emit_fatal(errors::MalformedCguName { span, user_path, crate_name });
+            }
 
-        if !user_path.starts_with(&crate_name) {
-            self.tcx.dcx().emit_fatal(errors::MalformedCguName {
-                span: attr.span(),
-                user_path,
-                crate_name,
-            });
-        }
+            // Split of the "special suffix" if there is one.
+            let (user_path, cgu_special_suffix) = if let Some(index) = user_path.rfind('.') {
+                (&user_path[..index], Some(&user_path[index + 1..]))
+            } else {
+                (&user_path[..], None)
+            };
 
-        // Split of the "special suffix" if there is one.
-        let (user_path, cgu_special_suffix) = if let Some(index) = user_path.rfind('.') {
-            (&user_path[..index], Some(&user_path[index + 1..]))
-        } else {
-            (&user_path[..], None)
-        };
+            let mut iter = user_path.split('-');
 
-        let mut iter = user_path.split('-');
+            // Remove the crate name
+            assert_eq!(iter.next().unwrap(), crate_name);
 
-        // Remove the crate name
-        assert_eq!(iter.next().unwrap(), crate_name);
+            let cgu_path_components = iter.collect::<Vec<_>>();
 
-        let cgu_path_components = iter.collect::<Vec<_>>();
+            let cgu_name_builder = &mut CodegenUnitNameBuilder::new(self.tcx);
+            let cgu_name = cgu_name_builder.build_cgu_name(
+                LOCAL_CRATE,
+                cgu_path_components,
+                cgu_special_suffix,
+            );
 
-        let cgu_name_builder = &mut CodegenUnitNameBuilder::new(self.tcx);
-        let cgu_name =
-            cgu_name_builder.build_cgu_name(LOCAL_CRATE, cgu_path_components, cgu_special_suffix);
+            debug!("mapping '{user_path}' to cgu name '{cgu_name}'");
 
-        debug!("mapping '{}' to cgu name '{}'", self.field(attr, sym::module), cgu_name);
+            if !self.available_cgus.contains(&cgu_name) {
+                let cgu_names: Vec<&str> =
+                    self.available_cgus.items().map(|cgu| cgu.as_str()).into_sorted_stable_ord();
+                self.tcx.dcx().emit_err(errors::NoModuleNamed {
+                    span,
+                    user_path,
+                    cgu_name,
+                    cgu_names: cgu_names.join(", "),
+                });
+            }
 
-        if !self.available_cgus.contains(&cgu_name) {
-            let cgu_names: Vec<&str> =
-                self.available_cgus.items().map(|cgu| cgu.as_str()).into_sorted_stable_ord();
-            self.tcx.dcx().emit_err(errors::NoModuleNamed {
-                span: attr.span(),
-                user_path,
+            self.cgu_reuse_tracker.set_expectation(
                 cgu_name,
-                cgu_names: cgu_names.join(", "),
-            });
+                user_path,
+                span,
+                expected_reuse,
+                comp_kind,
+            );
         }
-
-        self.cgu_reuse_tracker.set_expectation(
-            cgu_name,
-            user_path,
-            attr.span(),
-            expected_reuse,
-            comp_kind,
-        );
-    }
-
-    fn field(&self, attr: &hir::Attribute, name: Symbol) -> Symbol {
-        for item in attr.meta_item_list().unwrap_or_else(ThinVec::new) {
-            if item.has_name(name) {
-                if let Some(value) = item.value_str() {
-                    return value;
-                } else {
-                    self.tcx.dcx().emit_fatal(errors::FieldAssociatedValueExpected {
-                        span: item.span(),
-                        name,
-                    });
-                }
-            }
-        }
-
-        self.tcx.dcx().emit_fatal(errors::NoField { span: attr.span(), name });
     }
 
     /// Scan for a `cfg="foo"` attribute and check whether we have a
     /// cfg flag called `foo`.
-    fn check_config(&self, attr: &hir::Attribute) -> bool {
+    fn check_config(&self, value: Symbol) -> bool {
         let config = &self.tcx.sess.psess.config;
-        let value = self.field(attr, sym::cfg);
         debug!("check_config(config={:?}, value={:?})", config, value);
         if config.iter().any(|&(name, _)| name == value) {
             debug!("check_config: matched");

compiler/rustc_codegen_ssa/src/errors.rs

@@ -42,14 +42,6 @@ pub(crate) struct CguNotRecorded<'a> {
     pub cgu_name: &'a str,
 }
 
-#[derive(Diagnostic)]
-#[diag("unknown cgu-reuse-kind `{$kind}` specified")]
-pub(crate) struct UnknownReuseKind {
-    #[primary_span]
-    pub span: Span,
-    pub kind: Symbol,
-}
-
 #[derive(Diagnostic)]
 #[diag("found CGU-reuse attribute but `-Zquery-dep-graph` was not specified")]
 pub(crate) struct MissingQueryDepGraph {
@@ -61,11 +53,11 @@ pub(crate) struct MissingQueryDepGraph {
 #[diag(
     "found malformed codegen unit name `{$user_path}`. codegen units names must always start with the name of the crate (`{$crate_name}` in this case)"
 )]
-pub(crate) struct MalformedCguName {
+pub(crate) struct MalformedCguName<'a> {
     #[primary_span]
     pub span: Span,
-    pub user_path: String,
-    pub crate_name: String,
+    pub user_path: &'a str,
+    pub crate_name: &'a str,
 }
 
 #[derive(Diagnostic)]
@@ -78,22 +70,6 @@ pub(crate) struct NoModuleNamed<'a> {
     pub cgu_names: String,
 }
 
-#[derive(Diagnostic)]
-#[diag("associated value expected for `{$name}`")]
-pub(crate) struct FieldAssociatedValueExpected {
-    #[primary_span]
-    pub span: Span,
-    pub name: Symbol,
-}
-
-#[derive(Diagnostic)]
-#[diag("no field `{$name}`")]
-pub(crate) struct NoField {
-    #[primary_span]
-    pub span: Span,
-    pub name: Symbol,
-}
-
 #[derive(Diagnostic)]
 #[diag("failed to write lib.def file: {$error}")]
 pub(crate) struct LibDefWriteFailure {

compiler/rustc_hir/src/attrs/data_structures.rs

@@ -50,6 +50,19 @@ pub struct EiiDecl {
 }
 
 #[derive(Copy, Clone, PartialEq, Encodable, Decodable, Debug, HashStable_Generic, PrintAttribute)]
+pub enum CguKind {
+    No,
+    PreDashLto,
+    PostDashLto,
+    Any,
+}
+
+#[derive(Copy, Clone, PartialEq, Encodable, Decodable, Debug, HashStable_Generic, PrintAttribute)]
+pub enum CguFields {
+    PartitionReused { cfg: Symbol, module: Symbol },
+    PartitionCodegened { cfg: Symbol, module: Symbol },
+    ExpectedCguReuse { cfg: Symbol, module: Symbol, kind: CguKind },
+}
 pub enum InlineAttr {
     None,
     Hint,
@@ -1117,6 +1130,9 @@ pub enum AttributeKind {
     /// Represents `#[rustc_evaluate_where_clauses]`
     RustcEvaluateWhereClauses,
 
+    /// Represents `#[rustc_expected_cgu_reuse]`, `#[rustc_partition_codegened]` and `#[rustc_partition_reused]`.
+    RustcCguTestAttr(ThinVec<(Span, CguFields)>),
+
     /// Represents `#[rustc_has_incoherent_inherent_impls]`
     RustcHasIncoherentInherentImpls,
 

compiler/rustc_hir/src/attrs/encode_cross_crate.rs

@@ -146,6 +146,7 @@ impl AttributeKind {
             RustcOffloadKernel => Yes,
             RustcOutlives => No,
             RustcParenSugar(..) => No,
+            RustcCguTestAttr { .. } => No,
             RustcPassByValue(..) => Yes,
             RustcPassIndirectlyInNonRusticAbis(..) => No,
             RustcPreserveUbChecks => No,

compiler/rustc_passes/src/check_attr.rs

@@ -338,6 +338,7 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                     | AttributeKind::RustcOffloadKernel
                     | AttributeKind::RustcOutlives
                     | AttributeKind::RustcParenSugar(..)
+                    | AttributeKind::RustcCguTestAttr(..)
                     | AttributeKind::RustcPassByValue (..)
                     | AttributeKind::RustcPassIndirectlyInNonRusticAbis(..)
                     | AttributeKind::RustcPreserveUbChecks
@@ -407,9 +408,6 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                             | sym::rustc_autodiff
                             | sym::rustc_capture_analysis
                             | sym::rustc_mir
-                            | sym::rustc_partition_reused
-                            | sym::rustc_partition_codegened
-                            | sym::rustc_expected_cgu_reuse
                             // crate-level attrs, are checked below
                             | sym::feature
                             | sym::register_tool