array_const_fn_init/
lib.rs

1//! Initializes an array with constant values calculated by a `const fn`
2//!
3//! # Examples
4//!
5//! ```
6//! use array_const_fn_init::array_const_fn_init;
7//!
8//! const fn const_double_it(i: usize) -> usize {
9//!     i * 2
10//! }
11//! const ARRAY: [usize; 10] = array_const_fn_init![const_double_it; 10];
12//! assert_eq!(ARRAY, [0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
13//! ```
14//!
15//! ```
16//! use array_const_fn_init::array_const_fn_init;
17//!
18//! const fn const_vecs(i: usize) -> (u8, u8, u8) {
19//!     (i as u8, i as u8, i as u8)
20//! }
21//! const ARRAY: [(u8, u8, u8); 4] = array_const_fn_init![const_vecs; 4];
22//! assert_eq!(ARRAY, [(0, 0, 0), (1, 1, 1), (2, 2, 2), (3, 3, 3)]);
23//! ```
24extern crate proc_macro;
25
26use core::iter;
27use core::str::FromStr;
28use proc_macro::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream, TokenTree};
29
30#[proc_macro]
31pub fn array_const_fn_init(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
32    let mut iter = input.into_iter();
33
34    let func_name = match iter.next() {
35        Some(proc_macro::TokenTree::Ident(i)) => i.to_string(),
36        Some(other) => panic!("Expected function name, found {}", other),
37        None => panic!("Unexpected end of macro input"),
38    };
39
40    match iter.next() {
41        Some(proc_macro::TokenTree::Punct(ref p)) if p.as_char() == ';' => {}
42        Some(other) => panic!("Expected ';', found {}", other),
43        None => panic!("Unexpected end of macro input"),
44    };
45
46    let array_size: usize = match iter.next() {
47        Some(proc_macro::TokenTree::Literal(ref p)) => {
48            usize::from_str(&p.to_string()).expect("Expected <usize>")
49        }
50        Some(other) => panic!("Expected <usize>, found {}", other),
51        None => panic!("Unexpected end of macro input"),
52    };
53
54    match iter.next() {
55        None => {}
56        Some(_) => panic!("Unexpected trailing tokens in macro"),
57    }
58
59    let mut ts = TokenStream::new();
60    let span: Span = Span::call_site();
61    ts.extend({
62        let mut g = Group::new(Delimiter::Bracket, {
63            let mut ts = TokenStream::new();
64            (0..array_size).for_each(|i| {
65                ts.extend(iter::once(TokenTree::from(Ident::new(&func_name, span))));
66                ts.extend({
67                    let mut g = Group::new(Delimiter::Parenthesis, {
68                        let mut ts = TokenStream::new();
69                        let _span: Span = span;
70                        ts.extend(iter::once(TokenTree::from(Literal::usize_suffixed(i))));
71                        ts
72                    });
73                    g.set_span(span);
74                    Some(TokenTree::from(g))
75                });
76                ts.extend(iter::once(TokenTree::from(Punct::new(',', Spacing::Alone))));
77            });
78            ts
79        });
80        g.set_span(span);
81        Some(TokenTree::from(g))
82    });
83    ts
84}
array_const_fn_init/lib.rs

array_const_fn_init/
lib.rs
