From bb79d143ec10b13371911b9f554c1ebeec819a82 Mon Sep 17 00:00:00 2001
From: Tomasz Kulik <tomek.k@confio.gmbh>
Date: Mon, 2 Sep 2024 09:41:55 +0200
Subject: [PATCH] feat: Allow devs to assign version to a constant

---
 packages/derive/src/lib.rs | 102 +++++++++++++++++++++++++++++++++----
 1 file changed, 93 insertions(+), 9 deletions(-)

diff --git a/packages/derive/src/lib.rs b/packages/derive/src/lib.rs
index ddbb4bb13..41d4c8db7 100644
--- a/packages/derive/src/lib.rs
+++ b/packages/derive/src/lib.rs
@@ -1,5 +1,6 @@
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote, ToTokens};
+use syn::spanned::Spanned;
 use syn::{
     parse::{Parse, ParseStream},
     parse_quote,
@@ -111,6 +112,23 @@ impl Parse for Options {
 ///     todo!();
 /// }
 /// ```
+///
+/// It is also possible to assign the migrate version number to
+/// a given constant name:
+///
+/// ```
+/// # use cosmwasm_std::{
+/// #     DepsMut, entry_point, Env,
+/// #     Response, StdResult,
+/// # };
+/// #
+/// # type MigrateMsg = ();
+/// #[entry_point]
+/// #[migrate_version(CONTRACT_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,
@@ -135,7 +153,35 @@ fn expand_attributes(func: &mut ItemFn) -> syn::Result<TokenStream> {
             ));
         }
 
-        let version: syn::LitInt = attribute.parse_args()?;
+        let (const_name, version): (Option<syn::Path>, syn::LitInt) =
+            match attribute.parse_args()? {
+                syn::Expr::Assign(syn::ExprAssign { left, right, .. }) => match (*left, *right) {
+                    (
+                        syn::Expr::Path(syn::ExprPath { path, .. }),
+                        syn::Expr::Lit(syn::ExprLit {
+                            lit: syn::Lit::Int(version),
+                            ..
+                        }),
+                    ) => (Some(path), version),
+                    _ => {
+                        return Err(syn::Error::new(
+                            attribute.span(),
+                            "Expected version number or `CONST_NAME = {version_number}`",
+                        ))
+                    }
+                },
+                syn::Expr::Lit(syn::ExprLit {
+                    lit: syn::Lit::Int(version),
+                    ..
+                }) => (None, version),
+                _ => {
+                    return Err(syn::Error::new(
+                        attribute.span(),
+                        "Expected version number or `CONST_NAME = {version_number}`",
+                    ))
+                }
+            };
+
         // Enforce that the version is a valid u64 and non-zero
         let numeric_version = version.base10_parse::<u64>()?;
         if numeric_version == 0 {
@@ -144,6 +190,15 @@ fn expand_attributes(func: &mut ItemFn) -> syn::Result<TokenStream> {
                 "please start versioning with 1",
             ));
         }
+        let const_assignment = if let Some(const_name) = const_name {
+            quote! {
+                #[allow(unused)]
+                #[doc(hidden)]
+                pub const #const_name: u64 = #numeric_version;
+            }
+        } else {
+            quote! {}
+        };
 
         let version = version.base10_digits();
         let n = version.len();
@@ -160,11 +215,7 @@ fn expand_attributes(func: &mut ItemFn) -> syn::Result<TokenStream> {
             /// The format and even the existence of this value is an implementation detail, DO NOT RELY ON THIS!
             static __CW_MIGRATE_VERSION: [u8; #n] = *#version;
 
-            #[allow(unused)]
-            #[doc(hidden)]
-            const fn migrate_version() -> u64 {
-                #numeric_version
-            }
+            #const_assignment
         };
     }
 
@@ -284,11 +335,44 @@ mod test {
             /// The format and even the existence of this value is an implementation detail, DO NOT RELY ON THIS!
             static __CW_MIGRATE_VERSION: [u8; 1usize] = *b"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 contract_migrate_version_with_const_expansion() {
+        let code = quote! {
+            #[migrate_version(CONTRACT_VERSION = 66)]
+            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)]
-            const fn migrate_version() -> u64 {
-                2u64
-            }
+            #[cfg(target_arch = "wasm32")]
+            #[link_section = "cw_migrate_version"]
+            /// This is an internal constant exported as a custom section denoting the contract migrate version.
+            /// The format and even the existence of this value is an implementation detail, DO NOT RELY ON THIS!
+            static __CW_MIGRATE_VERSION: [u8; 2usize] = *b"66";
+
+            #[allow(unused)]
+            #[doc(hidden)]
+            pub const CONTRACT_VERSION: u64 = 66u64;
 
             fn migrate(deps: DepsMut, env: Env, msg: MigrateMsg) -> Response {
                 // Logic here