From 5779a3cdaf74c47e984d452e9eb0f5dddf6990d3 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Sun, 24 Oct 2021 19:58:41 +0200
Subject: [PATCH] feat: add ethabitype support for solidity style enums (#526)

* feat: tokenize solidity like enums

* test: add enum test

* rustfmt
---
 .../ethers-contract-derive/src/abi_ty.rs      | 29 ++++++++++++++-----
 ethers-contract/tests/common/derive.rs        | 18 ++++++++++++
 2 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/ethers-contract/ethers-contract-derive/src/abi_ty.rs b/ethers-contract/ethers-contract-derive/src/abi_ty.rs
index 06e38257..0bebc8c5 100644
--- a/ethers-contract/ethers-contract-derive/src/abi_ty.rs
+++ b/ethers-contract/ethers-contract-derive/src/abi_ty.rs
@@ -1,7 +1,7 @@
 //! Helper functions for deriving `EthAbiType`
 
 use ethers_contract_abigen::ethers_core_crate;
-use proc_macro2::{Ident, TokenStream};
+use proc_macro2::{Ident, Literal, TokenStream};
 use quote::{quote, quote_spanned};
 use syn::spanned::Spanned as _;
 use syn::{parse::Error, Data, DeriveInput, Fields, Variant};
@@ -198,6 +198,12 @@ fn tokenize_unit_type(name: &Ident) -> TokenStream {
     }
 }
 
+/// Derive for an enum
+///
+/// An enum can be a [solidity enum](https://docs.soliditylang.org/en/v0.5.3/types.html#enums) or a
+/// bundled set of different types.
+///
+/// Decoding works like untagged decoding
 fn tokenize_enum<'a>(
     enum_name: &Ident,
     variants: impl Iterator<Item = &'a Variant> + 'a,
@@ -206,15 +212,24 @@ fn tokenize_enum<'a>(
 
     let mut into_tokens = TokenStream::new();
     let mut from_tokens = TokenStream::new();
-    for variant in variants {
+    for (idx, variant) in variants.into_iter().enumerate() {
+        let var_ident = &variant.ident;
         if variant.fields.len() > 1 {
             return Err(Error::new(
                 variant.span(),
                 "EthAbiType cannot be derived for enum variants with multiple fields",
             ));
-        }
-        let var_ident = &variant.ident;
-        if let Some(field) = variant.fields.iter().next() {
+        } else if variant.fields.is_empty() {
+            let value = Literal::u8_unsuffixed(idx as u8);
+            from_tokens.extend(quote! {
+                 if let Ok(#value) = u8::from_token(token.clone()) {
+                    return Ok(#enum_name::#var_ident)
+                }
+            });
+            into_tokens.extend(quote! {
+                 #enum_name::#var_ident => #value.into_token(),
+            });
+        } else if let Some(field) = variant.fields.iter().next() {
             let ty = &field.ty;
             from_tokens.extend(quote! {
                 if let Ok(decoded) = #ty::from_token(token.clone()) {
@@ -226,8 +241,8 @@ fn tokenize_enum<'a>(
             });
         } else {
             into_tokens.extend(quote! {
-                 #enum_name::#var_ident(element) => # ethers_core::abi::Token::Tuple(::std::vec::Vec::new()),
-            });
+             #enum_name::#var_ident(element) => # ethers_core::abi::Token::Tuple(::std::vec::Vec::new()),
+        });
         }
     }
 
diff --git a/ethers-contract/tests/common/derive.rs b/ethers-contract/tests/common/derive.rs
index f483d8db..bcb529ec 100644
--- a/ethers-contract/tests/common/derive.rs
+++ b/ethers-contract/tests/common/derive.rs
@@ -489,3 +489,21 @@ fn can_derive_ethcall_with_nested_structs() {
     assert_tokenizeable::<FooCall>();
     assert_ethcall::<FooCall>();
 }
+
+#[test]
+fn can_derive_for_enum() {
+    #[derive(Debug, Clone, PartialEq, EthAbiType)]
+    enum ActionChoices {
+        GoLeft,
+        GoRight,
+        GoStraight,
+        SitStill,
+    }
+    assert_tokenizeable::<ActionChoices>();
+
+    let token = ActionChoices::GoLeft.into_token();
+    assert_eq!(
+        ActionChoices::GoLeft,
+        ActionChoices::from_token(token).unwrap()
+    );
+}