cell_map_macro/
lib.rs

1//! # `cell-map` Macros
2//!
3//! This crate provides implementations of macros for the [`cell_map`](todo) crate.
4
5// ------------------------------------------------------------------------------------------------
6// IMPORTS
7// ------------------------------------------------------------------------------------------------
8
9use proc_macro::TokenStream;
10use quote::quote;
11use syn::{parse_macro_input, DeriveInput};
12
13// ------------------------------------------------------------------------------------------------
14// DERIVES
15// ------------------------------------------------------------------------------------------------
16
17#[proc_macro_derive(Layer)]
18pub fn derive_layer(input: TokenStream) -> TokenStream {
19    let input = parse_macro_input!(input as DeriveInput);
20
21    // Check input is an enum
22    let variants = match input.data {
23        syn::Data::Enum(e) => e.variants,
24        _ => panic!("Layer can only be derived on enums"),
25    };
26
27    // Get the type name
28    let name = &input.ident;
29
30    // Map the varients into the match patterns we need for the to_index function
31    let var_to_index_patterns = variants.iter().enumerate().map(|(i, v)| {
32        let var_name = &v.ident;
33
34        quote! {
35            #name::#var_name => #i
36        }
37    });
38
39    // Map the varients into the match patterns we need for the from_index function
40    let var_from_index_patterns = variants.iter().enumerate().map(|(i, v)| {
41        let var_name = &v.ident;
42
43        quote! {
44            #i => #name::#var_name
45        }
46    });
47
48    let var_all_patterns = variants.iter().map(|v| {
49        let var_name = &v.ident;
50
51        quote! {
52            #name::#var_name
53        }
54    });
55
56    let first_var_name = &variants[0].ident;
57
58    let num_variants = variants.len();
59
60    let impled = quote! {
61        impl ::cell_map::Layer for #name {
62            const NUM_LAYERS: usize = #num_variants;
63
64            const FIRST: Self = Self::#first_var_name;
65
66            fn to_index(&self) -> usize {
67                match self {
68                    #(#var_to_index_patterns),*
69                }
70            }
71
72            fn from_index(index: usize) -> Self {
73                match index {
74                    #(#var_from_index_patterns),*,
75                    _ => panic!("Got a layer index of {} but there are only {} layers", index, #num_variants)
76                }
77            }
78
79            fn all() -> Vec<Self> {
80                vec![#(#var_all_patterns),*]
81            }
82        }
83    };
84
85    impled.into()
86}