fill_array/
lib.rs

1extern crate proc_macro;
2extern crate syn;
3extern crate quote;
4
5use proc_macro::TokenStream;
6use proc_macro2::{Ident, TokenStream as TokenStream2};
7use quote::quote;
8use syn::{parse_macro_input, Token, parse::{ParseStream, Parse}, Result, LitInt, Expr};
9
10#[allow(unused)]
11enum Length {
12    Literal(LitInt),
13    Constant(Ident),
14}
15
16impl Parse for Length {
17    fn parse(input: ParseStream) -> Result<Self> {
18        if input.peek(LitInt) {
19            let lit_int: LitInt = input.parse()?;
20            Ok(Length::Literal(lit_int))
21        } else {
22            let ident: Ident = input.parse()?;
23            Ok(Length::Constant(ident))
24        }
25    }
26}
27
28#[allow(unused)]
29struct FillInput {
30    item: Expr,
31    semi_token: Token![;],
32    length: Length,
33}
34
35impl Parse for FillInput {
36    fn parse(input: ParseStream) -> Result<Self> {
37        let item = input.parse()?;
38        let semi_token = input.parse()?;
39        let length = input.parse()?;
40        Ok(FillInput { item, semi_token, length })
41    }
42}
43
44#[proc_macro]
45pub fn fill(input: TokenStream) -> TokenStream {
46    let FillInput { item, length, .. } = parse_macro_input!(input as FillInput);
47
48    match length {
49        Length::Literal(lit) => {
50            let length = lit.base10_parse().expect("Incorrect length");
51
52            let mut inner_tokens = TokenStream2::new();
53            for _ in 0..length {
54                inner_tokens.extend(quote!{#item,})
55            }
56
57            let tokens = quote! {
58                [#inner_tokens]
59            };
60
61            tokens.into()
62        },
63        Length::Constant(ident) => {
64            let tokens = quote! {
65                core::array::from_fn::<_, #ident, _>(|_| #item)
66            };
67
68            tokens.into()
69        },
70    }
71}
fill_array/lib.rs

fill_array/
lib.rs
