From e37e43c3d13a85df39f3f76fd486e7582ecf5e6d Mon Sep 17 00:00:00 2001
From: steven-omaha <35634100+steven-omaha@users.noreply.github.com>
Date: Thu, 12 Jan 2023 23:27:32 +0100
Subject: [PATCH] clean up procedural macro

---
 pacdef_macro/src/lib.rs | 112 +++++++++++++++++++++++++++-------------
 1 file changed, 77 insertions(+), 35 deletions(-)

diff --git a/pacdef_macro/src/lib.rs b/pacdef_macro/src/lib.rs
index fd3d881..698cdea 100644
--- a/pacdef_macro/src/lib.rs
+++ b/pacdef_macro/src/lib.rs
@@ -1,39 +1,74 @@
 use proc_macro::TokenStream;
 use quote::quote;
 use syn::DeriveInput;
+use syn::__private::TokenStream2;
 
 #[proc_macro_derive(Register)]
 pub fn register(input: TokenStream) -> TokenStream {
     let input = syn::parse::<DeriveInput>(input).unwrap();
+
     let name = &input.ident;
+
     let enum_data = if let syn::Data::Enum(enum_data) = &input.data {
         enum_data
     } else {
         panic!("`Register` can only be used on enums");
     };
+
     let first_variant = &enum_data.variants[0].ident;
 
-    let variant_matches_backend = enum_data.variants.iter().map(|variant| {
-        let variant_name = &variant.ident;
-        quote! {
-            Self::#variant_name => Box::new(#variant_name::new()),
-        }
-    });
+    let variant_matches_backend = generate_variant_backend(enum_data);
+    let variant_matches_next = generate_variant_matches_next(enum_data);
+    let variant_imports = generate_variant_imports(enum_data);
 
-    let variant_matches_next = enum_data.variants.iter().enumerate().map(|(i, variant)| {
-        let variant_name = &variant.ident;
-        if i == enum_data.variants.len() - 1 {
-            quote! {
-                Self::#variant_name => None,
+    let expanded = compile_output(
+        name,
+        first_variant,
+        variant_matches_backend,
+        variant_matches_next,
+        variant_imports,
+    );
+
+    TokenStream::from(expanded)
+}
+
+fn compile_output<T, U, V>(
+    name: &syn::Ident,
+    first_variant: &syn::Ident,
+    variant_backend: T,
+    variant_next: U,
+    variant_imports: V,
+) -> TokenStream2
+where
+    T: Iterator<Item = TokenStream2>,
+    U: Iterator<Item = TokenStream2>,
+    V: Iterator<Item = TokenStream2>,
+{
+    let expanded = quote! {
+        #(#variant_imports)*
+
+        impl #name {
+            pub fn iter() -> BackendIter {
+                BackendIter {
+                    next: Some(Self::#first_variant),
+                }
             }
-        } else {
-            let next_variant = &enum_data.variants[i + 1].ident;
-            quote! {
-                Self::#variant_name => Some(Self::#next_variant),
+            fn get_backend(&self) -> Box<dyn Backend> {
+                match self {
+                    #(#variant_backend)*
+                }
+            }
+            fn next(&self) -> Option<Self> {
+                match self {
+                    #(#variant_next)*
+                }
             }
         }
-    });
+    };
+    expanded
+}
 
+fn generate_variant_imports(enum_data: &syn::DataEnum) -> impl Iterator<Item = TokenStream2> + '_ {
     let variant_imports = enum_data.variants.iter().map(|variant| {
         let variant_name = &variant.ident;
         // let variant_module = variant_name.clone();
@@ -48,27 +83,34 @@ pub fn register(input: TokenStream) -> TokenStream {
 
         }
     });
+    variant_imports
+}
 
-    let expanded = quote! {
-        #(#variant_imports)*
-
-        impl #name {
-            pub fn iter() -> BackendIter {
-                BackendIter {
-                    next: Some(Self::#first_variant),
-                }
+fn generate_variant_matches_next(
+    enum_data: &syn::DataEnum,
+) -> impl Iterator<Item = TokenStream2> + '_ {
+    let variant_matches_next = enum_data.variants.iter().enumerate().map(|(i, variant)| {
+        let variant_name = &variant.ident;
+        if i == enum_data.variants.len() - 1 {
+            quote! {
+                Self::#variant_name => None,
             }
-            fn get_backend(&self) -> Box<dyn Backend> {
-                match self {
-                    #(#variant_matches_backend)*
-                }
-            }
-            fn next(&self) -> Option<Self> {
-                match self {
-                    #(#variant_matches_next)*
-                }
+        } else {
+            let next_variant = &enum_data.variants[i + 1].ident;
+            quote! {
+                Self::#variant_name => Some(Self::#next_variant),
             }
         }
-    };
-    TokenStream::from(expanded)
+    });
+    variant_matches_next
+}
+
+fn generate_variant_backend(enum_data: &syn::DataEnum) -> impl Iterator<Item = TokenStream2> + '_ {
+    let variant_matches_backend = enum_data.variants.iter().map(|variant| {
+        let variant_name = &variant.ident;
+        quote! {
+            Self::#variant_name => Box::new(#variant_name::new()),
+        }
+    });
+    variant_matches_backend
 }