diff --git a/src/contract.rs b/src/contract.rs index a1d5a3d..eb51222 100644 --- a/src/contract.rs +++ b/src/contract.rs @@ -113,9 +113,17 @@ fn register_vesting_account( let denom_key = denom_to_key(deposit_denom.clone()); - // vesting_account existence check + let mut vesting_account: VestingAccount; if VESTING_ACCOUNTS.has(deps.storage, (recipient.as_str(), &denom_key)) { - return Err(StdError::generic_err("already exists")); + vesting_account = VESTING_ACCOUNTS.load(deps.storage, (recipient.as_str(), &denom_key))?; + } else { + vesting_account = VestingAccount{ + address: recipient.to_string(), + vesting_denom: deposit_denom.clone(), + vesting_amount: Uint128::zero(), + vesting_schedules: vec![], + claimed_amount: Uint128::zero(), + } } // validate vesting schedule @@ -150,6 +158,9 @@ fn register_vesting_account( "assert(deposit_amount == vesting_amount)", )); } + + vesting_account.vesting_schedules.push(vesting_schedule.clone()); + vesting_account.vesting_amount = vesting_account.vesting_amount + vesting_amount; } VestingSchedule::PeriodicVesting { start_time, @@ -179,14 +190,6 @@ fn register_vesting_account( return Err(StdError::generic_err("invalid start_time")); } - if end_time < start_time { - return Err(StdError::generic_err("assert(end_time > start_time)")); - } - - if vesting_interval == 0 { - return Err(StdError::generic_err("assert(vesting_interval != 0)")); - } - let time_period = end_time - start_time; if time_period != (time_period / vesting_interval) * vesting_interval { return Err(StdError::generic_err( @@ -201,19 +204,24 @@ fn register_vesting_account( "assert(deposit_amount = amount * ((end_time - start_time) / vesting_interval + 1))", )); } + + if end_time < start_time { + return Err(StdError::generic_err("assert(end_time => start_time)")); + } + + if vesting_interval == 0 { + return Err(StdError::generic_err("assert(vesting_interval != 0)")); + } + + vesting_account.vesting_schedules.push(vesting_schedule.clone()); + vesting_account.vesting_amount = vesting_account.vesting_amount + vesting_amount; } } VESTING_ACCOUNTS.save( deps.storage, (recipient.as_str(), &denom_key), - &VestingAccount { - address: recipient.to_string(), - vesting_denom: deposit_denom.clone(), - vesting_amount: deposit_amount, - vesting_schedule, - claimed_amount: Uint128::zero(), - }, + &vesting_account )?; Ok(Response::new().add_attributes(vec![ @@ -254,9 +262,7 @@ fn deregister_vesting_account( // remove vesting account VESTING_ACCOUNTS.remove(deps.storage, (address.as_str(), &denom_key)); - let vested_amount = account - .vesting_schedule - .vested_amount(env.block.time.seconds())?; + let vested_amount = get_vested_amount(account.vesting_schedules, env.block.time.seconds())?; let claimed_amount = account.claimed_amount; // transfer already vested but not claimed amount to @@ -350,9 +356,7 @@ fn claim( } let mut account = account.unwrap(); - let vested_amount = account - .vesting_schedule - .vested_amount(env.block.time.seconds())?; + let vested_amount = get_vested_amount(account.vesting_schedules.clone(), env.block.time.seconds())?; let claimed_amount = account.claimed_amount; let claimable_amount = vested_amount.checked_sub(claimed_amount)?; @@ -405,6 +409,18 @@ fn claim( .add_attributes(attrs)) } +fn get_vested_amount(vesting_schedules: Vec, block_time: u64)-> StdResult { + let schedules_len = vesting_schedules.len(); + let mut total_vested_amount = Uint128::zero(); + + for index in 0..schedules_len { + let vested_amount = vesting_schedules[index].vested_amount(block_time)?; + total_vested_amount = total_vested_amount + vested_amount; + } + + return Ok(total_vested_amount) +} + pub fn receive_cw20( deps: DepsMut, env: Env, @@ -475,15 +491,13 @@ fn vesting_account( .take(limit) { let (_, account) = item?; - let vested_amount = account - .vesting_schedule - .vested_amount(env.block.time.seconds())?; + let vested_amount = get_vested_amount(account.vesting_schedules.clone(), env.block.time.seconds())?; vestings.push(VestingData { vesting_denom: account.vesting_denom, vesting_amount: account.vesting_amount, vested_amount, - vesting_schedule: account.vesting_schedule, + vesting_schedules: account.vesting_schedules, claimable_amount: vested_amount.checked_sub(account.claimed_amount)?, }) } diff --git a/src/msg.rs b/src/msg.rs index e5cca4d..36b91d9 100644 --- a/src/msg.rs +++ b/src/msg.rs @@ -77,7 +77,7 @@ pub struct VestingData { pub vesting_denom: Denom, pub vesting_amount: Uint128, pub vested_amount: Uint128, - pub vesting_schedule: VestingSchedule, + pub vesting_schedules: Vec, pub claimable_amount: Uint128, } @@ -153,6 +153,23 @@ impl VestingSchedule { } } } + + // return (start_time, end_timg) + pub fn get_vesting_time(&self) -> StdResult<(u64, u64)> { + match self { + VestingSchedule::LinearVesting { + start_time, + end_time, + vesting_amount: _ + } => return Ok((start_time.parse::().unwrap(), end_time.parse::().unwrap())), + VestingSchedule::PeriodicVesting { + start_time, + end_time, + vesting_interval: _, + amount: _, + } => return Ok((start_time.parse::().unwrap(), end_time.parse::().unwrap())), + } + } } #[test] diff --git a/src/state.rs b/src/state.rs index 5e5f8e8..2932c7d 100644 --- a/src/state.rs +++ b/src/state.rs @@ -14,7 +14,7 @@ pub struct VestingAccount { pub address: String, pub vesting_denom: Denom, pub vesting_amount: Uint128, - pub vesting_schedule: VestingSchedule, + pub vesting_schedules: Vec, pub claimed_amount: Uint128, } diff --git a/src/testing.rs b/src/testing.rs index 59769d4..43d2347 100644 --- a/src/testing.rs +++ b/src/testing.rs @@ -181,11 +181,11 @@ fn register_vesting_account_with_native_token() { vesting_denom: Denom::Native("uusd".to_string()), vesting_amount: Uint128::new(1000000), vested_amount: Uint128::zero(), - vesting_schedule: VestingSchedule::LinearVesting { + vesting_schedules: vec![VestingSchedule::LinearVesting { start_time: "100".to_string(), end_time: "110".to_string(), vesting_amount: Uint128::new(1000000u128), - }, + }], claimable_amount: Uint128::zero(), }], } @@ -304,11 +304,11 @@ fn register_vesting_account_with_cw20_token() { vesting_denom: Denom::Cw20(Addr::unchecked("token0000")), vesting_amount: Uint128::new(1000000), vested_amount: Uint128::zero(), - vesting_schedule: VestingSchedule::LinearVesting { + vesting_schedules: vec![VestingSchedule::LinearVesting { start_time: "100".to_string(), end_time: "110".to_string(), vesting_amount: Uint128::new(1000000u128), - }, + }], claimable_amount: Uint128::zero(), }], } @@ -417,11 +417,11 @@ fn claim_native() { vesting_denom: Denom::Native("uusd".to_string()), vesting_amount: Uint128::new(1000000), vested_amount: Uint128::new(500000), - vesting_schedule: VestingSchedule::LinearVesting { + vesting_schedules: vec![VestingSchedule::LinearVesting { start_time: "100".to_string(), end_time: "110".to_string(), vesting_amount: Uint128::new(1000000u128), - }, + }], claimable_amount: Uint128::zero(), }], } @@ -585,11 +585,11 @@ fn claim_cw20() { vesting_denom: Denom::Cw20(Addr::unchecked("token0001")), vesting_amount: Uint128::new(1000000), vested_amount: Uint128::new(500000), - vesting_schedule: VestingSchedule::LinearVesting { + vesting_schedules: vec![VestingSchedule::LinearVesting { start_time: "100".to_string(), end_time: "110".to_string(), vesting_amount: Uint128::new(1000000u128), - }, + }], claimable_amount: Uint128::zero(), }], } @@ -718,22 +718,22 @@ fn query_vesting_account() { vesting_denom: Denom::Cw20(Addr::unchecked("token0001")), vesting_amount: Uint128::new(1000000), vested_amount: Uint128::new(500000), - vesting_schedule: VestingSchedule::LinearVesting { + vesting_schedules: vec![VestingSchedule::LinearVesting { start_time: "100".to_string(), end_time: "110".to_string(), vesting_amount: Uint128::new(1000000u128), - }, + }], claimable_amount: Uint128::new(500000), }, VestingData { vesting_denom: Denom::Native("uusd".to_string()), vesting_amount: Uint128::new(1000000), vested_amount: Uint128::new(500000), - vesting_schedule: VestingSchedule::LinearVesting { + vesting_schedules: vec![VestingSchedule::LinearVesting { start_time: "100".to_string(), end_time: "110".to_string(), vesting_amount: Uint128::new(1000000u128), - }, + }], claimable_amount: Uint128::new(500000), } ], @@ -761,11 +761,11 @@ fn query_vesting_account() { vesting_denom: Denom::Cw20(Addr::unchecked("token0001")), vesting_amount: Uint128::new(1000000), vested_amount: Uint128::new(500000), - vesting_schedule: VestingSchedule::LinearVesting { + vesting_schedules: vec![VestingSchedule::LinearVesting { start_time: "100".to_string(), end_time: "110".to_string(), vesting_amount: Uint128::new(1000000u128), - }, + }], claimable_amount: Uint128::new(500000), },], } @@ -792,13 +792,116 @@ fn query_vesting_account() { vesting_denom: Denom::Native("uusd".to_string()), vesting_amount: Uint128::new(1000000), vested_amount: Uint128::new(500000), - vesting_schedule: VestingSchedule::LinearVesting { + vesting_schedules: vec![VestingSchedule::LinearVesting { start_time: "100".to_string(), end_time: "110".to_string(), vesting_amount: Uint128::new(1000000u128), - }, + }], claimable_amount: Uint128::new(500000), }], } ); } + + + + +#[test] +fn add_multiple_schedules() { + let mut deps = mock_dependencies(&[]); + let _res = instantiate( + deps.as_mut(), + mock_env(), + mock_info("addr0000", &[]), + InstantiateMsg { + master_address: None, + }, + ) + .unwrap(); + + // init env to time 100 + let mut env = mock_env(); + env.block.time = Timestamp::from_seconds(100); + + let msg = ExecuteMsg::RegisterVestingAccount { + address: "addr0001".to_string(), + vesting_schedule: VestingSchedule::LinearVesting { + start_time: "100".to_string(), + end_time: "110".to_string(), + vesting_amount: Uint128::new(1000000u128), + }, + }; + + let info = mock_info("addr0000", &[Coin::new(1000000u128, "uusd")]); + let _ = execute(deps.as_mut(), env.clone(), info, msg).unwrap(); + + // valid + let msg = ExecuteMsg::RegisterVestingAccount { + address: "addr0001".to_string(), + vesting_schedule: VestingSchedule::LinearVesting { + start_time: "120".to_string(), + end_time: "130".to_string(), + vesting_amount: Uint128::new(1000000u128), + }, + }; + + let info = mock_info("addr0000", &[Coin::new(1000000u128, "uusd")]); + let _ = execute(deps.as_mut(), env.clone(), info, msg).unwrap(); + + // valid and the one between others + let msg = ExecuteMsg::RegisterVestingAccount { + address: "addr0001".to_string(), + vesting_schedule: VestingSchedule::LinearVesting { + start_time: "105".to_string(), + end_time: "125".to_string(), + vesting_amount: Uint128::new(1000000u128), + }, + }; + + let info = mock_info("addr0000", &[Coin::new(1000000u128, "uusd")]); + let _ = execute(deps.as_mut(), env.clone(), info, msg).unwrap(); + + env.block.time = Timestamp::from_seconds(115); + + assert_eq!( + from_binary::( + &query( + deps.as_ref(), + env.clone(), + QueryMsg::VestingAccount { + address: "addr0001".to_string(), + start_after: None, + limit: Some(1), + }, + ) + .unwrap() + ) + .unwrap(), + VestingAccountResponse { + address: "addr0001".to_string(), + vestings: vec![VestingData { + vesting_denom: Denom::Native("uusd".to_string()), + vesting_amount: Uint128::new(3000000), + vested_amount: Uint128::new(1500000), + vesting_schedules: vec![ + VestingSchedule::LinearVesting { + start_time: "100".to_string(), + end_time: "110".to_string(), + vesting_amount: Uint128::new(1000000u128), + }, + VestingSchedule::LinearVesting { + start_time: "120".to_string(), + end_time: "130".to_string(), + vesting_amount: Uint128::new(1000000u128), + }, + VestingSchedule::LinearVesting { + start_time: "105".to_string(), + end_time: "125".to_string(), + vesting_amount: Uint128::new(1000000u128), + }, + ], + claimable_amount: Uint128::new(1500000), + },], + } + ); +} \ No newline at end of file