From 59750612a2845603913dcb68ef969cb97fbfbcd6 Mon Sep 17 00:00:00 2001
From: Hirrolot <hirrolot@gmail.com>
Date: Thu, 3 Feb 2022 03:10:58 +0600
Subject: [PATCH] Support multiple fields in `DialogueState`

---
 src/dialogue_state.rs | 83 +++++++++++++++++++++++++++++++++++--------
 1 file changed, 69 insertions(+), 14 deletions(-)

diff --git a/src/dialogue_state.rs b/src/dialogue_state.rs
index fa74e3ae..bd626bd9 100644
--- a/src/dialogue_state.rs
+++ b/src/dialogue_state.rs
@@ -1,9 +1,9 @@
 use proc_macro2::{Ident, Span, TokenStream};
-use quote::{quote, ToTokens};
+use quote::{format_ident, quote, ToTokens};
 use syn::{
     parse::{Parse, ParseStream},
     spanned::Spanned,
-    Fields, GenericParam, ItemEnum, Path, Type,
+    Fields, FieldsNamed, GenericParam, ItemEnum, Path, Type,
 };
 
 pub fn expand(item: ItemEnum) -> Result<TokenStream, syn::Error> {
@@ -64,12 +64,13 @@ pub fn expand(item: ItemEnum) -> Result<TokenStream, syn::Error> {
         };
 
         branches.extend(match &variant.fields {
-            Fields::Named(_) => {
-                return Err(syn::Error::new(
-                    variant.span(),
-                    "Named fields does not allowed",
-                ))
-            }
+            Fields::Named(fields) => create_branch_multiple_fields_named(
+                &enum_ident,
+                &self_params,
+                &variant.ident,
+                &handler.func,
+                fields,
+            ),
             Fields::Unnamed(fields) => match fields.unnamed.len() {
                 1 => create_branch_one_field(
                     &enum_ident,
@@ -77,12 +78,13 @@ pub fn expand(item: ItemEnum) -> Result<TokenStream, syn::Error> {
                     &variant.ident,
                     &handler.func,
                 ),
-                len => {
-                    return Err(syn::Error::new(
-                        fields.span(),
-                        format!("Expected 1 field, found {}", len),
-                    ));
-                }
+                len => create_branch_multiple_fields(
+                    &enum_ident,
+                    &self_params,
+                    &variant.ident,
+                    &handler.func,
+                    len,
+                ),
             },
             Fields::Unit => create_branch_no_fields(
                 &enum_ident,
@@ -142,6 +144,59 @@ fn create_branch_one_field(
     }
 }
 
+fn create_branch_multiple_fields(
+    state: &Ident,
+    state_generics: impl ToTokens,
+    kind: &Ident,
+    handler: &Path,
+    fields_count: usize,
+) -> TokenStream {
+    let fields = gen_variant_field_names(fields_count);
+
+    quote! {
+        .branch(
+            dptree::filter_map(|state: #state #state_generics| async move {
+                match state { #state::#kind(#fields) => Some((#fields)), _ => None }
+            }).endpoint(#handler)
+        )
+    }
+}
+
+fn gen_variant_field_names(len: usize) -> TokenStream {
+    let mut fields = quote! {};
+
+    for i in 0..len {
+        let idx = format_ident!("_{}", i);
+        fields.extend(quote! { #idx, });
+    }
+
+    return fields;
+}
+
+fn create_branch_multiple_fields_named(
+    state: &Ident,
+    state_generics: impl ToTokens,
+    kind: &Ident,
+    handler: &Path,
+    fields_named: &FieldsNamed,
+) -> TokenStream {
+    let mut fields = quote! {};
+
+    for field in fields_named.named.iter() {
+        let ident =
+            field.ident.as_ref().expect("Named fields must have identifiers");
+        fields.extend(quote! { #ident, });
+    }
+
+    quote! {
+        .branch(
+            dptree::filter_map(|state: #state #state_generics| async move {
+                match state { #state::#kind { #fields } => Some((#fields)), _ => None }
+            }).endpoint(#handler)
+        )
+    }
+}
+
 fn parse_out_type(
     span: Span,
     attrs: &[syn::Attribute],