add NewTypeKey derive macro

CHANGELOG.md

@@ -5,6 +5,9 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased]
 
+## [2.1.0] - 2024-09-25
+- Add derive macro for `NewTypeKey`
+
 ## [2.0.0] - 2024-03-14
 
 ## [2.0.0-rc.0] - 2024-02-09

Cargo.toml

@@ -34,6 +34,7 @@ criterion = { version = "0.3", features = ["html_reports"] }
 # access to an entropy souce via getrandom such as wasm32-unknown-unknown
 rand = { version = "0.8", default-features = false }
 rand_xoshiro = { version = "0.6.0", default-features = false }
+derive_more = {version = "1.0.0", features = ["full"]}
 
 # We don't use the following dependencies directly. They're dependencies of our dependencies.
 # We specify them to tighten their version requirements so that builds with `-Zminimal-versions` work.

macros/Cargo.toml

@@ -13,4 +13,6 @@ documentation = "https://docs.cosmwasm.com"
 proc-macro = true
 
 [dependencies]
+proc-macro2 = "1.0.86"
+quote = "1.0.37"
 syn = { version = "2", features = ["full"] }

macros/README.md

@@ -20,3 +20,13 @@ struct TestIndexes<'a> {
     addr: UniqueIndex<'a, Addr, TestStruct, String>,
 }
 ```
+
+Auto generate the required impls to use a newtype as a key
+
+```rust
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
+#[derive(NewTypeKey)] // <- Add this line right here.
+struct TestKey(u64);
+
+// You can now use `TestKey` as a key in `Map`
+```

macros/src/index_list.rs

@@ -0,0 +1,36 @@
+use proc_macro::TokenStream;
+use syn::{
+    Ident,
+    __private::{quote::quote, Span},
+    parse_macro_input, ItemStruct,
+};
+
+pub fn index_list(attr: TokenStream, item: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(item as ItemStruct);
+
+    let ty = Ident::new(&attr.to_string(), Span::call_site());
+    let struct_ty = input.ident.clone();
+
+    let names = input
+        .fields
+        .clone()
+        .into_iter()
+        .map(|e| {
+            let name = e.ident.unwrap();
+            quote! { &self.#name }
+        })
+        .collect::<Vec<_>>();
+
+    let expanded = quote! {
+        #input
+
+        impl cw_storage_plus::IndexList<#ty> for #struct_ty<'_> {
+            fn get_indexes(&'_ self) -> Box<dyn Iterator<Item = &'_ dyn cw_storage_plus::Index<#ty>> + '_> {
+                let v: Vec<&dyn cw_storage_plus::Index<#ty>> = vec![#(#names),*];
+                Box::new(v.into_iter())
+            }
+        }
+    };
+
+    TokenStream::from(expanded)
+}

macros/src/lib.rs

@@ -2,43 +2,22 @@
 Procedural macros helper for interacting with cw-storage-plus and cosmwasm-storage.
 
 For more information on this package, please check out the
-[README](https://github.com/CosmWasm/cw-plus/blob/main/packages/storage-macro/README.md).
+[README](https://github.com/CosmWasm/cw-storage-plus/blob/main/macros/README.md).
 */
 
+mod index_list;
+mod newtype;
+
 use proc_macro::TokenStream;
-use syn::{
-    Ident,
-    __private::{quote::quote, Span},
-    parse_macro_input, ItemStruct,
-};
+
+// Re-export the procedural macro functions
 
 #[proc_macro_attribute]
 pub fn index_list(attr: TokenStream, item: TokenStream) -> TokenStream {
-    let input = parse_macro_input!(item as ItemStruct);
-
-    let ty = Ident::new(&attr.to_string(), Span::call_site());
-    let struct_ty = input.ident.clone();
-
-    let names = input
-        .fields
-        .clone()
-        .into_iter()
-        .map(|e| {
-            let name = e.ident.unwrap();
-            quote! { &self.#name }
-        })
-        .collect::<Vec<_>>();
-
-    let expanded = quote! {
-        #input
-
-        impl cw_storage_plus::IndexList<#ty> for #struct_ty<'_> {
-            fn get_indexes(&'_ self) -> Box<dyn Iterator<Item = &'_ dyn cw_storage_plus::Index<#ty>> + '_> {
-                let v: Vec<&dyn cw_storage_plus::Index<#ty>> = vec![#(#names),*];
-                Box::new(v.into_iter())
-            }
-        }
-    };
+    index_list::index_list(attr, item)
+}
 
-    TokenStream::from(expanded)
+#[proc_macro_derive(NewTypeKey)]
+pub fn cw_storage_newtype_key_derive(input: TokenStream) -> TokenStream {
+    newtype::cw_storage_newtype_key_derive(input)
 }

macros/src/newtype.rs

@@ -0,0 +1,82 @@
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::{parse_macro_input, Data, DeriveInput, Fields};
+
+pub fn cw_storage_newtype_key_derive(input: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(input as DeriveInput);
+
+    let expanded = impl_newtype(&input);
+
+    TokenStream::from(expanded)
+}
+
+fn impl_newtype(input: &DeriveInput) -> proc_macro2::TokenStream {
+    // Extract the struct name
+    let name = &input.ident;
+
+    // Extract the inner type
+    let inner_type = match &input.data {
+        Data::Struct(data_struct) => match &data_struct.fields {
+            Fields::Unnamed(fields) if fields.unnamed.len() == 1 => Some(&fields.unnamed[0].ty),
+            _ => None,
+        },
+        _ => None,
+    };
+
+    let inner_type = inner_type
+        .expect("NewTypeKey can only be derived for newtypes (tuple structs with one field)");
+
+    // Implement PrimaryKey
+    let impl_primary_key = quote! {
+        impl<'a> cw_storage_plus::PrimaryKey<'a> for #name
+        where
+            #inner_type: cw_storage_plus::PrimaryKey<'a>,
+        {
+            type Prefix = ();
+            type SubPrefix = ();
+            type Suffix = Self;
+            type SuperSuffix = Self;
+
+            fn key(&self) -> Vec<cw_storage_plus::Key> {
+                self.0.key()
+            }
+        }
+    };
+
+    // Implement Prefixer
+    let impl_prefixer = quote! {
+        impl<'a> cw_storage_plus::Prefixer<'a> for #name
+        where
+            #inner_type: cw_storage_plus::Prefixer<'a>,
+        {
+            fn prefix(&self) -> Vec<cw_storage_plus::Key> {
+                self.0.prefix()
+            }
+        }
+    };
+
+    // Implement KeyDeserialize
+    let impl_key_deserialize = quote! {
+        impl cw_storage_plus::KeyDeserialize for #name
+        where
+            #inner_type: cw_storage_plus::KeyDeserialize<Output = #inner_type>,
+        {
+            type Output = #name;
+            const KEY_ELEMS: u16 = 1;
+
+            #[inline(always)]
+            fn from_vec(value: Vec<u8>) -> cosmwasm_std::StdResult<Self::Output> {
+                <#inner_type as cw_storage_plus::KeyDeserialize>::from_vec(value).map(#name)
+            }
+        }
+    };
+
+    // Combine all implementations
+    let expanded = quote! {
+        #impl_primary_key
+        #impl_prefixer
+        #impl_key_deserialize
+    };
+
+    expanded
+}

src/lib.rs

@@ -83,3 +83,20 @@ extern crate cw_storage_macro;
 /// ```
 ///
 pub use cw_storage_macro::index_list;
+
+#[cfg(all(feature = "iterator", feature = "macro"))]
+/// Auto generate the required impls to use a newtype as a key
+/// # Example
+///
+/// ```rust
+/// use cw_storage_plus::NewTypeKey;
+/// use serde::{Serialize, Deserialize};
+///
+/// #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
+/// #[derive(NewTypeKey)] // <- Add this line right here.
+/// struct TestKey(u64);
+///
+/// // You can now use `TestKey` as a key in `Map`
+/// ```
+///
+pub use cw_storage_macro::NewTypeKey;