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::prelude::Allow;
12use crate::attributes::{AttributeOrder, OnDuplicate, SingleAttributeParser};
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 ATTRIBUTE_ORDER: AttributeOrder = AttributeOrder::KeepInnermost;
22    const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::Error;
23    const ALLOWED_TARGETS: AllowedTargets = AllowedTargets::AllowList(&[
24        Allow(Target::Fn),
25        Allow(Target::Method(MethodKind::Inherent)),
26        Allow(Target::Method(MethodKind::Trait { body: true })),
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
34    fn convert(cx: &mut AcceptContext<'_, '_, S>, args: &ArgParser) -> Option<AttributeKind> {
35        let list = match args {
36            ArgParser::NoArgs => return Some(AttributeKind::RustcAutodiff(None)),
37            ArgParser::List(list) => list,
38            ArgParser::NameValue(_) => {
39                cx.expected_list_or_no_args(cx.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.expected_at_least_one_argument(list.span);
49            return None;
50        };
51        let Some(mode) = mode.meta_item() else {
52            cx.expected_identifier(mode.span());
53            return None;
54        };
55        let Ok(()) = mode.args().no_args() else {
56            cx.expected_identifier(mode.span());
57            return None;
58        };
59        let Some(mode) = mode.path().word() else {
60            cx.expected_identifier(mode.span());
61            return None;
62        };
63        let Ok(mode) = DiffMode::from_str(mode.as_str()) else {
64            cx.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.expected_specific_argument(activity.span(), DiffActivity::all_activities());
85                return None;
86            };
87            let Ok(()) = activity.args().no_args() else {
88                cx.expected_specific_argument(activity.span(), DiffActivity::all_activities());
89                return None;
90            };
91            let Some(activity) = activity.path().word() else {
92                cx.expected_specific_argument(activity.span(), DiffActivity::all_activities());
93                return None;
94            };
95            let Ok(activity) = DiffActivity::from_str(activity.as_str()) else {
96                cx.expected_specific_argument(activity.span, DiffActivity::all_activities());
97                return None;
98            };
99
100            activities.push(activity);
101        }
102        let Some(ret_activity) = activities.pop() else {
103            cx.expected_specific_argument(
104                list.span.with_lo(list.span.hi()),
105                DiffActivity::all_activities(),
106            );
107            return None;
108        };
109
110        Some(AttributeKind::RustcAutodiff(Some(Box::new(RustcAutodiff {
111            mode,
112            width,
113            input_activity: activities,
114            ret_activity,
115        }))))
116    }
117}
rustc_attr_parsing/attributes/autodiff.rs

rustc_attr_parsing/attributes/
autodiff.rs
