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
33/// Declare a new const array without specify length.
34/// It helps when apply conditional compilation to part of a const array.
35/// Similar to [`auto_const_array`], but using attribute macro syntax.
36///
37/// ```rust
38/// use auto_const_array::auto_const_array_attr as auto_const_array;
39///
40/// /// Common array with public visibility.
41/// #[auto_const_array]
42/// pub const ARRAY_COMMON: [u8; _] = [1, 2, 4];
43/// /// Special array with cfg conditional compiling.
44/// #[auto_const_array]
45/// const ARRAY_WITH_ATTR: [u8; _] = [
46///     1,
47///     #[cfg(unix)]
48///     2,
49/// ];
50/// ```
51#[proc_macro_attribute]
52pub fn auto_const_array_attr(_attr: TokenStream, item: TokenStream) -> TokenStream {
53    const_array(item).unwrap_or_else(|e| TokenStream::from(e.to_compile_error()))
54}
55
56fn const_array(input: TokenStream) -> Result<TokenStream> {
57    let parser = Punctuated::<ConstArray, Token![;]>::parse_terminated;
58    let args = parser.parse(input)?;
59
60    let mut output = proc_macro2::TokenStream::new();
61    for array in args {
62        let len = array.len()?;
63        let ConstArray {
64            attrs,
65            visibility,
66            name,
67            ty,
68            val,
69            span,
70        } = array;
71        output.extend(quote_spanned! {
72            span => #(#attrs)* #visibility const #name: [#ty; #len] = #val;
73        });
74    }
75
76    Ok(TokenStream::from(output))
77}
78
79struct ConstArray {
80    attrs: Vec<Attribute>,
81    visibility: Visibility,
82    name: Ident,
83    ty: Type,
84    val: Expr,
85    span: Span,
86}
87
88impl Parse for ConstArray {
89    fn parse(input: ParseStream) -> Result<Self> {
90        let attrs = input.call(Attribute::parse_outer)?;
91        let visibility: Visibility = input.parse()?;
92        input.parse::<Token![const]>()?;
93        let name: Ident = input.parse()?;
94        input.parse::<Token![:]>()?;
95
96        let left_inner;
97        bracketed!(left_inner in input);
98        let ty = left_inner.parse()?;
99        left_inner.parse::<Token![;]>()?;
100        left_inner.parse::<Token![_]>()?;
101        input.parse::<Token![=]>()?;
102
103        let val: Expr = input.parse()?;
104        let span = input.span();
105        Ok(Self {
106            attrs,
107            visibility,
108            name,
109            ty,
110            val,
111            span,
112        })
113    }
114}
115
116impl ConstArray {
117    fn len(&self) -> Result<proc_macro2::TokenStream> {
118        let array = match self.val.clone() {
119            Expr::Array(array) => array,
120            _ => return Err(syn::Error::new(self.span, "value is not array")),
121        };
122        let mut output = quote_spanned! {
123            self.span =>
124                #[allow(unused_mut)]
125                let mut length = 0;
126        };
127        for expr in array.elems {
128            let attrs = match expr {
129                Expr::Array(inner) => inner.attrs,
130                Expr::Assign(inner) => inner.attrs,
131                // Expr::AssignOp(inner) => inner.attrs,
132                Expr::Async(inner) => inner.attrs,
133                Expr::Await(inner) => inner.attrs,
134                Expr::Binary(inner) => inner.attrs,
135                Expr::Block(inner) => inner.attrs,
136                // Expr::Box(inner) => inner.attrs,
137                Expr::Break(inner) => inner.attrs,
138                Expr::Call(inner) => inner.attrs,
139                Expr::Cast(inner) => inner.attrs,
140                Expr::Closure(inner) => inner.attrs,
141                Expr::Continue(inner) => inner.attrs,
142                Expr::Field(inner) => inner.attrs,
143                Expr::ForLoop(inner) => inner.attrs,
144                Expr::Group(inner) => inner.attrs,
145                Expr::If(inner) => inner.attrs,
146                Expr::Index(inner) => inner.attrs,
147                Expr::Let(inner) => inner.attrs,
148                Expr::Lit(inner) => inner.attrs,
149                Expr::Loop(inner) => inner.attrs,
150                Expr::Macro(inner) => inner.attrs,
151                Expr::Match(inner) => inner.attrs,
152                Expr::MethodCall(inner) => inner.attrs,
153                Expr::Paren(inner) => inner.attrs,
154                Expr::Path(inner) => inner.attrs,
155                Expr::Range(inner) => inner.attrs,
156                Expr::Reference(inner) => inner.attrs,
157                Expr::Repeat(inner) => inner.attrs,
158                Expr::Return(inner) => inner.attrs,
159                Expr::Struct(inner) => inner.attrs,
160                Expr::Try(inner) => inner.attrs,
161                Expr::TryBlock(inner) => inner.attrs,
162                Expr::Tuple(inner) => inner.attrs,
163                // Expr::Type(inner) => inner.attrs,
164                Expr::Unary(inner) => inner.attrs,
165                Expr::Unsafe(inner) => inner.attrs,
166                Expr::While(inner) => inner.attrs,
167                Expr::Yield(inner) => inner.attrs,
168                _ => return Err(syn::Error::new(self.span, "unsupported expr type")),
169            };
170            output.extend(quote_spanned! {
171                self.span =>
172                    #(#attrs)*
173                    {length += 1;}
174            })
175        }
176        Ok(quote_spanned! {
177            self.span => {
178                #output
179                length
180            }
181        })
182    }
183}