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#[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#[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::Async(inner) => inner.attrs,
133 Expr::Await(inner) => inner.attrs,
134 Expr::Binary(inner) => inner.attrs,
135 Expr::Block(inner) => inner.attrs,
136 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::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}