migration to update a number of users total allocations

Makefile.toml

@@ -6,8 +6,19 @@ extend = [
 [config]
 default_to_workspace = false
 
+[env]
+# Directory with wasm files used by integration tests (another directory can be used instead, for example 'artifacts' from rust-optimizer)
+ARTIFACTS_DIR_PATH = "target/wasm32-unknown-unknown/release"
+# If you bump this version, verify RUST_VERSION correctness
+RUST_OPTIMIZER_VERSION = "0.13.0"
+# Use rust version from rust-optimizer Dockerfile (see https://github.com/CosmWasm/rust-optimizer/blob/main/Dockerfile#L1)
+# to be sure that we compile / test against the same version
+RUST_VERSION = "1.69.0"
+
 [tasks.all-actions]
 dependencies = [
+  "install-stable",
+  "install-nightly",
   "fmt",
   "clippy",
   "build",
@@ -17,6 +28,22 @@ dependencies = [
   "rust-optimizer",
 ]
 
+[tasks.install-stable]
+script = '''
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain ${RUST_VERSION}
+rustup target add wasm32-unknown-unknown --toolchain ${RUST_VERSION}
+rustup component add rustfmt --toolchain ${RUST_VERSION}
+rustup component add clippy --toolchain ${RUST_VERSION}
+rustup component add llvm-tools-preview --toolchain ${RUST_VERSION}
+'''
+
+[tasks.install-nightly]
+script = '''
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain nightly
+rustup target add wasm32-unknown-unknown --toolchain nightly
+rustup component add rustfmt --toolchain nightly
+'''
+
 [tasks.build]
 command = "cargo"
 args = ["build", "--release", "--target", "wasm32-unknown-unknown", "--locked"]

contracts/vesting/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name          = "mars-vesting"
 description   = "Smart contract managing token vesting for Mars protocol contributors"
-version       = "1.1.0"
+version       = "1.1.1"
 authors       = { workspace = true }
 edition       = { workspace = true }
 rust-version  = { workspace = true }

contracts/vesting/src/contract.rs

@@ -1,8 +1,8 @@
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
-    coins, to_binary, Addr, BankMsg, Binary, CosmosMsg, Deps, DepsMut, Empty, Env, MessageInfo,
-    Order, Response, Uint128,
+    coins, to_binary, Addr, BankMsg, Binary, CosmosMsg, Deps, DepsMut, Env, MessageInfo, Order,
+    Response, Uint128,
 };
 use cw2::set_contract_version;
 use cw_storage_plus::Bound;
@@ -11,9 +11,10 @@ use cw_utils::must_pay;
 use crate::{
     error::{Error, Result},
     helpers::{compute_position_response, compute_withdrawable},
-    migrations::v1_1_0,
+    migrations::{v1_1_0, v1_1_1},
     msg::{
-        Config, ExecuteMsg, Position, PositionResponse, QueryMsg, Schedule, VotingPowerResponse,
+        Config, ExecuteMsg, MigrateMsg, Position, PositionResponse, QueryMsg, Schedule,
+        VotingPowerResponse,
     },
     state::{CONFIG, POSITIONS},
 };
@@ -306,6 +307,9 @@ pub fn query_positions(
 //--------------------------------------------------------------------------------------------------
 
 #[cfg_attr(not(feature = "library"), entry_point)]
-pub fn migrate(deps: DepsMut, _: Env, _: Empty) -> Result<Response> {
-    v1_1_0::migrate(deps)
+pub fn migrate(deps: DepsMut, _: Env, msg: MigrateMsg) -> Result<Response> {
+    match msg {
+        MigrateMsg::V1_0_0ToV1_1_0 {} => v1_1_0::migrate(deps),
+        MigrateMsg::V1_1_0ToV1_1_1(updates) => v1_1_1::migrate(deps, updates),
+    }
 }

contracts/vesting/src/error.rs

@@ -17,6 +17,9 @@ pub enum Error {
 
     #[error("{0}")]
     Version(#[from] cw2::VersionError),
+
+    #[error("{0}")]
+    Overflow(#[from] cosmwasm_std::OverflowError),
 }
 
 pub(crate) type Result<T> = core::result::Result<T, Error>;

contracts/vesting/src/migrations/mod.rs

@@ -1 +1,2 @@
 pub mod v1_1_0;
+pub mod v1_1_1;

contracts/vesting/src/migrations/v1_1_1.rs

@@ -0,0 +1,51 @@
+use cosmwasm_std::{coins, BankMsg, CosmosMsg, DepsMut, Response, Uint128};
+use cw2::set_contract_version;
+
+use crate::{
+    contract::{CONTRACT_NAME, CONTRACT_VERSION},
+    error::Result,
+    msg::V1_1_1Updates,
+    state::{CONFIG, POSITIONS},
+};
+
+const FROM_VERSION: &str = "1.1.0";
+
+pub fn migrate(deps: DepsMut, msg: V1_1_1Updates) -> Result<Response> {
+    // make sure we're migrating the correct contract and from the correct version
+    cw2::assert_contract_version(deps.as_ref().storage, CONTRACT_NAME, FROM_VERSION)?;
+
+    let mut total_reclaim: Uint128 = Uint128::new(0);
+
+    for position_alteration in msg.position_alterations {
+        let mut position = POSITIONS.load(deps.storage, &position_alteration.addr)?;
+
+        // Determine amount to send back to owner
+        let reclaim = position.total.checked_sub(position_alteration.total_new)?;
+        total_reclaim = total_reclaim.checked_add(reclaim)?;
+
+        // Confirm state is as expected
+        assert!(position.total == position_alteration.total_old);
+        assert!(position.withdrawn.is_zero());
+        assert!(reclaim == position_alteration.reclaim);
+
+        // Set new total vesting amount
+        position.total = position_alteration.total_new;
+        POSITIONS.save(deps.storage, &position_alteration.addr, &position)?;
+    }
+
+    // Additoinal check that the total amount reclaimed back is as expected
+    assert!(total_reclaim == msg.total_reclaim);
+
+    set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
+
+    let cfg = CONFIG.load(deps.storage)?;
+
+    Ok(Response::new()
+        .add_message(CosmosMsg::Bank(BankMsg::Send {
+            to_address: cfg.owner.to_string(),
+            amount: coins(total_reclaim.u128(), cfg.denom),
+        }))
+        .add_attribute("action", "migrate")
+        .add_attribute("from_version", FROM_VERSION)
+        .add_attribute("to_version", CONTRACT_VERSION))
+}

contracts/vesting/src/msg.rs

@@ -133,3 +133,29 @@ pub struct PositionResponse {
     /// This vesting position's vesting schedule
     pub vest_schedule: Schedule,
 }
+
+#[cw_serde]
+pub enum MigrateMsg {
+    V1_0_0ToV1_1_0 {},
+    V1_1_0ToV1_1_1(V1_1_1Updates),
+}
+
+#[cw_serde]
+pub struct V1_1_1Updates {
+    /// Array of positions for alteration
+    pub position_alterations: Vec<PositionAlteration>,
+    /// Total amount of MARS to be reclaimed
+    pub total_reclaim: Uint128,
+}
+
+#[cw_serde]
+pub struct PositionAlteration {
+    /// Address of user to alter
+    pub addr: Addr,
+    /// Total amount of MARS allocated previously
+    pub total_old: Uint128,
+    /// Total amount of MARS allocated now
+    pub total_new: Uint128,
+    /// Total amount of MARS to be reclaimed
+    pub reclaim: Uint128,
+}

contracts/vesting/tests/tests.rs

@@ -8,11 +8,11 @@ use cw_utils::PaymentError;
 use mars_vesting::{
     contract::{execute, instantiate, migrate, query},
     error::Error,
-    migrations::v1_1_0::v1_0_0_state,
     msg::{
-        Config, ExecuteMsg, Position, PositionResponse, QueryMsg, Schedule, VotingPowerResponse,
+        Config, ExecuteMsg, MigrateMsg, Position, PositionAlteration, PositionResponse, QueryMsg,
+        Schedule, V1_1_1Updates, VotingPowerResponse,
     },
-    state::{CONFIG, POSITIONS},
+    state::POSITIONS,
 };
 
 pub const MOCK_DENOM: &str = "umars";
@@ -592,7 +592,7 @@ fn invalid_contract_version() {
 
     let old_contract_version = ContractVersion {
         contract: "crates.io:mars-vesting".to_string(),
-        version: "0.9.0".to_string(),
+        version: "1.0.0".to_string(),
     };
 
     set_contract_version(
@@ -602,42 +602,63 @@ fn invalid_contract_version() {
     )
     .unwrap();
 
-    let err = migrate(deps.as_mut(), env, Empty {}).unwrap_err();
+    let update_msg = V1_1_1Updates {
+        position_alterations: vec![],
+        total_reclaim: Uint128::new(123),
+    };
+
+    let err = migrate(deps.as_mut(), env, MigrateMsg::V1_1_0ToV1_1_1(update_msg)).unwrap_err();
     assert_eq!(
         Error::Version(VersionError::WrongVersion {
-            expected: "1.0.0".to_string(),
-            found: "0.9.0".to_string()
+            expected: "1.1.0".to_string(),
+            found: "1.0.0".to_string()
         }),
         err
     );
 }
 
 #[test]
 fn proper_migration() {
-    let mut deps = mock_dependencies();
-    cw2::set_contract_version(deps.as_mut().storage, "crates.io:mars-vesting", "1.0.0").unwrap();
+    let mut deps = setup_test();
+    cw2::set_contract_version(deps.as_mut().storage, "crates.io:mars-vesting", "1.1.0").unwrap();
 
-    let old_owner = "spiderman_246";
-    v1_0_0_state::OWNER.save(deps.as_mut().storage, &Addr::unchecked(old_owner)).unwrap();
+    execute(
+        deps.as_mut(),
+        mock_env(),
+        mock_info("owner", &[coin(456, "umars")]),
+        ExecuteMsg::CreatePosition {
+            user: "larry".to_string(),
+            vest_schedule: Schedule {
+                start_time: 1614600000, // 2021-03-01
+                cliff: 31536000,        // 1 year
+                duration: 126144000,    // 4 years
+            },
+        },
+    )
+    .unwrap();
 
-    let old_schedule = Schedule {
-        start_time: 1614600000,
-        cliff: 31536000,
-        duration: 126144000,
+    let update_msg = V1_1_1Updates {
+        position_alterations: vec![PositionAlteration {
+            addr: Addr::unchecked("larry"),
+            total_old: Uint128::new(456),
+            total_new: Uint128::new(333),
+            reclaim: Uint128::new(123),
+        }],
+        total_reclaim: Uint128::new(123),
     };
-    v1_0_0_state::UNLOCK_SCHEDULE.save(deps.as_mut().storage, &old_schedule).unwrap();
 
-    let res = migrate(deps.as_mut(), mock_env(), Empty {}).unwrap();
+    let res = migrate(deps.as_mut(), mock_env(), MigrateMsg::V1_1_0ToV1_1_1(update_msg)).unwrap();
 
-    assert_eq!(res.messages, vec![]);
+    assert_eq!(
+        res.messages,
+        vec![SubMsg::new(CosmosMsg::Bank(BankMsg::Send {
+            to_address: "owner".to_string(),
+            amount: coins(123, "umars")
+        }))]
+    );
     assert!(res.data.is_none());
     assert_eq!(
         res.attributes,
-        vec![attr("action", "migrate"), attr("from_version", "1.0.0"), attr("to_version", "1.1.0"),]
+        vec![attr("action", "migrate"), attr("from_version", "1.1.0"), attr("to_version", "1.1.1"),]
     );
-
-    let config = CONFIG.load(deps.as_ref().storage).unwrap();
-    assert_eq!(config.denom, v1_0_0_state::VEST_DENOM.to_string());
-    assert_eq!(config.owner.to_string(), old_owner.to_string());
-    assert_eq!(config.unlock_schedule, old_schedule);
 }

schemas/mars-vesting/mars-vesting.json

@@ -1,6 +1,6 @@
 {
   "contract_name": "mars-vesting",
-  "contract_version": "1.1.0",
+  "contract_version": "1.1.1",
   "idl_version": "1.0.0",
   "instantiate": {
     "$schema": "http://json-schema.org/draft-07/schema#",