[BREAKING] fix: allow vault to withdraw all collateral

crates/nomination/Cargo.toml

@@ -82,6 +82,7 @@ runtime-benchmarks = [
   "frame-support/runtime-benchmarks",
   "frame-system/runtime-benchmarks",
 
+  "traits/runtime-benchmarks",
   "vault-registry/runtime-benchmarks",
 
   "orml-tokens",

crates/nomination/src/benchmarking.rs

@@ -184,7 +184,12 @@ pub mod benchmarks {
         let balance_before = <orml_tokens::Pallet<T>>::reserved_balance(collateral_currency, &vault_id.account_id);
 
         #[extrinsic_call]
-        _(RawOrigin::Signed(nominator.clone()), vault_id.clone(), amount, None);
+        _(
+            RawOrigin::Signed(nominator.clone()),
+            vault_id.clone(),
+            Some(amount),
+            None,
+        );
 
         let balance_after = <orml_tokens::Pallet<T>>::reserved_balance(collateral_currency, &vault_id.account_id);
         assert_eq!(balance_before - amount, balance_after);

crates/nomination/src/ext.rs

@@ -22,7 +22,7 @@ pub(crate) mod vault_registry {
 
     pub fn is_allowed_to_withdraw_collateral<T: crate::Config>(
         vault_id: &DefaultVaultId<T>,
-        amount: &Amount<T>,
+        amount: Option<Amount<T>>,
     ) -> Result<bool, DispatchError> {
         <vault_registry::Pallet<T>>::is_allowed_to_withdraw_collateral(vault_id, amount)
     }
@@ -55,10 +55,10 @@ pub(crate) mod vault_registry {
         pub fn withdraw_collateral<T: crate::Config>(
             vault_id: &DefaultVaultId<T>,
             nominator_id: &T::AccountId,
-            amount: &Amount<T>,
+            maybe_amount: Option<Amount<T>>,
             nonce: Option<<T as frame_system::Config>::Index>,
-        ) -> Result<(), DispatchError> {
-            <vault_registry::PoolManager<T>>::withdraw_collateral(vault_id, nominator_id, amount, nonce)
+        ) -> Result<Amount<T>, DispatchError> {
+            <vault_registry::PoolManager<T>>::withdraw_collateral(vault_id, nominator_id, maybe_amount, nonce)
         }
 
         pub fn kick_nominators<T: crate::Config>(vault_id: &DefaultVaultId<T>) -> Result<Amount<T>, DispatchError> {
@@ -77,6 +77,7 @@ pub(crate) mod staking {
     pub fn nonce<T: crate::Config>(vault_id: &DefaultVaultId<T>) -> T::Index {
         T::VaultStaking::nonce(vault_id)
     }
+
     pub fn compute_stake<T: vault_registry::Config>(
         vault_id: &DefaultVaultId<T>,
         nominator_id: &T::AccountId,

crates/nomination/src/lib.rs

@@ -197,7 +197,7 @@ pub mod pallet {
         pub fn withdraw_collateral(
             origin: OriginFor<T>,
             vault_id: DefaultVaultId<T>,
-            amount: BalanceOf<T>,
+            amount: Option<BalanceOf<T>>,
             index: Option<T::Index>,
         ) -> DispatchResultWithPostInfo {
             let nominator_id = ensure_signed(origin)?;
@@ -228,20 +228,20 @@ impl<T: Config> Pallet<T> {
     pub fn _withdraw_collateral(
         vault_id: &DefaultVaultId<T>,
         nominator_id: &T::AccountId,
-        amount: BalanceOf<T>,
+        maybe_amount: Option<BalanceOf<T>>,
         index: T::Index,
     ) -> DispatchResult {
         let nonce = ext::staking::nonce::<T>(vault_id);
         let index = sp_std::cmp::min(index, nonce);
 
-        let amount = Amount::new(amount, vault_id.collateral_currency());
+        let maybe_amount = maybe_amount.map(|x| Amount::<T>::new(x, vault_id.collateral_currency()));
 
         // nominators are always allowed to withdraw from stale staking pools
         if index == nonce {
             // we can only withdraw nominated collateral if the vault is still
             // above the secure threshold for issued + to_be_issued tokens
             ensure!(
-                ext::vault_registry::is_allowed_to_withdraw_collateral::<T>(vault_id, &amount)?,
+                ext::vault_registry::is_allowed_to_withdraw_collateral::<T>(vault_id, maybe_amount.clone())?,
                 Error::<T>::CannotWithdrawCollateral
             );
 
@@ -251,13 +251,18 @@ impl<T: Config> Pallet<T> {
             }
         }
 
-        ext::vault_registry::decrease_total_backing_collateral(&vault_id.currencies, &amount)?;
-
         // withdraw `amount` of stake from the vault staking pool
-        ext::vault_registry::pool_manager::withdraw_collateral::<T>(vault_id, nominator_id, &amount, Some(index))?;
+        let amount = ext::vault_registry::pool_manager::withdraw_collateral::<T>(
+            vault_id,
+            nominator_id,
+            maybe_amount,
+            Some(index),
+        )?;
         amount.unlock_on(&vault_id.account_id)?;
         amount.transfer(&vault_id.account_id, &nominator_id)?;
 
+        ext::vault_registry::decrease_total_backing_collateral(&vault_id.currencies, &amount)?;
+
         Self::deposit_event(Event::<T>::WithdrawCollateral {
             vault_id: vault_id.clone(),
             nominator_id: nominator_id.clone(),
@@ -326,7 +331,7 @@ impl<T: Config> Pallet<T> {
         ensure!(Self::is_opted_in(&vault_id), Error::<T>::VaultNotOptedInToNomination);
         let total_nominated_collateral = Self::get_total_nominated_collateral(&vault_id)?;
         ensure!(
-            ext::vault_registry::is_allowed_to_withdraw_collateral::<T>(&vault_id, &total_nominated_collateral)?,
+            ext::vault_registry::is_allowed_to_withdraw_collateral::<T>(&vault_id, Some(total_nominated_collateral))?,
             Error::<T>::CollateralizationTooLow
         );
 

crates/nomination/src/tests.rs

@@ -2,6 +2,7 @@ use crate::{ext, mock::*};
 use currency::Amount;
 use frame_support::{assert_err, assert_ok};
 use mocktopus::mocking::*;
+use sp_arithmetic::FixedI128;
 
 #[test]
 fn should_not_deposit_against_invalid_vault() {
@@ -12,6 +13,7 @@ fn should_not_deposit_against_invalid_vault() {
         );
     })
 }
+
 fn collateral(amount: u128) -> Amount<Test> {
     Amount::new(amount, DEFAULT_COLLATERAL_CURRENCY)
 }
@@ -32,3 +34,64 @@ fn should_deposit_against_valid_vault() {
         assert_ok!(Nomination::_deposit_collateral(&ALICE, &BOB.account_id, 100));
     })
 }
+
+#[test]
+fn should_not_withdraw_collateral() {
+    use orml_traits::MultiCurrency;
+
+    run_test(|| {
+        assert_ok!(Tokens::deposit(
+            ALICE.currencies.collateral,
+            &ALICE.account_id,
+            24_609_778_406_619_232
+        ));
+        VaultRegistry::_set_system_collateral_ceiling(ALICE.currencies, u128::MAX);
+        assert_ok!(VaultRegistry::register_public_key(
+            RuntimeOrigin::signed(ALICE.account_id),
+            vault_registry::BtcPublicKey::dummy()
+        ));
+        assert_ok!(VaultRegistry::register_vault(
+            RuntimeOrigin::signed(ALICE.account_id),
+            ALICE.currencies,
+            24_609_778_406_619_232
+        ));
+
+        staking::SlashPerToken::<Test>::insert(0, ALICE, FixedI128::from_inner(25_210_223_519_649_666));
+        staking::SlashTally::<Test>::insert(
+            0,
+            (ALICE, ALICE.account_id),
+            FixedI128::from_inner(100_834_580_684_768_029_667_333_677_168),
+        );
+        staking::Stake::<Test>::insert(
+            0,
+            (ALICE, ALICE.account_id),
+            FixedI128::from_inner(3_999_749_570_096_999_994_120_799_432_121),
+        );
+        staking::TotalCurrentStake::<Test>::insert(
+            0,
+            ALICE,
+            FixedI128::from_inner(3_999_749_570_097_000_000_000_000_000_000),
+        );
+        staking::TotalStake::<Test>::insert(
+            0,
+            ALICE,
+            FixedI128::from_inner(3_999_749_570_096_999_994_120_799_432_121),
+        );
+
+        // should not withdraw all
+        assert_err!(
+            Nomination::_withdraw_collateral(&ALICE, &ALICE.account_id, Some(3999749570097), 0),
+            staking::Error::<Test>::InsufficientFunds
+        );
+
+        // should withdraw all
+        assert_ok!(Nomination::_withdraw_collateral(&ALICE, &ALICE.account_id, None, 0));
+
+        // stake is now zero
+        assert_ok!(ext::staking::compute_stake::<Test>(&ALICE, &ALICE.account_id), 0);
+        assert_ok!(
+            VaultRegistry::get_backing_collateral(&ALICE),
+            Amount::new(0, ALICE.collateral_currency())
+        );
+    });
+}

crates/replace/src/ext.rs

@@ -133,7 +133,7 @@ pub(crate) mod vault_registry {
 
     pub fn is_allowed_to_withdraw_collateral<T: crate::Config>(
         vault_id: &DefaultVaultId<T>,
-        amount: &Amount<T>,
+        amount: Option<Amount<T>>,
     ) -> Result<bool, DispatchError> {
         <vault_registry::Pallet<T>>::is_allowed_to_withdraw_collateral(vault_id, amount)
     }

crates/replace/src/lib.rs

@@ -598,7 +598,7 @@ impl<T: Config> Pallet<T> {
         // if the new_vault locked additional collateral especially for this replace,
         // release it if it does not cause them to be undercollateralized
         if !ext::vault_registry::is_vault_liquidated::<T>(&new_vault_id)?
-            && ext::vault_registry::is_allowed_to_withdraw_collateral::<T>(&new_vault_id, &collateral)?
+            && ext::vault_registry::is_allowed_to_withdraw_collateral::<T>(&new_vault_id, Some(collateral.clone()))?
         {
             ext::vault_registry::force_withdraw_collateral::<T>(&new_vault_id, &collateral)?;
         }

crates/reward/src/lib.rs

@@ -359,7 +359,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 
 pub trait RewardsApi<PoolId, StakeId, Balance>
 where
-    Balance: Saturating + PartialOrd,
+    Balance: Saturating + PartialOrd + Copy,
 {
     type CurrencyId;
 
@@ -389,8 +389,10 @@ where
     fn withdraw_stake(pool_id: &PoolId, stake_id: &StakeId, amount: Balance) -> DispatchResult;
 
     /// Withdraw all stake for an account.
-    fn withdraw_all_stake(pool_id: &PoolId, stake_id: &StakeId) -> DispatchResult {
-        Self::withdraw_stake(pool_id, stake_id, Self::get_stake(pool_id, stake_id)?)
+    fn withdraw_all_stake(pool_id: &PoolId, stake_id: &StakeId) -> Result<Balance, DispatchError> {
+        let amount = Self::get_stake(pool_id, stake_id)?;
+        Self::withdraw_stake(pool_id, stake_id, amount)?;
+        Ok(amount)
     }
 
     /// Return the stake associated with the `pool_id`.
@@ -418,7 +420,7 @@ impl<T, I, Balance> RewardsApi<T::PoolId, T::StakeId, Balance> for Pallet<T, I>
 where
     T: Config<I>,
     I: 'static,
-    Balance: BalanceToFixedPoint<SignedFixedPoint<T, I>> + Saturating + PartialOrd,
+    Balance: BalanceToFixedPoint<SignedFixedPoint<T, I>> + Saturating + PartialOrd + Copy,
     <T::SignedFixedPoint as FixedPointNumber>::Inner: TryInto<Balance>,
 {
     type CurrencyId = T::CurrencyId;
@@ -490,7 +492,7 @@ where
 
 impl<PoolId, StakeId, Balance> RewardsApi<PoolId, StakeId, Balance> for ()
 where
-    Balance: Saturating + PartialOrd + Default,
+    Balance: Saturating + PartialOrd + Default + Copy,
 {
     type CurrencyId = ();
 

crates/staking/src/lib.rs

@@ -641,6 +641,12 @@ impl<T: Config> Pallet<T> {
         checked_sub_mut!(TotalStake<T>, nonce, vault_id, &amount);
         checked_sub_mut!(TotalCurrentStake<T>, nonce, vault_id, &amount);
 
+        if Self::total_stake_at_index(nonce, vault_id).is_zero() {
+            // may be non-zero due to rounding, will truncate to zero
+            // but cleanup anyway
+            TotalCurrentStake::<T>::remove(nonce, vault_id);
+        }
+
         <SlashTally<T>>::mutate(nonce, (vault_id, nominator_id), |slash_tally| {
             let slash_per_token = Self::slash_per_token_at_index(nonce, vault_id);
             let slash_per_token_mul_amount = slash_per_token.checked_mul(&amount).ok_or(ArithmeticError::Overflow)?;
@@ -772,7 +778,7 @@ impl<T: Config> Pallet<T> {
 impl<T, Balance> RewardsApi<(Option<T::Index>, DefaultVaultId<T>), T::AccountId, Balance> for Pallet<T>
 where
     T: Config,
-    Balance: BalanceToFixedPoint<SignedFixedPoint<T>> + Saturating + PartialOrd,
+    Balance: BalanceToFixedPoint<SignedFixedPoint<T>> + Saturating + PartialOrd + Copy,
     <T::SignedFixedPoint as FixedPointNumber>::Inner: TryInto<Balance>,
 {
     type CurrencyId = T::CurrencyId;
@@ -845,6 +851,21 @@ where
         )
     }
 
+    fn withdraw_all_stake(
+        (nonce, vault_id): &(Option<T::Index>, DefaultVaultId<T>),
+        nominator_id: &T::AccountId,
+    ) -> Result<Balance, DispatchError> {
+        let nonce = nonce.unwrap_or(Pallet::<T>::nonce(vault_id));
+        // use the precise stake to avoid rounding issues
+        let amount = Self::compute_precise_stake_at_index(nonce, vault_id, nominator_id)?;
+        Pallet::<T>::withdraw_stake(vault_id, nominator_id, amount, Some(nonce))?;
+        amount
+            .truncate_to_inner()
+            .ok_or(Error::<T>::TryIntoIntError)?
+            .try_into()
+            .map_err(|_| Error::<T>::TryIntoIntError.into())
+    }
+
     fn get_total_stake((_, vault_id): &(Option<T::Index>, DefaultVaultId<T>)) -> Result<Balance, DispatchError> {
         Pallet::<T>::total_current_stake(vault_id)?
             .try_into()

crates/vault-registry/src/ext.rs

@@ -40,10 +40,16 @@ pub(crate) mod staking {
     pub fn withdraw_stake<T: crate::Config>(
         vault_id: &DefaultVaultId<T>,
         nominator_id: &T::AccountId,
-        amount: &Amount<T>,
+        maybe_amount: Option<Amount<T>>,
         nonce: Option<<T as frame_system::Config>::Index>,
-    ) -> DispatchResult {
-        T::VaultStaking::withdraw_stake(&(nonce, vault_id.clone()), nominator_id, amount.amount())
+    ) -> Result<Amount<T>, DispatchError> {
+        if let Some(amount) = maybe_amount {
+            T::VaultStaking::withdraw_stake(&(nonce, vault_id.clone()), nominator_id, amount.amount())?;
+            Ok(amount)
+        } else {
+            let balance = T::VaultStaking::withdraw_all_stake(&(nonce, vault_id.clone()), nominator_id)?;
+            Ok(Amount::new(balance, vault_id.collateral_currency()))
+        }
     }
 
     pub fn slash_stake<T: crate::Config>(vault_id: &DefaultVaultId<T>, amount: &Amount<T>) -> DispatchResult {

crates/vault-registry/src/lib.rs

@@ -825,22 +825,26 @@ impl<T: Config> Pallet<T> {
         Self::decrease_total_backing_collateral(&vault_id.currencies, amount)?;
 
         // Withdraw `amount` of stake from the pool
-        PoolManager::<T>::withdraw_collateral(vault_id, &vault_id.account_id, amount, None)?;
+        PoolManager::<T>::withdraw_collateral(vault_id, &vault_id.account_id, Some(amount.clone()), None)?;
 
         Ok(())
     }
 
     /// Checks if the vault would be above the secure threshold after withdrawing collateral
     pub fn is_allowed_to_withdraw_collateral(
         vault_id: &DefaultVaultId<T>,
-        amount: &Amount<T>,
+        amount: Option<Amount<T>>,
     ) -> Result<bool, DispatchError> {
         let vault = Self::get_rich_vault_from_id(vault_id)?;
 
-        let new_collateral = match Self::get_backing_collateral(vault_id)?.checked_sub(&amount) {
-            Ok(x) => x,
-            Err(x) if x == ArithmeticError::Underflow.into() => return Ok(false),
-            Err(x) => return Err(x),
+        let new_collateral = if let Some(amount) = amount {
+            match Self::get_backing_collateral(vault_id)?.checked_sub(&amount) {
+                Ok(x) => x,
+                Err(x) if x == ArithmeticError::Underflow.into() => return Ok(false),
+                Err(x) => return Err(x),
+            }
+        } else {
+            Amount::<T>::zero(vault_id.collateral_currency())
         };
 
         ensure!(

crates/vault-registry/src/pool_manager.rs

@@ -19,14 +19,16 @@ impl<T: Config> PoolManager<T> {
     pub fn withdraw_collateral(
         vault_id: &DefaultVaultId<T>,
         nominator_id: &T::AccountId,
-        amount: &Amount<T>,
+        maybe_amount: Option<Amount<T>>,
         nonce: Option<<T as frame_system::Config>::Index>,
-    ) -> Result<(), DispatchError> {
+    ) -> Result<Amount<T>, DispatchError> {
         ext::fee::distribute_all_vault_rewards::<T>(vault_id)?;
-        ext::staking::withdraw_stake(vault_id, nominator_id, amount, nonce)?;
+        let amount = ext::staking::withdraw_stake(vault_id, nominator_id, maybe_amount, nonce)?;
 
         // also propagate to reward & capacity pools
-        Self::update_reward_stake(vault_id)
+        Self::update_reward_stake(vault_id)?;
+
+        Ok(amount)
     }
 
     pub fn slash_collateral(vault_id: &DefaultVaultId<T>, amount: &Amount<T>) -> Result<(), DispatchError> {

crates/vault-registry/src/tests.rs

@@ -202,17 +202,17 @@ fn should_check_withdraw_collateral() {
 
         // should allow withdraw all
         assert_ok!(
-            VaultRegistry::is_allowed_to_withdraw_collateral(&DEFAULT_ID, &amount(DEFAULT_COLLATERAL)),
+            VaultRegistry::is_allowed_to_withdraw_collateral(&DEFAULT_ID, Some(amount(DEFAULT_COLLATERAL))),
             true,
         );
         // should allow withdraw above minimum
         assert_ok!(
-            VaultRegistry::is_allowed_to_withdraw_collateral(&DEFAULT_ID, &amount(DEFAULT_COLLATERAL / 4)),
+            VaultRegistry::is_allowed_to_withdraw_collateral(&DEFAULT_ID, Some(amount(DEFAULT_COLLATERAL / 4))),
             true,
         );
         // should not allow withdraw above zero, below minimum
         assert_err!(
-            VaultRegistry::is_allowed_to_withdraw_collateral(&DEFAULT_ID, &amount(DEFAULT_COLLATERAL / 4 * 3)),
+            VaultRegistry::is_allowed_to_withdraw_collateral(&DEFAULT_ID, Some(amount(DEFAULT_COLLATERAL / 4 * 3))),
             TestError::InsufficientVaultCollateralAmount,
         );
     });