withdraw Ring/Kton automatically

srml/kton/src/lib.rs

@@ -26,7 +26,7 @@ use system::ensure_signed;
 
 use darwinia_support::{
 	traits::LockableCurrency,
-	types::{CompositeLock, LockUpdateStrategy},
+	types::{CompositeLock, Lock},
 };
 use imbalance::{NegativeImbalance, PositiveImbalance};
 
@@ -105,7 +105,7 @@ decl_storage! {
 
 		pub ReservedBalance get(reserved_balance): map T::AccountId => T::Balance;
 
-		pub Locks get(locks): map T::AccountId => CompositeLock<T::Balance, T::Moment>;
+		pub Locks get(locks): map T::AccountId => BalanceLocks<T::Balance, T::Moment>;
 
 		pub TotalLock get(total_lock): T::Balance;
 
@@ -374,71 +374,26 @@ impl<T: Trait> LockableCurrency<T::AccountId> for Module<T>
 where
 	T::Balance: MaybeSerializeDeserialize + Debug,
 {
-	type LockUpdateStrategy = LockUpdateStrategy<T::Balance, Self::Moment>;
+	type Lock = Lock<T::Balance, Self::Moment>;
 	type Moment = T::Moment;
 	type WithdrawReasons = WithdrawReasons;
 
-	fn update_lock(who: &T::AccountId, lock_update_strategy: Self::LockUpdateStrategy) -> Self::Balance {
-		let now = <timestamp::Module<T>>::now();
-		let mut expired_locks_amount = Self::Balance::default();
-		let mut composite_lock = Self::locks(who);
-
-		if let Some(staking_amount) = lock_update_strategy.staking_amount {
-			composite_lock.staking_amount = staking_amount;
-		}
-		if let Some(lock) = lock_update_strategy.lock {
-			let at = lock.at;
-			let mut new_lock = Some(lock);
-			composite_lock.locks = composite_lock
-				.locks
-				.into_iter()
-				.filter_map(|lock| {
-					if lock.at == at {
-						new_lock.take()
-					} else if lock.valid_at(now) {
-						Some(lock)
-					} else {
-						// TODO: check overflow?
-						expired_locks_amount += lock.amount;
-						None
-					}
-				})
-				.collect::<Vec<_>>();
-			if let Some(lock) = new_lock {
-				composite_lock.locks.push(lock);
-			}
-		} else if lock_update_strategy.check_expired {
-			composite_lock.locks.retain(|lock| {
-				if lock.valid_at(now) {
-					true
-				} else {
-					expired_locks_amount += lock.amount;
-					false
-				}
-			});
-		}
-
-		<Locks<T>>::insert(who, composite_lock);
+	fn update_lock(who: &T::AccountId, lock: Option<Self::Lock>) -> Self::Balance {
+		let at = <timestamp::Module<T>>::now();
+		let mut locks = Self::locks(who);
+		let expired_locks_amount = if let Some(lock) = lock {
+			locks.update_lock(lock, at)
+		} else {
+			locks.remove_expired_lock(at)
+		};
+		<Locks<T>>::insert(who, locks);
 
 		expired_locks_amount
 	}
 
-	fn remove_lock(who: &T::AccountId, at: Self::Moment) -> Self::Balance {
-		let now = <timestamp::Module<T>>::now();
-		let mut expired_locks_amount = Self::Balance::default();
-
-		<Locks<T>>::mutate(who, |composite_lock| {
-			composite_lock.locks.retain(|lock| {
-				if lock.valid_at(now) && lock.at != at {
-					true
-				} else {
-					expired_locks_amount += lock.amount;
-					false
-				}
-			});
-		});
-
-		expired_locks_amount
+	fn remove_locks(who: &T::AccountId, lock: Self::Lock) -> Self::Balance {
+		let at = <timestamp::Module<T>>::now();
+		locks.remove_locks(at, lock)
 	}
 
 	fn can_withdraw(who: &T::AccountId, reasons: Self::WithdrawReasons, new_balance: Self::Balance) -> bool {

srml/support/src/traits.rs

@@ -23,7 +23,7 @@ pub trait OnAccountBalanceChanged<AccountId, Balance> {
 
 /// A more powerful lockable currency.
 pub trait LockableCurrency<AccountId>: Currency<AccountId> {
-	type LockUpdateStrategy;
+	type Lock;
 	/// The quantity used to denote time; usually just a `BlockNumber`.
 	/// In Darwinia we prefer using `TimeStamp/u64`.
 	type Moment;
@@ -36,15 +36,15 @@ pub trait LockableCurrency<AccountId>: Currency<AccountId> {
 	/// - Remove the expired locks on account `who`.
 	/// - Update the global staking amount.
 	/// - The function will return the sum of expired locks' amount.
-	fn update_lock(who: &AccountId, strategy: Self::LockUpdateStrategy) -> Self::Balance;
+	fn update_lock(who: &AccountId, strategy: Self::Lock) -> Self::Balance;
 
 	// TODO: reserve
 	// fn extend_lock();
 
 	/// Remove an existing lock.
 	///
 	/// The function will return the sum of expired locks' amount.
-	fn remove_lock(who: &AccountId, at: Self::Moment) -> Self::Balance;
+	fn remove_locks(who: &AccountId, lock: Self::Lock) -> Self::Balance;
 
 	fn can_withdraw(who: &AccountId, reasons: Self::WithdrawReasons, new_balance: Self::Balance) -> bool;
 

srml/support/src/types.rs

@@ -1,54 +1,83 @@
 use codec::{Decode, Encode};
 use rstd::vec::Vec;
-use sr_primitives::{traits::Zero, RuntimeDebug};
+use sr_primitives::{traits::SimpleArithmetic, RuntimeDebug};
 use srml_support::traits::WithdrawReasons;
 
 pub type TimeStamp = u64;
 
-#[derive(Clone, PartialEq, Default, Encode, Decode, RuntimeDebug)]
-pub struct CompositeLock<Balance, Moment> {
-	pub staking_amount: Balance,
-	pub locks: Vec<BalanceLock<Balance, Moment>>,
-}
+#[derive(Clone, Default, Encode, Decode, RuntimeDebug)]
+pub struct BalanceLocks<Balance, Moment>(Vec<Lock<Balance, Moment>>);
 
-impl<Balance, Moment> CompositeLock<Balance, Moment>
+impl<Balance, Moment> BalanceLocks<Balance, Moment>
 where
-	Balance: Zero,
+	Balance: Clone + Copy + Default + SimpleArithmetic,
+	Moment: Clone + Copy + PartialOrd,
 {
-	pub fn is_empty(&self) -> bool {
-		self.staking_amount.is_zero() && self.locks.is_empty()
+	#[inline]
+	fn update_lock(&mut self, lock: Lock<Balance, Moment>, at: Moment) -> Balance {
+		let expired_locks_amount = self.remove_expired_locks(at);
+		self.0.push(lock);
+
+		expired_locks_amount
 	}
-}
 
-pub struct LockUpdateStrategy<Balance, Moment> {
-	/// if `lock` is set, `check_expired` will be ignored
-	pub check_expired: bool,
-	pub staking_amount: Option<Balance>,
-	pub lock: Option<BalanceLock<Balance, Moment>>,
-}
+	fn remove_locks(&mut self, at: Moment, lock: Lock<Balance, Moment>) -> Balance {
+		let mut expired_locks_amount = Self::Balance::zero();
 
-impl<Balance, Moment> LockUpdateStrategy<Balance, Moment> {
-	pub fn new() -> Self {
-		Self {
-			check_expired: false,
-			staking_amount: None,
-			lock: None,
-		}
+		<Locks<T>>::mutate(who, |locks| {
+			locks.retain(|lock_| {
+				if lock_.valid_at(now) && lock == lock {
+					true
+				} else {
+					expired_locks_amount += lock.amount;
+					false
+				}
+			});
+		});
+
+		expired_locks_amount
 	}
 
-	pub fn with_check_expired(mut self, check_expired: bool) -> Self {
-		self.check_expired = check_expired;
-		self
+	fn remove_expired_locks(&mut self, at: Moment) -> Balance {
+		let mut expired_locks_amount = Balance::default();
+		self.0.retain(|lock| {
+			if lock.valid_at(at) {
+				true
+			} else {
+				expired_locks_amount += lock.amount();
+				false
+			}
+		});
+
+		expired_locks_amount
 	}
+}
 
-	pub fn with_staking_amount(mut self, staking_amount: Balance) -> Self {
-		self.staking_amount = Some(staking_amount);
-		self
+#[derive(Clone, RuntimeDebug)]
+pub enum Lock<Balance, Moment> {
+	Staking(Balance),
+	Unbonding(BalanceLock<Balance, Moment>),
+}
+
+impl<Balance, Moment> Lock<Balance, Moment>
+where
+	Balance: Copy,
+	Moment: PartialOrd,
+{
+	#[inline]
+	fn valid_at(&self, at: Moment) -> bool {
+		match self {
+			Lock::Staking(_) => true,
+			Lock::Unbonding(balance_lock) => balance_lock.at > at,
+		}
 	}
 
-	pub fn with_lock(mut self, lock: BalanceLock<Balance, Moment>) -> Self {
-		self.lock = Some(lock);
-		self
+	#[inline]
+	fn amount(&self) -> Balance {
+		match self {
+			Lock::Staking(balance) => *balance,
+			Lock::Unbonding(balance_lock) => balance_lock.amount,
+		}
 	}
 }
 
@@ -58,12 +87,3 @@ pub struct BalanceLock<Balance, Moment> {
 	pub at: Moment,
 	pub reasons: WithdrawReasons,
 }
-
-impl<Balance, Moment> BalanceLock<Balance, Moment>
-where
-	Moment: PartialOrd,
-{
-	pub fn valid_at(&self, at: Moment) -> bool {
-		self.at > at
-	}
-}