Add ensure_can_withdraw to BasicCurrency & MultiCurrency. (#63)

currencies/src/lib.rs

@@ -5,7 +5,7 @@ use frame_support::{
 	traits::{Currency as PalletCurrency, ExistenceRequirement, Get, WithdrawReason},
 };
 use rstd::{convert::TryInto, marker, result};
-use sp_runtime::traits::StaticLookup;
+use sp_runtime::traits::{CheckedSub, StaticLookup};
 // FIXME: `pallet/frame-` prefix should be used for all pallet modules, but currently `frame_system`
 // would cause compiling error in `decl_module!` and `construct_runtime!`
 // #3295 https://github.com/paritytech/substrate/issues/3295
@@ -55,6 +55,7 @@ decl_error! {
 	/// Error for currencies module.
 	pub enum Error {
 		AmountIntoBalanceFailed,
+		BalanceTooLow,
 	}
 }
 
@@ -118,6 +119,18 @@ impl<T: Trait> MultiCurrency<T::AccountId> for Module<T> {
 		}
 	}
 
+	fn ensure_can_withdraw(
+		currency_id: Self::CurrencyId,
+		who: &T::AccountId,
+		amount: Self::Balance,
+	) -> result::Result<(), Self::Error> {
+		if currency_id == T::GetNativeCurrencyId::get() {
+			T::NativeCurrency::ensure_can_withdraw(who, amount)
+		} else {
+			T::MultiCurrency::ensure_can_withdraw(currency_id, who, amount)
+		}
+	}
+
 	fn transfer(
 		currency_id: Self::CurrencyId,
 		from: &T::AccountId,
@@ -198,6 +211,10 @@ where
 		<Module<T>>::balance(GetCurrencyId::get(), who)
 	}
 
+	fn ensure_can_withdraw(who: &T::AccountId, amount: Self::Balance) -> result::Result<(), Self::Error> {
+		<Module<T>>::ensure_can_withdraw(GetCurrencyId::get(), who, amount)
+	}
+
 	fn transfer(from: &T::AccountId, to: &T::AccountId, amount: Self::Balance) -> result::Result<(), Self::Error> {
 		<Module<T> as MultiCurrency<T::AccountId>>::transfer(GetCurrencyId::get(), from, to, amount)
 	}
@@ -259,6 +276,18 @@ where
 		BalanceConvert::from(Currency::total_balance(who)).into()
 	}
 
+	fn ensure_can_withdraw(who: &AccountId, amount: Self::Balance) -> Result<(), Self::Error> {
+		let new_balance_pallet = {
+			let new_balance = Self::balance(who)
+				.checked_sub(&amount)
+				.ok_or(ErrorConvert::from(Error::BalanceTooLow.into()).into())?;
+			BalanceConvert::from(new_balance).into()
+		};
+		let amount_pallet = BalanceConvert::from(amount).into();
+		Currency::ensure_can_withdraw(who, amount_pallet, WithdrawReason::Transfer.into(), new_balance_pallet)
+			.map_err(|err| ErrorConvert::from(err).into())
+	}
+
 	fn transfer(from: &AccountId, to: &AccountId, amount: Self::Balance) -> result::Result<(), Self::Error> {
 		let amount_pallet = BalanceConvert::from(amount).into();
 		Currency::transfer(from, to, amount_pallet, ExistenceRequirement::AllowDeath)

tokens/src/lib.rs

@@ -103,8 +103,8 @@ decl_error! {
 impl<T: Trait> Module<T> {}
 
 impl<T: Trait> MultiCurrency<T::AccountId> for Module<T> {
-	type Balance = T::Balance;
 	type CurrencyId = T::CurrencyId;
+	type Balance = T::Balance;
 	type Error = Error;
 
 	fn total_issuance(currency_id: Self::CurrencyId) -> Self::Balance {
@@ -115,6 +115,18 @@ impl<T: Trait> MultiCurrency<T::AccountId> for Module<T> {
 		<Balance<T>>::get(currency_id, who)
 	}
 
+	fn ensure_can_withdraw(
+		currency_id: Self::CurrencyId,
+		who: &T::AccountId,
+		amount: Self::Balance,
+	) -> result::Result<(), Self::Error> {
+		if Self::balance(currency_id, who).checked_sub(&amount).is_some() {
+			Ok(())
+		} else {
+			Err(Error::BalanceTooLow)
+		}
+	}
+
 	fn transfer(
 		currency_id: Self::CurrencyId,
 		from: &T::AccountId,

tokens/src/tests.rs

@@ -105,3 +105,19 @@ fn update_balance_should_work() {
 			assert_noop!(Tokens::update_balance(TEST_TOKEN_ID, &BOB, -60), Error::BalanceTooLow);
 		});
 }
+
+#[test]
+fn ensure_can_withdraw_should_work() {
+	ExtBuilder::default()
+		.one_hundred_for_alice_n_bob()
+		.build()
+		.execute_with(|| {
+			assert_noop!(
+				Tokens::ensure_can_withdraw(TEST_TOKEN_ID, &ALICE, 101),
+				Error::BalanceTooLow
+			);
+
+			assert_ok!(Tokens::ensure_can_withdraw(TEST_TOKEN_ID, &ALICE, 1));
+			assert_eq!(Tokens::balance(TEST_TOKEN_ID, &ALICE), 100);
+		});
+}

traits/src/lib.rs

@@ -33,6 +33,13 @@ pub trait MultiCurrency<AccountId> {
 	/// The combined balance of `who` under `currency_id`.
 	fn balance(currency_id: Self::CurrencyId, who: &AccountId) -> Self::Balance;
 
+	/// A dry-run of `withdraw`. Returns `Ok` iff the account is able to make a withdrawal of the given amount.
+	fn ensure_can_withdraw(
+		currency_id: Self::CurrencyId,
+		who: &AccountId,
+		amount: Self::Balance,
+	) -> result::Result<(), Self::Error>;
+
 	// Public mutables
 
 	/// Transfer some amount from one account to another.
@@ -101,6 +108,9 @@ pub trait BasicCurrency<AccountId> {
 	/// The balance of `who`.
 	fn balance(who: &AccountId) -> Self::Balance;
 
+	/// A dry-run of `withdraw`. Returns `Ok` iff the account is able to make a withdrawal of the given amount.
+	fn ensure_can_withdraw(who: &AccountId, amount: Self::Balance) -> result::Result<(), Self::Error>;
+
 	// Public mutables
 
 	/// Transfer some amount from one account to another.