From ba233b601651f7ec896fdcc8af56b3a09016ede7 Mon Sep 17 00:00:00 2001 From: Drazen Urch Date: Fri, 11 Feb 2022 08:44:15 +0000 Subject: [PATCH] Fix hardcoded period logic (#1104) --- contracts/vesting/src/vesting/account/mod.rs | 34 +++++++++++-------- .../src/vesting/account/vesting_account.rs | 32 ++++++++++------- contracts/vesting/src/vesting/mod.rs | 20 +++++++++-- 3 files changed, 56 insertions(+), 30 deletions(-) diff --git a/contracts/vesting/src/vesting/account/mod.rs b/contracts/vesting/src/vesting/account/mod.rs index 0e2a48df41..5c14cbba42 100644 --- a/contracts/vesting/src/vesting/account/mod.rs +++ b/contracts/vesting/src/vesting/account/mod.rs @@ -22,6 +22,13 @@ fn generate_storage_key(storage: &mut dyn Storage) -> Result Ok(key) } +#[derive(Debug, PartialEq)] +pub enum Period { + Before, + In(usize), + After, +} + #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)] pub struct Account { owner_address: Addr, @@ -92,24 +99,23 @@ impl Account { } } - pub fn get_current_vesting_period(&self, block_time: Timestamp) -> usize { + pub fn get_current_vesting_period(&self, block_time: Timestamp) -> Period { // Returns the index of the next vesting period. Unless the current time is somehow in the past or vesting has not started yet. // In case vesting is over it will always return NUM_VESTING_PERIODS. - let period = match self - .periods - .iter() - .map(|period| period.start_time) - .collect::>() - .binary_search(&block_time.seconds()) - { - Ok(u) => u, - Err(u) => u, - }; - if period > 0 { - period - 1 + if block_time.seconds() < self.periods.first().unwrap().start_time { + Period::Before + } else if self.periods.last().unwrap().end_time() < block_time { + Period::After } else { - 0 + let mut index = 0; + for period in &self.periods { + if block_time < period.end_time() { + break; + } + index += 1; + } + Period::In(index) } } diff --git a/contracts/vesting/src/vesting/account/vesting_account.rs b/contracts/vesting/src/vesting/account/vesting_account.rs index 5abe3ef8ff..0c3832b382 100644 --- a/contracts/vesting/src/vesting/account/vesting_account.rs +++ b/contracts/vesting/src/vesting/account/vesting_account.rs @@ -4,7 +4,7 @@ use crate::traits::VestingAccount; use config::defaults::DENOM; use cosmwasm_std::{Addr, Coin, Env, Order, Storage, Timestamp, Uint128}; -use super::Account; +use super::{Account, Period}; impl VestingAccount for Account { fn locked_coins( @@ -61,17 +61,15 @@ impl VestingAccount for Account { let period = self.get_current_vesting_period(block_time); let amount = match period { - // We're in the first period, or the vesting has not started yet. - 0 => Coin { + Period::Before => Coin { amount: Uint128::new(0), denom: DENOM.to_string(), }, - // We always have 8 vesting periods, so periods 1-7 are special - 1..=7 => Coin { - amount: Uint128::new(self.tokens_per_period()? * period as u128), + Period::In(idx) => Coin { + amount: Uint128::new(self.tokens_per_period()? * idx as u128), denom: DENOM.to_string(), }, - _ => Coin { + Period::After => Coin { amount: self.coin.amount, denom: DENOM.to_string(), }, @@ -111,8 +109,12 @@ impl VestingAccount for Account { ) -> Result { let block_time = block_time.unwrap_or(env.block.time); let period = self.get_current_vesting_period(block_time); - let max_vested = self.tokens_per_period()? * period as u128; - let start_time = self.periods[period].start_time; + let max_vested = self.get_vested_coins(Some(block_time), env)?; + let start_time = match period { + Period::Before => 0, + Period::After => u64::MAX, + Period::In(idx) => self.periods[idx as usize].start_time, + }; let coin = DELEGATIONS .sub_prefix(self.storage_key()) @@ -123,7 +125,7 @@ impl VestingAccount for Account { acc + amount }); - let amount = Uint128::new(coin.u128().min(max_vested)); + let amount = Uint128::new(coin.u128().min(max_vested.amount.u128())); Ok(Coin { amount, @@ -157,8 +159,12 @@ impl VestingAccount for Account { ) -> Result { let block_time = block_time.unwrap_or(env.block.time); let period = self.get_current_vesting_period(block_time); - let max_vested = self.tokens_per_period()? * period as u128; - let start_time = self.periods[period].start_time; + let max_vested = self.get_vested_coins(Some(block_time), env)?; + let start_time = match period { + Period::Before => 0, + Period::After => u64::MAX, + Period::In(idx) => self.periods[idx as usize].start_time, + }; let amount = if let Some(bond) = self .load_mixnode_pledge(storage)? @@ -173,7 +179,7 @@ impl VestingAccount for Account { Uint128::zero() }; - let amount = Uint128::new(amount.u128().min(max_vested)); + let amount = Uint128::new(amount.u128().min(max_vested.amount.u128())); Ok(Coin { amount, diff --git a/contracts/vesting/src/vesting/mod.rs b/contracts/vesting/src/vesting/mod.rs index 7b6c260f7f..8293ef2063 100644 --- a/contracts/vesting/src/vesting/mod.rs +++ b/contracts/vesting/src/vesting/mod.rs @@ -48,6 +48,7 @@ mod tests { use crate::traits::DelegatingAccount; use crate::traits::VestingAccount; use crate::traits::{GatewayBondingAccount, MixnodeBondingAccount}; + use crate::vesting::Period; use config::defaults::DENOM; use cosmwasm_std::testing::{mock_env, mock_info}; use cosmwasm_std::{coins, Addr, Coin, Timestamp, Uint128}; @@ -182,12 +183,12 @@ mod tests { assert_eq!(account.periods().len(), num_vesting_periods as usize); let current_period = account.get_current_vesting_period(Timestamp::from_seconds(0)); - assert_eq!(0, current_period); + assert_eq!(Period::Before, current_period); let block_time = Timestamp::from_seconds(account.start_time().seconds() + vesting_period + 1); let current_period = account.get_current_vesting_period(block_time); - assert_eq!(current_period, 1); + assert_eq!(current_period, Period::In(1)); let vested_coins = account.get_vested_coins(Some(block_time), &env).unwrap(); let vesting_coins = account.get_vesting_coins(Some(block_time), &env).unwrap(); assert_eq!( @@ -207,7 +208,7 @@ mod tests { let block_time = Timestamp::from_seconds(account.start_time().seconds() + 5 * vesting_period + 1); let current_period = account.get_current_vesting_period(block_time); - assert_eq!(current_period, 5); + assert_eq!(current_period, Period::In(5)); let vested_coins = account.get_vested_coins(Some(block_time), &env).unwrap(); let vesting_coins = account.get_vesting_coins(Some(block_time), &env).unwrap(); assert_eq!( @@ -224,6 +225,19 @@ mod tests { / num_vesting_periods as u128 ) ); + let vesting_over_period = num_vesting_periods + 1; + let block_time = Timestamp::from_seconds( + account.start_time().seconds() + vesting_over_period * vesting_period + 1, + ); + let current_period = account.get_current_vesting_period(block_time); + assert_eq!(current_period, Period::After); + let vested_coins = account.get_vested_coins(Some(block_time), &env).unwrap(); + let vesting_coins = account.get_vesting_coins(Some(block_time), &env).unwrap(); + assert_eq!( + vested_coins.amount, + Uint128::new(account.get_original_vesting().amount.u128()) + ); + assert_eq!(vesting_coins.amount, Uint128::zero()); } #[test]