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;
14use crate::parser::{ArgParser, MetaItemOrLitParser};
15use crate::target_checking::AllowedTargets;
16
17pub(crate) struct RustcAutodiffParser;
18
19impl SingleAttributeParser 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<'_, '_>, 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        cx.expect_no_args(mode.args())?;
56        let Some(mode) = mode.path().word() else {
57            cx.adcx().expected_identifier(mode.span());
58            return None;
59        };
60        let Ok(mode) = DiffMode::from_str(mode.as_str()) else {
61            cx.adcx().expected_specific_argument(mode.span, DiffMode::all_modes());
62            return None;
63        };
64
65        // Parse width
66        let width = if let Some(width) = items.peek()
67            && let MetaItemOrLitParser::Lit(width) = width
68            && let LitKind::Int(width, _) = width.kind
69            && let Ok(width) = width.0.try_into()
70        {
71            _ = items.next();
72            width
73        } else {
74            1
75        };
76
77        // Parse activities
78        let mut activities = ThinVec::new();
79        for activity in items {
80            let MetaItemOrLitParser::MetaItemParser(activity) = activity else {
81                cx.adcx()
82                    .expected_specific_argument(activity.span(), DiffActivity::all_activities());
83                return None;
84            };
85            cx.expect_no_args(activity.args())?;
86            let Some(activity) = activity.path().word() else {
87                cx.adcx()
88                    .expected_specific_argument(activity.span(), DiffActivity::all_activities());
89                return None;
90            };
91            let Ok(activity) = DiffActivity::from_str(activity.as_str()) else {
92                cx.adcx().expected_specific_argument(activity.span, DiffActivity::all_activities());
93                return None;
94            };
95
96            activities.push(activity);
97        }
98        let Some(ret_activity) = activities.pop() else {
99            cx.adcx().expected_specific_argument(
100                list.span.with_lo(list.span.hi()),
101                DiffActivity::all_activities(),
102            );
103            return None;
104        };
105
106        Some(AttributeKind::RustcAutodiff(Some(Box::new(RustcAutodiff {
107            mode,
108            width,
109            input_activity: activities,
110            ret_activity,
111        }))))
112    }
113}
rustc_attr_parsing/attributes/autodiff.rs

rustc_attr_parsing/attributes/
autodiff.rs
