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

rustc_attr_parsing/attributes/
autodiff.rs
