Add macro to add the contract state version as a custom section

packages/derive/src/lib.rs

@@ -4,7 +4,7 @@
     parse::{Parse, ParseStream},
     parse_quote,
     punctuated::Punctuated,
-    Token,
+    ItemFn, Token,
 };
 
 macro_rules! maybe {
@@ -93,6 +93,24 @@
 ///
 /// where `InstantiateMsg`, `ExecuteMsg`, and `QueryMsg` are contract defined
 /// types that implement `DeserializeOwned + JsonSchema`.
+///
+/// ## Set the version of the state of your contract
+///
+/// The VM will use this as a hint whether it needs to run the migrate function of your contract or not.
+///
+/// ```
+/// # use cosmwasm_std::{
+/// #     DepsMut, entry_point, Env,
+/// #     Response, StdResult,
+/// # };
+/// #
+/// # type MigrateMsg = ();
+/// #[entry_point]
+/// #[set_state_version(2)]
+/// pub fn migrate(deps: DepsMut, env: Env, msg: MigrateMsg) -> StdResult<Response> {
+///     todo!();
+/// }
+/// ```
 #[proc_macro_attribute]
 pub fn entry_point(
     attr: proc_macro::TokenStream,
@@ -101,32 +119,68 @@
     entry_point_impl(attr.into(), item.into()).into()
 }
 
-fn entry_point_impl(attr: TokenStream, mut item: TokenStream) -> TokenStream {
-    let cloned = item.clone();
-    let function: syn::ItemFn = maybe!(syn::parse2(cloned));
+fn expand_attributes(func: &mut ItemFn) -> syn::Result<TokenStream> {
+    let attributes = std::mem::take(&mut func.attrs);
+    let mut stream = TokenStream::new();
+    for attribute in attributes {
+        if !attribute.path().is_ident("set_state_version") {
+            func.attrs.push(attribute);
+            continue;
+        }
+
+        if func.sig.ident != "migrate" {
+            return Err(syn::Error::new_spanned(
+                &attribute,
+                "you only want to add this macro to your migrate function",
+            ));
+        }
+
+        let version: syn::LitInt = attribute.parse_args()?;
+        let version = version.base10_digits();
+
+        stream = quote! {
+            #stream
+
+            #[allow(unused)]
+            #[doc(hidden)]
+            #[link_section = "cw_state_version"]
+            /// This is an internal constant exported as a custom section denoting the contract state version.
+            /// The format and even the existence of this value is an implementation detail, DO NOT RELY ON THIS!
+            static __CW_STATE_VERSION: &str = #version;
+        };
+    }
+
+    Ok(stream)
+}
+
+fn entry_point_impl(attr: TokenStream, item: TokenStream) -> TokenStream {
+    let mut function: syn::ItemFn = maybe!(syn::parse2(item));
     let Options { crate_path } = maybe!(syn::parse2(attr));
 
+    let attribute_code = maybe!(expand_attributes(&mut function));
+
     // The first argument is `deps`, the rest is region pointers
-    let args = function.sig.inputs.len() - 1;
-    let fn_name = function.sig.ident;
+    let args = function.sig.inputs.len().saturating_sub(1);
+    let fn_name = &function.sig.ident;
     let wasm_export = format_ident!("__wasm_export_{fn_name}");
     let do_call = format_ident!("do_{fn_name}");
 
     let decl_args = (0..args).map(|item| format_ident!("ptr_{item}"));
     let call_args = decl_args.clone();
 
-    let new_code = quote! {
+    quote! {
+        #attribute_code
+
+        #function
+
         #[cfg(target_arch = "wasm32")]
         mod #wasm_export { // new module to avoid conflict of function name
             #[no_mangle]
             extern "C" fn #fn_name(#( #decl_args : u32 ),*) -> u32 {
                 #crate_path::#do_call(&super::#fn_name, #( #call_args ),*)
             }
         }
-    };
-
-    item.extend(new_code);
-    item
+    }
 }
 
 #[cfg(test)]
@@ -136,6 +190,40 @@
 
     use crate::entry_point_impl;
 
+    #[test]
+    fn contract_state_version_expansion() {
+        let code = quote! {
+            #[set_state_version(2)]
+            fn migrate(deps: DepsMut, env: Env, msg: MigrateMsg) -> Response {
+                // Logic here
+            }
+        };
+
+        let actual = entry_point_impl(TokenStream::new(), code);
+        let expected = quote! {
+            #[allow(unused)]
+            #[doc(hidden)]
+            #[link_section = "cw_state_version"]
+            /// This is an internal constant exported as a custom section denoting the contract state version.
+            /// The format and even the existence of this value is an implementation detail, DO NOT RELY ON THIS!
+            static __CW_STATE_VERSION: &str = "2";
+
+            fn migrate(deps: DepsMut, env: Env, msg: MigrateMsg) -> Response {
+                // Logic here
+            }
+
+            #[cfg(target_arch = "wasm32")]
+            mod __wasm_export_migrate {
+                #[no_mangle]
+                extern "C" fn migrate(ptr_0: u32, ptr_1: u32) -> u32 {
+                    ::cosmwasm_std::do_migrate(&super::migrate, ptr_0, ptr_1)
+                }
+            }
+        };
+
+        assert_eq!(actual.to_string(), expected.to_string());
+    }
+
     #[test]
     fn default_expansion() {
         let code = quote! {