Enum spread allocate

crates/storage/derive/src/lib.rs

@@ -137,11 +137,11 @@ synstructure::decl_derive!(
 
 synstructure::decl_derive!(
     [SpreadAllocate] =>
-    /// Derives `ink_storage`'s `SpreadAllocate` trait for the given `struct`.
+    /// Derives `ink_storage`'s `SpreadAllocate` trait for the given `struct` or `enum`.
     ///
     /// # Note
     ///
-    /// As of now `enum` types are not supported!
+    /// As of now `enum` types are only supported if they derive the `Default` trait
     ///
     /// # Examples
     ///

crates/storage/derive/src/spread_allocate.rs

@@ -14,17 +14,17 @@
 
 use proc_macro2::TokenStream as TokenStream2;
 use quote::quote;
+use synstructure::VariantInfo;
 
 /// Derives `ink_storage`'s `SpreadAllocate` trait for the given type.
 pub fn spread_allocate_derive(mut s: synstructure::Structure) -> TokenStream2 {
     s.bind_with(|_| synstructure::BindStyle::Move)
         .add_bounds(synstructure::AddBounds::Generics)
         .underscore_const(true);
+
     match s.ast().data {
         syn::Data::Struct(_) => derive_struct(s),
-        syn::Data::Enum(_) => {
-            panic!("cannot derive `SpreadAllocate` for `enum` types")
-        }
+        syn::Data::Enum(_) => derive_enum_struct(s),
         syn::Data::Union(_) => {
             panic!("cannot derive `SpreadAllocate` for `union` types")
         }
@@ -35,12 +35,8 @@ pub fn spread_allocate_derive(mut s: synstructure::Structure) -> TokenStream2 {
 fn derive_struct(s: synstructure::Structure) -> TokenStream2 {
     assert!(s.variants().len() == 1, "can only operate on structs");
     let variant = &s.variants()[0];
-    let allocate_body = variant.construct(|field, _index| {
-        let ty = &field.ty;
-        quote! {
-            <#ty as ::ink_storage::traits::SpreadAllocate>::allocate_spread(__key_ptr)
-        }
-    });
+
+    let allocate_body = allocate_body(variant);
     s.gen_impl(quote! {
         gen impl ::ink_storage::traits::SpreadAllocate for @Self {
             fn allocate_spread(__key_ptr: &mut ::ink_primitives::KeyPtr) -> Self {
@@ -49,3 +45,40 @@ fn derive_struct(s: synstructure::Structure) -> TokenStream2 {
         }
     })
 }
+
+/// Derives `ink_storage`'s `SpreadAllocate` trait for the given `enum`.
+fn derive_enum_struct(s: synstructure::Structure) -> TokenStream2 {
+    let default_index = search_variants_for_default(s.variants());
+    let variant = &s.variants()[default_index];
+    let allocate_body = allocate_body(variant);
+    s.gen_impl(quote! {
+        gen impl ::ink_storage::traits::SpreadAllocate for @Self
+            where Self:Default {
+                fn allocate_spread(__key_ptr: &mut ::ink_primitives::KeyPtr) -> Self {
+                    #allocate_body
+            }
+        }
+    })
+}
+
+fn allocate_body(variant: &VariantInfo) -> TokenStream2 {
+    variant.construct(|field, _index| {
+        let ty = &field.ty;
+        quote! {
+            <#ty as ::ink_storage::traits::SpreadAllocate>::allocate_spread(__key_ptr)
+        }
+    })
+}
+
+fn search_variants_for_default(variants: &[VariantInfo]) -> usize {
+    for (index, variant) in variants.iter().enumerate() {
+        for attr in variant.ast().attrs {
+            for segment in &attr.path.segments {
+                if segment.ident == "default" {
+                    return index
+                }
+            }
+        }
+    }
+    panic!("can only derive `SpreadAllocate` for `enum` types that implement `Default`")
+}

crates/storage/derive/src/tests/spread_allocate.rs

@@ -14,13 +14,47 @@
 
 use crate::spread_allocate_derive;
 use syn::parse_quote;
+use synstructure::{
+    quote,
+};
 
 #[test]
-#[should_panic(expected = "cannot derive `SpreadAllocate` for `enum` types")]
-fn enum_fails() {
-    spread_allocate_derive(synstructure::Structure::new(&parse_quote! {
+#[should_panic(
+    expected = "can only derive `SpreadAllocate` for `enum` types that implement `Default`"
+)]
+fn enum_without_default_fails() {
+    let parsed = &parse_quote! { #[derive(Default)]
+        //enum Enum { A, B, #[default] C }
         enum Enum { A, B, C }
-    }));
+    };
+
+    let mut s = synstructure::Structure::new(parsed);
+
+    s.add_where_predicate(syn::parse_quote!(Self: Default));
+
+    spread_allocate_derive(s);
+}
+
+#[test]
+fn enum_with_default_works() {
+    let parsed = &parse_quote! { #[derive(Default)]
+        enum Enum { A, B, #[default] C }
+    };
+
+    let s = synstructure::Structure::new(parsed);
+    let result = spread_allocate_derive(s);
+    assert_eq!(result.to_string(),
+    quote! {
+        const _ : () = {
+            impl::ink_storage::traits::SpreadAllocate for Enum
+            where Self : Default {
+                fn allocate_spread (__key_ptr : & mut :: ink_primitives :: KeyPtr) -> Self {
+                    Enum :: C
+                }
+            }
+        };
+    }.to_string()
+    );
 }
 
 #[test]