Skip to content

Commit

Permalink
Merge pull request #20 from sei-protocol/partial-withdrawal
Browse files Browse the repository at this point in the history 
Support partial withdraw for unlocked tokens
  • Loading branch information
@codchen
codchen authored Jul 9, 2024
2 parents 8990f0e + 34c6d0f commit f39cabb
Show file tree
Hide file tree
Showing 4 changed files with 101 additions and 39 deletions.
26 changes: 13 additions & 13 deletions src/contract.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use cw3::{
};
use cw_utils::{Threshold, ThresholdError};

use crate::data_structure::{EmptyStruct};
use crate::data_structure::EmptyStruct;
use crate::error::ContractError;
use crate::msg::{
AdminListResponse, ExecuteMsg, InstantiateMsg, MigrateMsg, OpListResponse, QueryMsg, ShowConfigResponse,
Expand Down Expand Up @@ -258,8 +258,8 @@ pub fn execute(
ExecuteMsg::Undelegate { validator, amount } => {
execute_undelegate(deps.as_ref(), info, validator, amount)
}
ExecuteMsg::InitiateWithdrawUnlocked {} => {
execute_initiate_withdraw_unlocked(deps, env, info)
ExecuteMsg::InitiateWithdrawUnlocked { amount } => {
execute_initiate_withdraw_unlocked(deps, env, info, amount)
}
ExecuteMsg::UpdateOp { op, remove } => execute_update_op(deps, info, op, remove),
ExecuteMsg::InitiateWithdrawReward {} => execute_initiate_withdraw_reward(deps, env, info),
Expand Down Expand Up @@ -341,9 +341,10 @@ fn execute_initiate_withdraw_unlocked(
deps: DepsMut,
env: Env,
info: MessageInfo,
amount: u128,
) -> Result<Response<Empty>, ContractError> {
authorize_op(deps.storage, info.sender)?;
let vested_amount = collect_vested(deps.storage, env.block.time)?;
let vested_amount = collect_vested(deps.storage, env.block.time, amount)?;
WITHDRAWN_UNLOCKED.update(deps.storage, |old| -> Result<u128, StdError> {
Ok(old + vested_amount)
})?;
Expand Down Expand Up @@ -662,7 +663,7 @@ fn execute_internal_update_unlocked_distribution_address(
unlocked_distribution_address: Addr,
) -> Result<Response<Empty>, ContractError> {
authorize_self_call(env, info)?;
UNLOCK_DISTRIBUTION_ADDRESS.save(deps.storage, &unlocked_distribution_address);
UNLOCK_DISTRIBUTION_ADDRESS.save(deps.storage, &unlocked_distribution_address)?;
Ok(Response::new())
}

Expand All @@ -673,7 +674,7 @@ fn execute_internal_update_staking_reward_distribution_address(
staking_reward_distribution_address: Addr,
) -> Result<Response<Empty>, ContractError> {
authorize_self_call(env, info)?;
STAKING_REWARD_ADDRESS.save(deps.storage, &staking_reward_distribution_address);
STAKING_REWARD_ADDRESS.save(deps.storage, &staking_reward_distribution_address)?;
Ok(Response::new())
}

Expand Down Expand Up @@ -801,11 +802,10 @@ fn query_total_vested(deps: Deps, env: Env) -> StdResult<ShowTotalVestedResponse
#[cfg(test)]
mod tests {
use cosmwasm_std::testing::{MOCK_CONTRACT_ADDR, mock_dependencies, mock_env, mock_info};
use cosmwasm_std::{from_binary, Addr, Coin, Decimal, FullDelegation, Timestamp, Validator, Storage};
use cosmwasm_std::{from_binary, Addr, Coin, Decimal, FullDelegation, Timestamp, Validator};

use cw2::{get_contract_version, ContractVersion};
use cw_utils::{Duration, Expiration, ThresholdResponse};
use schemars::_private::NoSerialize;

use crate::data_structure::Tranche;
use crate::state::get_number_of_ops;
Expand Down Expand Up @@ -1035,7 +1035,7 @@ mod tests {
let info = mock_info(VOTER5, &[Coin::new(48000000, "usei".to_string())]);
setup_test_case(deps.as_mut(), info.clone()).unwrap();

let msg = ExecuteMsg::InitiateWithdrawUnlocked {};
let msg = ExecuteMsg::InitiateWithdrawUnlocked { amount: 12000000 };
let mut env = mock_env();
let mut block = env.block;
block.time = block.time.plus_seconds(31536000);
Expand All @@ -1055,7 +1055,7 @@ mod tests {
let info = mock_info(OWNER, &[Coin::new(48000000, "usei".to_string())]);
setup_test_case(deps.as_mut(), info.clone()).unwrap();

let msg = ExecuteMsg::InitiateWithdrawUnlocked {};
let msg = ExecuteMsg::InitiateWithdrawUnlocked { amount: 12000000};
let mut env = mock_env();
let mut block = env.block;
block.time = block.time.plus_seconds(31536000);
Expand Down Expand Up @@ -1161,7 +1161,7 @@ mod tests {
remove: false
};
let internal_info = mock_info(MOCK_CONTRACT_ADDR, &[]);
execute(deps.as_mut(), mock_env(), internal_info, internal_update.clone());
execute(deps.as_mut(), mock_env(), internal_info, internal_update.clone()).unwrap();
let result = match ADMINS.load(deps.as_ref().storage, &Addr::unchecked(new_admin.clone())) {
Ok(_) => Ok(()),
Err(_) => Err(ContractError::Unauthorized {}),
Expand Down Expand Up @@ -1214,7 +1214,7 @@ mod tests {
unlocked_distribution_address: new_addr.clone()
};
let internal_info = mock_info(MOCK_CONTRACT_ADDR, &[]);
execute(deps.as_mut(), mock_env(), internal_info, internal_update.clone());
execute(deps.as_mut(), mock_env(), internal_info, internal_update.clone()).unwrap();
assert_eq!(UNLOCK_DISTRIBUTION_ADDRESS.load(deps.as_ref().storage).unwrap(), new_addr);
}

Expand Down Expand Up @@ -1247,7 +1247,7 @@ mod tests {
staking_reward_distribution_address: new_addr.clone()
};
let internal_info = mock_info(MOCK_CONTRACT_ADDR, &[]);
execute(deps.as_mut(), mock_env(), internal_info, internal_update.clone());
execute(deps.as_mut(), mock_env(), internal_info, internal_update.clone()).unwrap();
assert_eq!(STAKING_REWARD_ADDRESS.load(deps.as_ref().storage).unwrap(), new_addr);
}

Expand Down
2 changes: 2 additions & 0 deletions src/error.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,8 @@ pub enum ContractError {
#[error("Semver parsing error: {0}")]
SemVer(String),

#[error("No sufficient vested amount")]
NoSufficientUnlockedTokens {},
}

impl From<semver::Error> for ContractError {
Expand Down
4 changes: 3 additions & 1 deletion src/msg.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,9 @@ pub enum ExecuteMsg {
validator: String,
amount: u128,
},
InitiateWithdrawUnlocked {},
InitiateWithdrawUnlocked {
amount: u128,
},
InitiateWithdrawReward {},
UpdateOp {
op: Addr,
Expand Down
108 changes: 83 additions & 25 deletions src/vesting.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,29 +5,45 @@ use crate::{
ContractError,
};

pub fn collect_vested(storage: &mut dyn Storage, now: Timestamp) -> Result<u128, ContractError> {
pub fn collect_vested(storage: &mut dyn Storage, now: Timestamp, requested_amount: u128) -> Result<u128, ContractError> {
let vesting_ts = VESTING_TIMESTAMPS.load(storage)?;
let mut earliest_unvested_ts_idx = 0;
for ts in vesting_ts.iter() {
if *ts <= now {
earliest_unvested_ts_idx += 1;
let vesting_amounts = VESTING_AMOUNTS.load(storage)?;
let mut vested_amount = 0u128;
let mut amount_to_subtract = 0u128;
let mut remaining_first_idx: usize = 0;
for (i, (ts, amount)) in vesting_ts.iter().zip(vesting_amounts.iter()).enumerate() {
if *ts > now {
break;
}
vested_amount += *amount;
if vested_amount >= requested_amount {
amount_to_subtract = *amount + requested_amount - vested_amount;
if vested_amount == requested_amount {
remaining_first_idx += 1;
amount_to_subtract = 0;
}
vested_amount = requested_amount;
break;
}
remaining_first_idx = i + 1;
}
if earliest_unvested_ts_idx == 0 {
return Ok(0);
if vested_amount < requested_amount {
return Err(ContractError::NoSufficientUnlockedTokens{});
}
let vesting_amounts = VESTING_AMOUNTS.load(storage)?;
let vested_amounts = vesting_amounts[..earliest_unvested_ts_idx].to_vec();
let remaining_vesting_amounts = vesting_amounts[earliest_unvested_ts_idx..].to_vec();
let mut total_vested_amount = 0u128;
for vested_amount in vested_amounts.iter() {
total_vested_amount += *vested_amount;
if remaining_first_idx >= vesting_amounts.len() {
VESTING_AMOUNTS.save(storage, &vec![])?;
VESTING_TIMESTAMPS.save(storage, &vec![])?;
} else {
let mut remaining_amounts = vesting_amounts[remaining_first_idx..].to_vec();
if amount_to_subtract > 0 {
remaining_amounts[0] -= amount_to_subtract;
}
VESTING_AMOUNTS.save(storage, &remaining_amounts)?;
let remaining_ts = vesting_ts[remaining_first_idx..].to_vec();
VESTING_TIMESTAMPS.save(storage, &remaining_ts)?;
}
let remaining_vesting_ts = vesting_ts[earliest_unvested_ts_idx..].to_vec();
VESTING_AMOUNTS.save(storage, &remaining_vesting_amounts)?;
VESTING_TIMESTAMPS.save(storage, &remaining_vesting_ts)?;

Ok(total_vested_amount)
Ok(vested_amount)
}

pub fn total_vested_amount(storage: &dyn Storage, now: Timestamp) -> StdResult<u128> {
Expand Down Expand Up @@ -67,6 +83,7 @@ mod tests {
use cosmwasm_std::{Addr, Response};

use crate::state::{DENOM, UNLOCK_DISTRIBUTION_ADDRESS, VESTING_AMOUNTS, VESTING_TIMESTAMPS};
use crate::ContractError;

use super::{collect_vested, distribute_vested};

Expand All @@ -78,11 +95,12 @@ mod tests {
VESTING_TIMESTAMPS.save(deps_mut.storage, &vec![]).unwrap();
VESTING_AMOUNTS.save(deps_mut.storage, &vec![]).unwrap();

assert_eq!(0, collect_vested(deps_mut.storage, now).unwrap());
assert_eq!(0, collect_vested(deps_mut.storage, now, 0).unwrap());
assert_eq!(ContractError::NoSufficientUnlockedTokens {}, collect_vested(deps_mut.storage, now, 10).expect_err("should error"));
}

#[test]
fn test_vest_single() {
fn test_vest_single_full() {
let mut deps = mock_dependencies();
let deps_mut = deps.as_mut();
let now = mock_env().block.time;
Expand All @@ -91,11 +109,41 @@ mod tests {
.unwrap();
VESTING_AMOUNTS.save(deps_mut.storage, &vec![10]).unwrap();

assert_eq!(10, collect_vested(deps_mut.storage, now).unwrap());
assert_eq!(10, collect_vested(deps_mut.storage, now, 10).unwrap());
assert_eq!(VESTING_TIMESTAMPS.load(deps_mut.storage).unwrap(), vec![]);
assert_eq!(VESTING_AMOUNTS.load(deps_mut.storage).unwrap(), vec![]);
}

#[test]
fn test_vest_single_more() {
let mut deps = mock_dependencies();
let deps_mut = deps.as_mut();
let now = mock_env().block.time;
VESTING_TIMESTAMPS
.save(deps_mut.storage, &vec![now])
.unwrap();
VESTING_AMOUNTS.save(deps_mut.storage, &vec![10]).unwrap();

assert_eq!(ContractError::NoSufficientUnlockedTokens {}, collect_vested(deps_mut.storage, now, 15).expect_err("should error"));
assert_eq!(VESTING_TIMESTAMPS.load(deps_mut.storage).unwrap(), vec![now]);
assert_eq!(VESTING_AMOUNTS.load(deps_mut.storage).unwrap(), vec![10]);
}

#[test]
fn test_vest_single_partial() {
let mut deps = mock_dependencies();
let deps_mut = deps.as_mut();
let now = mock_env().block.time;
VESTING_TIMESTAMPS
.save(deps_mut.storage, &vec![now])
.unwrap();
VESTING_AMOUNTS.save(deps_mut.storage, &vec![10]).unwrap();

assert_eq!(5, collect_vested(deps_mut.storage, now, 5).unwrap());
assert_eq!(VESTING_TIMESTAMPS.load(deps_mut.storage).unwrap(), vec![now]);
assert_eq!(VESTING_AMOUNTS.load(deps_mut.storage).unwrap(), vec![5]);
}

#[test]
fn test_not_vest_single() {
let mut deps = mock_dependencies();
Expand All @@ -106,7 +154,7 @@ mod tests {
.unwrap();
VESTING_AMOUNTS.save(deps_mut.storage, &vec![10]).unwrap();

assert_eq!(0, collect_vested(deps_mut.storage, now).unwrap());
assert_eq!(ContractError::NoSufficientUnlockedTokens {}, collect_vested(deps_mut.storage, now, 10).expect_err("should error"));
assert_eq!(
VESTING_TIMESTAMPS.load(deps_mut.storage).unwrap(),
vec![now.plus_seconds(1)]
Expand All @@ -129,14 +177,24 @@ mod tests {
.save(deps_mut.storage, &vec![10, 9, 11])
.unwrap();

assert_eq!(19, collect_vested(deps_mut.storage, now).unwrap());
assert_eq!(18, collect_vested(deps_mut.storage, now, 18).unwrap());
assert_eq!(
VESTING_TIMESTAMPS.load(deps_mut.storage).unwrap(),
vec![now, now.plus_seconds(1)]
);
assert_eq!(
VESTING_AMOUNTS.load(deps_mut.storage).unwrap(),
vec![1u128, 11u128]
);

assert_eq!(2, collect_vested(deps_mut.storage, now.plus_seconds(1), 2).unwrap());
assert_eq!(
VESTING_TIMESTAMPS.load(deps_mut.storage).unwrap(),
vec![now.plus_seconds(1)]
);
assert_eq!(
VESTING_AMOUNTS.load(deps_mut.storage).unwrap(),
vec![11u128]
vec![10u128]
);
}

Expand All @@ -155,7 +213,7 @@ mod tests {
.save(deps_mut.storage, &vec![10, 9, 11])
.unwrap();

assert_eq!(30, collect_vested(deps_mut.storage, now).unwrap());
assert_eq!(30, collect_vested(deps_mut.storage, now, 30).unwrap());
assert_eq!(VESTING_TIMESTAMPS.load(deps_mut.storage).unwrap(), vec![]);
assert_eq!(VESTING_AMOUNTS.load(deps_mut.storage).unwrap(), vec![]);
}
Expand All @@ -179,7 +237,7 @@ mod tests {
.save(deps_mut.storage, &vec![10, 9, 11])
.unwrap();

assert_eq!(0, collect_vested(deps_mut.storage, now).unwrap());
assert_eq!(ContractError::NoSufficientUnlockedTokens {}, collect_vested(deps_mut.storage, now, 30).expect_err("should error"));
assert_eq!(
VESTING_TIMESTAMPS.load(deps_mut.storage).unwrap(),
vec![
Expand Down

0 comments on commit f39cabb

Please sign in to comment.