Fix hardcoded period logic (#1104)

contracts/vesting/src/vesting/account/mod.rs

@@ -22,6 +22,13 @@ fn generate_storage_key(storage: &mut dyn Storage) -> Result<u32, ContractError>
     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::<Vec<u64>>()
-            .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)
         }
     }
 

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<Coin, ContractError> {
         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<Coin, ContractError> {
         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,

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]