auto_const_array/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::Span;
3use quote::quote_spanned;
4use syn::{
5    bracketed,
6    parse::{Parse, ParseStream, Parser},
7    punctuated::Punctuated,
8    Attribute, Expr, Ident, Result, Token, Type, Visibility,
9};
10
11/// Declare a new const array without specify length.
12/// It helps when apply conditional compilation to part of a const array.
13///
14/// # Syntax
15/// The macro wraps any number of const array declarations(with length `_`).
16///
17/// ```rust
18/// use auto_const_array::auto_const_array;
19/// auto_const_array! {
20///    // Additional attributes and docs are supported.
21///    /// Common array with public visibility.
22///    #[allow(unused)]
23///    pub const ARRAY_COMMON: [u8; _] = [1, 2, 4];
24///    /// Special array with cfg conditional compiling.
25///    const ARRAY_WITH_ATTR: [u8; _] = [1, #[cfg(unix)] 2]
26/// }
27/// ```
28#[proc_macro]
29pub fn auto_const_array(input: TokenStream) -> TokenStream {
30    const_array(input).unwrap_or_else(|e| TokenStream::from(e.to_compile_error()))
31}
32
33fn const_array(input: TokenStream) -> Result<TokenStream> {
34    let parser = Punctuated::<ConstArray, Token![;]>::parse_terminated;
35    let args = parser.parse(input)?;
36
37    let mut output = proc_macro2::TokenStream::new();
38    for array in args {
39        let len = array.len()?;
40        let ConstArray {
41            attrs,
42            visibility,
43            name,
44            ty,
45            val,
46            span,
47        } = array;
48        output.extend(quote_spanned! {
49            span => #(#attrs)* #visibility const #name: [#ty; #len] = #val;
50        });
51    }
52
53    Ok(TokenStream::from(output))
54}
55
56struct ConstArray {
57    attrs: Vec<Attribute>,
58    visibility: Visibility,
59    name: Ident,
60    ty: Type,
61    val: Expr,
62    span: Span,
63}
64
65impl Parse for ConstArray {
66    fn parse(input: ParseStream) -> Result<Self> {
67        let attrs = input.call(Attribute::parse_outer)?;
68        let visibility: Visibility = input.parse()?;
69        input.parse::<Token![const]>()?;
70        let name: Ident = input.parse()?;
71        input.parse::<Token![:]>()?;
72
73        let left_inner;
74        bracketed!(left_inner in input);
75        let ty = left_inner.parse()?;
76        left_inner.parse::<Token![;]>()?;
77        left_inner.parse::<Token![_]>()?;
78        input.parse::<Token![=]>()?;
79
80        let val: Expr = input.parse()?;
81        let span = input.span();
82        Ok(Self {
83            attrs,
84            visibility,
85            name,
86            ty,
87            val,
88            span,
89        })
90    }
91}
92
93impl ConstArray {
94    fn len(&self) -> Result<proc_macro2::TokenStream> {
95        let array = match self.val.clone() {
96            Expr::Array(array) => array,
97            _ => return Err(syn::Error::new(self.span, "value is not array")),
98        };
99        let mut output = quote_spanned! {
100            self.span =>
101                #[allow(unused_mut)]
102                let mut length = 0;
103        };
104        for expr in array.elems {
105            let attrs = match expr {
106                Expr::Array(inner) => inner.attrs,
107                Expr::Assign(inner) => inner.attrs,
108                // Expr::AssignOp(inner) => inner.attrs,
109                Expr::Async(inner) => inner.attrs,
110                Expr::Await(inner) => inner.attrs,
111                Expr::Binary(inner) => inner.attrs,
112                Expr::Block(inner) => inner.attrs,
113                // Expr::Box(inner) => inner.attrs,
114                Expr::Break(inner) => inner.attrs,
115                Expr::Call(inner) => inner.attrs,
116                Expr::Cast(inner) => inner.attrs,
117                Expr::Closure(inner) => inner.attrs,
118                Expr::Continue(inner) => inner.attrs,
119                Expr::Field(inner) => inner.attrs,
120                Expr::ForLoop(inner) => inner.attrs,
121                Expr::Group(inner) => inner.attrs,
122                Expr::If(inner) => inner.attrs,
123                Expr::Index(inner) => inner.attrs,
124                Expr::Let(inner) => inner.attrs,
125                Expr::Lit(inner) => inner.attrs,
126                Expr::Loop(inner) => inner.attrs,
127                Expr::Macro(inner) => inner.attrs,
128                Expr::Match(inner) => inner.attrs,
129                Expr::MethodCall(inner) => inner.attrs,
130                Expr::Paren(inner) => inner.attrs,
131                Expr::Path(inner) => inner.attrs,
132                Expr::Range(inner) => inner.attrs,
133                Expr::Reference(inner) => inner.attrs,
134                Expr::Repeat(inner) => inner.attrs,
135                Expr::Return(inner) => inner.attrs,
136                Expr::Struct(inner) => inner.attrs,
137                Expr::Try(inner) => inner.attrs,
138                Expr::TryBlock(inner) => inner.attrs,
139                Expr::Tuple(inner) => inner.attrs,
140                // Expr::Type(inner) => inner.attrs,
141                Expr::Unary(inner) => inner.attrs,
142                Expr::Unsafe(inner) => inner.attrs,
143                Expr::While(inner) => inner.attrs,
144                Expr::Yield(inner) => inner.attrs,
145                _ => return Err(syn::Error::new(self.span, "unsupported expr type")),
146            };
147            output.extend(quote_spanned! {
148                self.span =>
149                    #(#attrs)*
150                    {length += 1;}
151            })
152        }
153        Ok(quote_spanned! {
154            self.span => {
155                #output
156                length
157            }
158        })
159    }
160}