about summary refs log tree commit diff
path: root/library/compiler-builtins/crates/libm-macros/src/enums.rs
blob: b4646f984d471d8f7c6196c126556ec0932a821b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
use heck::ToUpperCamelCase;
use proc_macro2 as pm2;
use proc_macro2::{Ident, Span};
use quote::quote;
use syn::spanned::Spanned;
use syn::{Fields, ItemEnum, Variant};

use crate::{ALL_OPERATIONS, base_name};

/// Implement `#[function_enum]`, see documentation in `lib.rs`.
pub fn function_enum(
    mut item: ItemEnum,
    attributes: pm2::TokenStream,
) -> syn::Result<pm2::TokenStream> {
    expect_empty_enum(&item)?;
    let attr_span = attributes.span();
    let mut attr = attributes.into_iter();

    // Attribute should be the identifier of the `BaseName` enum.
    let Some(tt) = attr.next() else {
        return Err(syn::Error::new(attr_span, "expected one attribute"));
    };

    let pm2::TokenTree::Ident(base_enum) = tt else {
        return Err(syn::Error::new(tt.span(), "expected an identifier"));
    };

    if let Some(tt) = attr.next() {
        return Err(syn::Error::new(
            tt.span(),
            "unexpected token after identifier",
        ));
    }

    let enum_name = &item.ident;
    let mut as_str_arms = Vec::new();
    let mut from_str_arms = Vec::new();
    let mut base_arms = Vec::new();

    for func in ALL_OPERATIONS.iter() {
        let fn_name = func.name;
        let ident = Ident::new(&fn_name.to_upper_camel_case(), Span::call_site());
        let bname_ident = Ident::new(&base_name(fn_name).to_upper_camel_case(), Span::call_site());

        // Match arm for `fn as_str(self)` matcher
        as_str_arms.push(quote! { Self::#ident => #fn_name });
        from_str_arms.push(quote! { #fn_name => Self::#ident });

        // Match arm for `fn base_name(self)` matcher
        base_arms.push(quote! { Self::#ident => #base_enum::#bname_ident });

        let variant = Variant {
            attrs: Vec::new(),
            ident,
            fields: Fields::Unit,
            discriminant: None,
        };

        item.variants.push(variant);
    }

    let variants = item.variants.iter();

    let res = quote! {
        // Instantiate the enum
        #item

        impl #enum_name {
            /// All variants of this enum.
            pub const ALL: &[Self] = &[
                #( Self::#variants, )*
            ];

            /// The stringified version of this function name.
            pub const fn as_str(self) -> &'static str {
                match self {
                    #( #as_str_arms , )*
                }
            }

            /// If `s` is the name of a function, return it.
            pub fn from_str(s: &str) -> Option<Self> {
                let ret = match s {
                    #( #from_str_arms , )*
                    _ => return None,
                };
                Some(ret)
            }

            /// The base name enum for this function.
            pub const fn base_name(self) -> #base_enum {
                match self {
                    #( #base_arms, )*
                }
            }

            /// Return information about this operation.
            pub fn math_op(self) -> &'static crate::op::MathOpInfo {
                crate::op::ALL_OPERATIONS.iter().find(|op| op.name == self.as_str()).unwrap()
            }
        }
    };

    Ok(res)
}

/// Implement `#[base_name_enum]`, see documentation in `lib.rs`.
pub fn base_name_enum(
    mut item: ItemEnum,
    attributes: pm2::TokenStream,
) -> syn::Result<pm2::TokenStream> {
    expect_empty_enum(&item)?;
    if !attributes.is_empty() {
        let sp = attributes.span();
        return Err(syn::Error::new(sp.span(), "no attributes expected"));
    }

    let mut base_names: Vec<_> = ALL_OPERATIONS
        .iter()
        .map(|func| base_name(func.name))
        .collect();
    base_names.sort_unstable();
    base_names.dedup();

    let item_name = &item.ident;
    let mut as_str_arms = Vec::new();

    for base_name in base_names {
        let ident = Ident::new(&base_name.to_upper_camel_case(), Span::call_site());

        // Match arm for `fn as_str(self)` matcher
        as_str_arms.push(quote! { Self::#ident => #base_name });

        let variant = Variant {
            attrs: Vec::new(),
            ident,
            fields: Fields::Unit,
            discriminant: None,
        };

        item.variants.push(variant);
    }

    let res = quote! {
        // Instantiate the enum
        #item

        impl #item_name {
            /// The stringified version of this base name.
            pub const fn as_str(self) -> &'static str {
                match self {
                    #( #as_str_arms ),*
                }
            }
        }
    };

    Ok(res)
}

/// Verify that an enum is empty, otherwise return an error
fn expect_empty_enum(item: &ItemEnum) -> syn::Result<()> {
    if !item.variants.is_empty() {
        Err(syn::Error::new(
            item.variants.span(),
            "expected an empty enum",
        ))
    } else {
        Ok(())
    }
}