rustc_attr_parsing/attributes/
autodiff.rs

1use std::str::FromStr;
2
3use rustc_ast::LitKind;
4use rustc_ast::expand::autodiff_attrs::{DiffActivity, DiffMode};
5use rustc_feature::{AttributeStability, AttributeTemplate, template};
6use rustc_hir::attrs::{AttributeKind, RustcAutodiff};
7use rustc_hir::{MethodKind, Target};
8use rustc_span::{Symbol, sym};
9use thin_vec::ThinVec;
10
11use crate::attributes::SingleAttributeParser;
12use crate::attributes::prelude::Allow;
13use crate::context::AcceptContext;
14use crate::parser::{ArgParser, MetaItemOrLitParser};
15use crate::target_checking::AllowedTargets;
16use crate::unstable;
17
18pub(crate) struct RustcAutodiffParser;
19
20impl SingleAttributeParser for RustcAutodiffParser {
21    const PATH: &[Symbol] = &[sym::rustc_autodiff];
22    const ALLOWED_TARGETS: AllowedTargets = AllowedTargets::AllowList(&[
23        Allow(Target::Fn),
24        Allow(Target::Method(MethodKind::Inherent)),
25        Allow(Target::Method(MethodKind::Trait { body: true })),
26        Allow(Target::Method(MethodKind::Trait { body: false })),
27        Allow(Target::Method(MethodKind::TraitImpl)),
28    ]);
29    const TEMPLATE: AttributeTemplate = ::rustc_feature::AttributeTemplate {
    word: false,
    list: Some(&["MODE", "WIDTH", "INPUT_ACTIVITIES", "OUTPUT_ACTIVITY"]),
    one_of: &[],
    name_value_str: None,
    docs: Some("https://doc.rust-lang.org/std/autodiff/index.html"),
}template!(
30        List: &["MODE", "WIDTH", "INPUT_ACTIVITIES", "OUTPUT_ACTIVITY"],
31        "https://doc.rust-lang.org/std/autodiff/index.html"
32    );
33    const STABILITY: AttributeStability = AttributeStability::Unstable {
    gate_name: rustc_span::sym::rustc_attrs,
    gate_check: rustc_feature::Features::rustc_attrs,
    notes: &[],
}unstable!(rustc_attrs);
34
35    fn convert(cx: &mut AcceptContext<'_, '_>, args: &ArgParser) -> Option<AttributeKind> {
36        let list = match args {
37            ArgParser::NoArgs => return Some(AttributeKind::RustcAutodiff(None)),
38            ArgParser::List(list) => list,
39            ArgParser::NameValue(_) => {
40                let attr_span = cx.attr_span;
41                cx.adcx().expected_list_or_no_args(attr_span);
42                return None;
43            }
44        };
45
46        let mut items = list.mixed().peekable();
47
48        // Parse name
49        let Some(mode) = items.next() else {
50            cx.adcx().expected_at_least_one_argument(list.span);
51            return None;
52        };
53        let Some(mode) = mode.meta_item() else {
54            cx.adcx().expected_identifier(mode.span());
55            return None;
56        };
57        cx.expect_no_args(mode.args())?;
58        let Some(mode) = mode.path().word() else {
59            cx.adcx().expected_identifier(mode.span());
60            return None;
61        };
62        let Ok(mode) = DiffMode::from_str(mode.as_str()) else {
63            cx.adcx().expected_specific_argument(mode.span, DiffMode::all_modes());
64            return None;
65        };
66
67        // Parse width
68        let width = if let Some(width) = items.peek()
69            && let MetaItemOrLitParser::Lit(width) = width
70            && let LitKind::Int(width, _) = width.kind
71            && let Ok(width) = width.0.try_into()
72        {
73            _ = items.next();
74            width
75        } else {
76            1
77        };
78
79        // Parse activities
80        let mut activities = ThinVec::new();
81        for activity in items {
82            let MetaItemOrLitParser::MetaItemParser(activity) = activity else {
83                cx.adcx()
84                    .expected_specific_argument(activity.span(), DiffActivity::all_activities());
85                return None;
86            };
87            cx.expect_no_args(activity.args())?;
88            let Some(activity) = activity.path().word() else {
89                cx.adcx()
90                    .expected_specific_argument(activity.span(), DiffActivity::all_activities());
91                return None;
92            };
93            let Ok(activity) = DiffActivity::from_str(activity.as_str()) else {
94                cx.adcx().expected_specific_argument(activity.span, DiffActivity::all_activities());
95                return None;
96            };
97
98            activities.push(activity);
99        }
100        let Some(ret_activity) = activities.pop() else {
101            cx.adcx().expected_specific_argument(
102                list.span.with_lo(list.span.hi()),
103                DiffActivity::all_activities(),
104            );
105            return None;
106        };
107
108        Some(AttributeKind::RustcAutodiff(Some(Box::new(RustcAutodiff {
109            mode,
110            width,
111            input_activity: activities,
112            ret_activity,
113        }))))
114    }
115}
rustc_attr_parsing/attributes/autodiff.rs

rustc_attr_parsing/attributes/
autodiff.rs
