From d0356bcd589b52ddcb36e97dd6ae3410ce41eac8 Mon Sep 17 00:00:00 2001
From: Tomasz Kulik <tomek.k@confio.gmbh>
Date: Wed, 14 Aug 2024 12:13:14 +0200
Subject: [PATCH] chore: Add new functions for migrate

---
 packages/std/src/lib.rs        |  2 +-
 packages/std/src/types.rs      | 11 ++++++++
 packages/vm/src/calls.rs       | 47 +++++++++++++++++++++++++++++++++-
 packages/vm/src/environment.rs |  1 -
 packages/vm/src/lib.rs         |  5 ++--
 5 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/packages/std/src/lib.rs b/packages/std/src/lib.rs
index 962409d9e2..9648b34c4b 100644
--- a/packages/std/src/lib.rs
+++ b/packages/std/src/lib.rs
@@ -116,7 +116,7 @@ pub use crate::stdack::StdAck;
 pub use crate::storage::MemoryStorage;
 pub use crate::timestamp::Timestamp;
 pub use crate::traits::{Api, HashFunction, Querier, QuerierResult, QuerierWrapper, Storage};
-pub use crate::types::{BlockInfo, ContractInfo, Env, MessageInfo, TransactionInfo};
+pub use crate::types::{BlockInfo, ContractInfo, Env, MessageInfo, MigrateInfo, TransactionInfo};
 
 // Exposed in wasm build only
 
diff --git a/packages/std/src/types.rs b/packages/std/src/types.rs
index f85d42e8a0..029d795acf 100644
--- a/packages/std/src/types.rs
+++ b/packages/std/src/types.rs
@@ -109,3 +109,14 @@ pub struct MessageInfo {
 pub struct ContractInfo {
     pub address: Addr,
 }
+
+/// The structure contains additional information related to the
+/// contract's state migration procedure - the sender address and
+/// the contract's state version currently stored on the blockchain.
+/// The `old_state_version` is optional, since there is no guarantee
+/// that the currently stored contract's binary contains that information.
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
+pub struct MigrateInfo {
+    pub sender: Addr,
+    pub old_state_version: Option<u64>,
+}
diff --git a/packages/vm/src/calls.rs b/packages/vm/src/calls.rs
index 96df03672c..70f2a7837a 100644
--- a/packages/vm/src/calls.rs
+++ b/packages/vm/src/calls.rs
@@ -3,7 +3,7 @@ use wasmer::Value;
 
 use cosmwasm_std::{
     ContractResult, CustomMsg, Env, IbcBasicResponse, IbcDestinationCallbackMsg,
-    IbcSourceCallbackMsg, MessageInfo, QueryResponse, Reply, Response,
+    IbcSourceCallbackMsg, MessageInfo, MigrateInfo, QueryResponse, Reply, Response,
 };
 #[cfg(feature = "stargate")]
 use cosmwasm_std::{
@@ -163,6 +163,26 @@ where
     Ok(result)
 }
 
+pub fn call_migrate2<A, S, Q, U>(
+    instance: &mut Instance<A, S, Q>,
+    env: &Env,
+    msg: &[u8],
+    migrate_info: &MigrateInfo,
+) -> VmResult<ContractResult<Response<U>>>
+where
+    A: BackendApi + 'static,
+    S: Storage + 'static,
+    Q: Querier + 'static,
+    U: DeserializeOwned + CustomMsg,
+{
+    let env = to_vec(env)?;
+    let migrate_info = to_vec(migrate_info)?;
+    let data = call_migrate2_raw(instance, &env, msg, &migrate_info)?;
+    let result: ContractResult<Response<U>> =
+        from_slice(&data, deserialization_limits::RESULT_MIGRATE)?;
+    Ok(result)
+}
+
 pub fn call_sudo<A, S, Q, U>(
     instance: &mut Instance<A, S, Q>,
     env: &Env,
@@ -441,6 +461,31 @@ where
     )
 }
 
+/// Calls Wasm export "migrate" and returns raw data from the contract.
+/// The result is length limited to prevent abuse but otherwise unchecked.
+/// The difference between this function and [call_migrate_raw] is the
+/// additional argument - `migrate_info`. It contains additional data
+/// related to the contract's state migration procedure.
+pub fn call_migrate2_raw<A, S, Q>(
+    instance: &mut Instance<A, S, Q>,
+    env: &[u8],
+    msg: &[u8],
+    migrate_info: &[u8],
+) -> VmResult<Vec<u8>>
+where
+    A: BackendApi + 'static,
+    S: Storage + 'static,
+    Q: Querier + 'static,
+{
+    instance.set_storage_readonly(false);
+    call_raw(
+        instance,
+        "migrate",
+        &[env, migrate_info, msg],
+        read_limits::RESULT_MIGRATE,
+    )
+}
+
 /// Calls Wasm export "sudo" and returns raw data from the contract.
 /// The result is length limited to prevent abuse but otherwise unchecked.
 pub fn call_sudo_raw<A, S, Q>(
diff --git a/packages/vm/src/environment.rs b/packages/vm/src/environment.rs
index 8b8ae2b823..ca9c50dfde 100644
--- a/packages/vm/src/environment.rs
+++ b/packages/vm/src/environment.rs
@@ -274,7 +274,6 @@ impl<A: BackendApi, S: Storage, Q: Querier> Environment<A, S, Q> {
         })?;
         let function_arity = func.param_arity(store);
         if args.len() != function_arity {
-            // TODO tkulik: when return error, try to call it with different args vector
             return Err(VmError::func_arity_mismatch(function_arity));
         };
         self.increment_call_depth()?;
diff --git a/packages/vm/src/lib.rs b/packages/vm/src/lib.rs
index 699c03735f..954646f5fb 100644
--- a/packages/vm/src/lib.rs
+++ b/packages/vm/src/lib.rs
@@ -29,8 +29,9 @@ pub use crate::cache::{
 pub use crate::calls::{
     call_execute, call_execute_raw, call_ibc_destination_callback,
     call_ibc_destination_callback_raw, call_ibc_source_callback, call_ibc_source_callback_raw,
-    call_instantiate, call_instantiate_raw, call_migrate, call_migrate_raw, call_query,
-    call_query_raw, call_reply, call_reply_raw, call_sudo, call_sudo_raw,
+    call_instantiate, call_instantiate_raw, call_migrate, call_migrate2, call_migrate2_raw,
+    call_migrate_raw, call_query, call_query_raw, call_reply, call_reply_raw, call_sudo,
+    call_sudo_raw,
 };
 #[cfg(feature = "stargate")]
 pub use crate::calls::{