Skip to main content

rotary_permutator_derive/
lib.rs

1extern crate proc_macro;
2extern crate syn;
3#[macro_use]
4extern crate quote;
5
6use proc_macro::TokenStream;
7
8#[proc_macro_derive(EnumRotor)]
9pub fn derive_rotor(input: TokenStream) -> TokenStream {
10    let input = syn::parse_macro_input!(input as syn::DeriveInput);
11    let enum_ident = &input.ident;
12
13    let rotor_engine_struct = quote! {
14        pub struct RotorEngine<T> {
15            blocks: Vec<T>,
16            trigger: Option<usize>,
17            final_trigger: bool,
18        }
19    };
20
21    let init_rotor_engine = gen_init_rotor_engine_impl(&input);
22    let rotor_engine_rot_at = gen_rotor_engine_rot_at(&input);
23    let rotor_enum_impl = quote! {
24        impl #enum_ident {
25            #init_rotor_engine
26            #rotor_engine_rot_at
27        }
28
29    };
30
31    let iterator_for_rotor_engine = gen_iterator_for_rotator_engine(&input);
32
33    quote! {
34        #rotor_engine_struct
35        #rotor_enum_impl
36        #iterator_for_rotor_engine
37
38    }
39    .into()
40}
41
42fn gen_init_rotor_engine_impl(ast: &syn::DeriveInput) -> proc_macro2::TokenStream {
43    let enum_ident = &ast.ident;
44    quote! {
45        pub fn init_rotor_engine(size: usize) -> RotorEngine<#enum_ident> {
46            let mut blocks = Vec::new();
47            for _i in 0..size{
48                blocks.push(#enum_ident::default());
49            }
50            RotorEngine {
51                blocks,
52                trigger: None,
53                final_trigger: false,
54            }
55        }
56    }
57}
58
59fn gen_rotor_engine_rot_at(ast: &syn::DeriveInput) -> proc_macro2::TokenStream {
60    let enum_ident = &ast.ident;
61
62    let mut match_elems_impl = quote! {};
63    let (first_elem, last_elem) = match &ast.data {
64        syn::Data::Enum(syn::DataEnum { variants, .. }) => {
65            let mut last_elem = &variants[0].ident;
66            for i in 0..variants.len() {
67                if i != variants.len() - 1 {
68                    let current_elem = &variants[i + 1].ident;
69                    let stmt = quote! {
70                        #enum_ident::#last_elem => rotor_engine.blocks[index] = #enum_ident::#current_elem,
71                    };
72                    match_elems_impl.extend(stmt);
73                    last_elem = current_elem;
74                }
75            }
76            (&variants[0].ident, &variants[variants.len() - 1].ident)
77        }
78        _ => panic!("only enums supported"),
79    };
80
81    quote! {
82        fn rot_at(rotor_engine: &mut RotorEngine<#enum_ident>, index: usize) {
83            match rotor_engine.blocks[index] {
84                #match_elems_impl
85                #enum_ident::#last_elem => {
86
87                    if index > 0 {
88                        rotor_engine.trigger = Some(index - 1);
89                        rotor_engine.blocks[index] = #enum_ident::#first_elem;
90                    } else {
91                        rotor_engine.final_trigger = true;
92                    }
93                }
94            }
95        }
96    }
97}
98
99fn gen_iterator_for_rotator_engine(ast: &syn::DeriveInput) -> proc_macro2::TokenStream {
100    let enum_ident = &ast.ident;
101    quote! {
102        impl std::iter::Iterator for RotorEngine<#enum_ident> {
103            type Item = Vec<#enum_ident>;
104            fn next(&mut self) -> Option<Self::Item> {
105
106                while let Some(idx) = self.trigger {
107                    self.trigger = None;
108                    #enum_ident::rot_at(self, idx);
109                }
110                let capture = self.blocks.clone();
111                #enum_ident::rot_at(self, self.blocks.len() - 1);
112
113                if self.final_trigger {
114                    return None;
115                }
116                Some(capture)
117            }
118        }
119    }
120}