Add StorageKeys derive macro

macros/Cargo.toml

@@ -10,5 +10,9 @@ version = "0.12.0"
 proc-macro = true
 
 [dependencies]
+proc-macro2 = "1.0"
 quote = "1.0"
-syn = "1.0"
+syn = {version="1.0", features = ["full", "extra-traits"]}
+
+[dev-dependencies]
+itertools = "0.10.1"

macros/src/lib.rs

@@ -7,8 +7,10 @@
 #![deny(rustdoc::private_intra_doc_links)]
 
 use proc_macro::TokenStream;
-use quote::quote;
-use syn::{parse_macro_input, ItemFn};
+use proc_macro2::{Span as Span2, TokenStream as TokenStream2};
+use quote::{quote, ToTokens};
+use syn::punctuated::Punctuated;
+use syn::{parse_macro_input, ItemFn, ItemStruct};
 
 /// Generate WASM binding for a transaction main entrypoint function.
 ///
@@ -149,3 +151,209 @@ pub fn validity_predicate(
     };
     TokenStream::from(gen)
 }
+
+#[proc_macro_derive(StorageKeys)]
+pub fn derive_storage_keys(struct_def: TokenStream) -> TokenStream {
+    derive_storage_keys_inner(struct_def.into()).into()
+}
+
+#[inline]
+// TODO: use this crate for errors: https://crates.io/crates/proc-macro-error
+// TODO: emit allow_deadcode in ALL and VALUES
+fn derive_storage_keys_inner(struct_def: TokenStream2) -> TokenStream2 {
+    let struct_def: ItemStruct = syn::parse2(struct_def)
+        .expect("Expected a struct in the StorageKeys derive");
+
+    // type check the struct - all fields must be of type `&'static str`
+    let fields = match &struct_def.fields {
+        syn::Fields::Named(fields) => &fields.named,
+        _ => panic!(
+            "Only named struct fields are accepted in StorageKeys derives"
+        ),
+    };
+
+    let mut idents = vec![];
+
+    for field in fields {
+        let field_type = field.ty.to_token_stream().to_string();
+        if field_type != "& 'static str" {
+            panic!(
+                "Expected `&'static str` field type in StorageKeys derive, \
+                 but got `{field_type}` instead"
+            );
+        }
+        idents.push(field.ident.clone().expect("Expected a named field"));
+    }
+
+    idents.sort();
+
+    let ident_list = create_punctuated(&idents, |ident| ident.clone());
+    let values_list = create_punctuated(&idents, |ident| {
+        let storage_key = {
+            let mut toks = TokenStream2::new();
+            ident.to_tokens(&mut toks);
+            toks.to_string()
+        };
+        syn::FieldValue {
+            attrs: vec![],
+            member: syn::Member::Named(ident.clone()),
+            colon_token: Some(syn::token::Colon {
+                spans: [Span2::call_site()],
+            }),
+            expr: syn::Expr::Lit(syn::ExprLit {
+                attrs: vec![],
+                lit: syn::Lit::Str(syn::LitStr::new(
+                    storage_key.as_str(),
+                    Span2::call_site(),
+                )),
+            }),
+        }
+    });
+
+    let struct_def_ident = &struct_def.ident;
+
+    quote! {
+        impl #struct_def_ident {
+            const ALL: &[&'static str] = {
+                let #struct_def_ident {
+                    #ident_list
+                } = Self::VALUES;
+
+                &[ #ident_list ]
+            };
+
+            const VALUES: #struct_def_ident = Self {
+                #values_list
+            };
+        }
+    }
+}
+
+#[inline]
+fn create_punctuated<F, M>(
+    idents: &[syn::Ident],
+    mut map: F,
+) -> Punctuated<M, syn::token::Comma>
+where
+    F: FnMut(&syn::Ident) -> M,
+{
+    idents.iter().fold(Punctuated::new(), |mut accum, ident| {
+        accum.push(map(ident));
+        accum
+    })
+}
+
+#[cfg(test)]
+mod test_proc_macros {
+    use itertools::Itertools;
+    use syn::ItemImpl;
+
+    use super::*;
+
+    /// Test if the `ALL` slice generated in `StorageKeys` macro
+    /// derives is sorted in ascending order.
+    #[test]
+    fn test_storage_keys_derive_sorted_slice() {
+        let test_struct = quote! {
+            struct Keys {
+                word: &'static str,
+                is: &'static str,
+                bird: &'static str,
+                the: &'static str,
+            }
+        };
+        let all = {
+            let test_impl: ItemImpl =
+                syn::parse2(derive_storage_keys_inner(test_struct))
+                    .expect("Test failed");
+            test_impl
+                .items
+                .iter()
+                .find_map(|i| match i {
+                    syn::ImplItem::Const(e)
+                        if e.ident.to_token_stream().to_string() == "ALL" =>
+                    {
+                        match &e.expr {
+                            syn::Expr::Block(e) => {
+                                match e
+                                    .block
+                                    .stmts
+                                    .last()
+                                    .expect("Must have slice")
+                                {
+                                    syn::Stmt::Expr(syn::Expr::Reference(
+                                        e,
+                                    )) => match &*e.expr {
+                                        syn::Expr::Array(e) => Some(e.clone()),
+                                        t => panic!(
+                                            "Expected array, but got {t:?}"
+                                        ),
+                                    },
+                                    t => panic!(
+                                        "Expected reference, but got {t:?}"
+                                    ),
+                                }
+                            }
+                            t => panic!("Expected block, but got {t:?}"),
+                        }
+                    }
+                    t => panic!("Expected const, but got {t:?}"),
+                })
+                .unwrap()
+        };
+        let string = all
+            .elems
+            .into_iter()
+            .map(|e| e.to_token_stream().to_string())
+            .join(" ");
+        assert_eq!(string, "bird is the word");
+    }
+
+    /// Test if we reject structs with non static string fields in
+    /// `StorageKeys` macro derives.
+    #[test]
+    #[should_panic(
+        expected = "Expected `&'static str` field type in StorageKeys derive"
+    )]
+    fn test_typecheck_storage_keys_derive() {
+        derive_storage_keys_inner(quote! {
+            struct Keys {
+                x: &'static str,
+                y: i32,
+                z: u64,
+            }
+        });
+    }
+
+    /// Test that the create storage keys produces
+    /// the expected code.
+    #[test]
+    fn test_derive_storage_keys() {
+        let test_struct = quote! {
+            struct Keys {
+                param1: &'static str,
+                param2: &'static str,
+            }
+        };
+        let test_impl: ItemImpl =
+            syn::parse2(derive_storage_keys_inner(test_struct))
+                .expect("Test failed");
+
+        let expected_impl = quote! {
+            impl Keys {
+                const ALL: &[&'static str] = {
+                    let Keys { param1, param2 } = Self::VALUES;
+                    &[param1, param2]
+                };
+                const VALUES: Keys = Self {
+                    param1: "param1",
+                    param2: "param2"
+                };
+            }
+        };
+        let expected_impl: ItemImpl =
+            syn::parse2(expected_impl).expect("Test failed");
+
+        assert_eq!(test_impl, expected_impl);
+    }
+}