refactor migrate functions

substrate/frame/delegated-staking/src/lib.rs

@@ -201,6 +201,52 @@ impl<T: Config> Delegatee for Pallet<T> {
 		}
 	}
 
+	fn withdraw(
+		delegator: &Self::AccountId,
+		delegatee: &Self::AccountId,
+		value: Self::Balance,
+	) -> sp_runtime::DispatchResult {
+		<Delegators<T>>::mutate_exists(delegator, |maybe_delegate| match maybe_delegate {
+			Some((current_delegatee, delegate_balance)) => {
+				ensure!(&current_delegatee.clone() == delegatee, Error::<T>::NotDelegatee);
+				ensure!(*delegate_balance >= value, Error::<T>::NotAllowed);
+
+				delegate_balance.saturating_reduce(value);
+
+				if *delegate_balance == BalanceOf::<T>::zero() {
+					*maybe_delegate = None;
+				}
+				Ok(())
+			},
+			None => {
+				// delegator does not exist
+				return Err(Error::<T>::NotAllowed)
+			},
+		})?;
+
+		<Delegatees<T>>::mutate(delegatee, |maybe_register| match maybe_register {
+			Some(ledger) => {
+				ledger.balance.saturating_reduce(value);
+				Ok(())
+			},
+			None => {
+				// Delegatee not found
+				return Err(Error::<T>::NotDelegatee)
+			},
+		})?;
+
+		let released = T::Currency::release(
+			&HoldReason::Delegating.into(),
+			&delegator,
+			value,
+			Precision::BestEffort,
+		)?;
+
+		defensive_assert!(released == value, "hold should have been released fully");
+
+		Ok(())
+	}
+
 	fn apply_slash(
 		delegatee: &Self::AccountId,
 		delegator: &Self::AccountId,
@@ -243,8 +289,13 @@ impl<T: Config> Delegatee for Pallet<T> {
 
 		// try transferring the staked amount. This should never fail but if it does, it indicates
 		// bad state and we abort.
-		T::Currency::transfer(new_delegatee, proxy_delegator, stake.total, Preservation::Expendable)
-			.map_err(|_| Error::<T>::BadState)?;
+		T::Currency::transfer(
+			new_delegatee,
+			proxy_delegator,
+			stake.total,
+			Preservation::Expendable,
+		)
+		.map_err(|_| Error::<T>::BadState)?;
 
 		// delegate from new delegator to staker.
 		Self::accept_delegations(new_delegatee, payee)?;
@@ -302,58 +353,27 @@ impl<T: Config> Delegator for Pallet<T> {
 		todo!()
 	}
 
-	fn withdraw(
-		delegator: &Self::AccountId,
+	/// Move funds from proxy delegator to actual delegator.
+	// TODO: Keep track of proxy delegator and only allow movement from proxy -> new delegator
+	fn delegator_migrate(
+		existing_delegator: &Self::AccountId,
+		new_delegator: &Self::AccountId,
 		delegatee: &Self::AccountId,
 		value: Self::Balance,
 	) -> sp_runtime::DispatchResult {
-		<Delegators<T>>::mutate_exists(delegator, |maybe_delegate| match maybe_delegate {
-			Some((current_delegatee, delegate_balance)) => {
-				ensure!(&current_delegatee.clone() == delegatee, Error::<T>::InvalidDelegation);
-				ensure!(*delegate_balance >= value, Error::<T>::InvalidDelegation);
-
-				delegate_balance.saturating_reduce(value);
-
-				if *delegate_balance == BalanceOf::<T>::zero() {
-					*maybe_delegate = None;
-				}
-				Ok(())
-			},
-			None => {
-				// this should never happen
-				return Err(Error::<T>::InvalidDelegation)
-			},
-		})?;
-
-		<Delegatees<T>>::mutate(delegatee, |maybe_register| match maybe_register {
-			Some(ledger) => {
-				ledger.balance.saturating_reduce(value);
-				Ok(())
-			},
-			None => {
-				// this should never happen
-				return Err(Error::<T>::InvalidDelegation)
-			},
-		})?;
+		ensure!(value >= T::Currency::minimum_balance(), Error::<T>::NotEnoughFunds);
 
-		let released = T::Currency::release(
-			&HoldReason::Delegating.into(),
-			&delegator,
-			value,
-			Precision::BestEffort,
-		)?;
+		// ensure delegatee exists.
+		ensure!(!<Delegatees<T>>::contains_key(delegatee), Error::<T>::NotDelegatee);
 
-		defensive_assert!(released == value, "hold should have been released fully");
+		// remove delegation of `value` from `existing_delegator`.
+		Self::withdraw(existing_delegator, delegatee, value)?;
 
-		Ok(())
-	}
+		// transfer the withdrawn value to `new_delegator`.
+		T::Currency::transfer(existing_delegator, new_delegator, value, Preservation::Expendable)
+			.map_err(|_| Error::<T>::BadState)?;
 
-	fn delegator_migrate(
-		delegator_from: &Self::AccountId,
-		delegator_to: &Self::AccountId,
-		delegatee: &Self::AccountId,
-		value: Self::Balance,
-	) -> sp_runtime::DispatchResult {
-		todo!()
+		// add the above removed delegation to `new_delegator`.
+		Self::delegate(new_delegator, delegatee, value)
 	}
 }

substrate/primitives/staking/src/delegation.rs

@@ -17,23 +17,21 @@
 
 use codec::{FullCodec, MaxEncodedLen};
 use scale_info::TypeInfo;
-use sp_runtime::{
-	DispatchResult, Saturating,
-};
-use sp_std::{ops::Sub};
+use sp_runtime::{DispatchResult, Saturating};
+use sp_std::ops::Sub;
 
 /// Allows an account to accept stake delegations and manage its operations.
 pub trait Delegatee {
 	/// Balance type used by the staking system.
 	type Balance: Sub<Output = Self::Balance>
-	+ Ord
-	+ PartialEq
-	+ Default
-	+ Copy
-	+ MaxEncodedLen
-	+ FullCodec
-	+ TypeInfo
-	+ Saturating;
+		+ Ord
+		+ PartialEq
+		+ Default
+		+ Copy
+		+ MaxEncodedLen
+		+ FullCodec
+		+ TypeInfo
+		+ Saturating;
 
 	/// AccountId type used by the staking system.
 	type AccountId: Clone + sp_std::fmt::Debug;
@@ -45,9 +43,6 @@ pub trait Delegatee {
 	fn accept_delegations(delegatee: &Self::AccountId, payee: &Self::AccountId) -> DispatchResult;
 
 	/// Stop accepting new delegations on this account.
-	///
-	/// The account would continue to be a delegatee until all delegations to this account has been
-	/// withdrawn.
 	fn block_delegations(delegatee: &Self::AccountId) -> DispatchResult;
 
 	/// Remove oneself as Delegatee.
@@ -58,6 +53,13 @@ pub trait Delegatee {
 	/// Update bond whenever there is a new delegate funds that are not staked.
 	fn update_bond(delegatee: &Self::AccountId) -> DispatchResult;
 
+	/// Request removal of delegated stake.
+	fn withdraw(
+		delegatee: &Self::AccountId,
+		delegator: &Self::AccountId,
+		value: Self::Balance,
+	) -> DispatchResult;
+
 	/// Applies a pending slash on delegatee by passing a delegator account who should be slashed
 	/// and the value to be slashed. Optionally also takes a reporter account who will be rewarded
 	/// from part of the slash imbalance.
@@ -83,20 +85,19 @@ pub trait Delegatee {
 		proxy_delegator: &Self::AccountId,
 		payee: &Self::AccountId,
 	) -> DispatchResult;
-
 }
 
 /// Allows an account to delegate their stakes to a delegatee.
 pub trait Delegator {
 	type Balance: Sub<Output = Self::Balance>
-	+ Ord
-	+ PartialEq
-	+ Default
-	+ Copy
-	+ MaxEncodedLen
-	+ FullCodec
-	+ TypeInfo
-	+ Saturating;
+		+ Ord
+		+ PartialEq
+		+ Default
+		+ Copy
+		+ MaxEncodedLen
+		+ FullCodec
+		+ TypeInfo
+		+ Saturating;
 
 	/// AccountId type used by the staking system.
 	type AccountId: Clone + sp_std::fmt::Debug;
@@ -115,13 +116,6 @@ pub trait Delegator {
 		value: Self::Balance,
 	) -> DispatchResult;
 
-	/// Request removal of delegated stake.
-	fn withdraw (
-		delegator: &Self::AccountId,
-		delegatee: &Self::AccountId,
-		value: Self::Balance,
-	) -> DispatchResult;
-
 	/// Swap a delegated `value` from `delegator_from` to `delegator_to`, with delegatee remaining
 	/// the same.
 	///
@@ -135,4 +129,4 @@ pub trait Delegator {
 		delegatee: &Self::AccountId,
 		value: Self::Balance,
 	) -> DispatchResult;
-}
+}