11use super :: * ;
2+ use dispatch:: RawOrigin ;
23use frame_support:: {
34 storage:: IterableStorageDoubleMap ,
45 traits:: {
@@ -9,6 +10,7 @@ use frame_support::{
910 Imbalance ,
1011 } ,
1112} ;
13+ use num_traits:: Zero ;
1214
1315impl < T : Config > Pallet < T > {
1416 /// ---- The implementation for the extrinsic become_delegate: signals that this hotkey allows delegated stake.
@@ -560,6 +562,13 @@ impl<T: Config> Pallet<T> {
560562 if !Self :: hotkey_account_exists ( hotkey) {
561563 Stake :: < T > :: insert ( hotkey, coldkey, 0 ) ;
562564 Owner :: < T > :: insert ( hotkey, coldkey) ;
565+
566+ // Update OwnedHotkeys map
567+ let mut hotkeys = OwnedHotkeys :: < T > :: get ( coldkey) ;
568+ if !hotkeys. contains ( hotkey) {
569+ hotkeys. push ( hotkey. clone ( ) ) ;
570+ OwnedHotkeys :: < T > :: insert ( coldkey, hotkeys) ;
571+ }
563572 }
564573 }
565574
@@ -781,6 +790,31 @@ impl<T: Config> Pallet<T> {
781790 Ok ( credit)
782791 }
783792
793+ pub fn kill_coldkey_account (
794+ coldkey : & T :: AccountId ,
795+ amount : <<T as Config >:: Currency as fungible:: Inspect < <T as system:: Config >:: AccountId > >:: Balance ,
796+ ) -> Result < u64 , DispatchError > {
797+ if amount == 0 {
798+ return Ok ( 0 ) ;
799+ }
800+
801+ let credit = T :: Currency :: withdraw (
802+ coldkey,
803+ amount,
804+ Precision :: Exact ,
805+ Preservation :: Expendable ,
806+ Fortitude :: Force ,
807+ )
808+ . map_err ( |_| Error :: < T > :: BalanceWithdrawalError ) ?
809+ . peek ( ) ;
810+
811+ if credit == 0 {
812+ return Err ( Error :: < T > :: ZeroBalanceAfterWithdrawn . into ( ) ) ;
813+ }
814+
815+ Ok ( credit)
816+ }
817+
784818 pub fn unstake_all_coldkeys_from_hotkey_account ( hotkey : & T :: AccountId ) {
785819 // Iterate through all coldkeys that have a stake on this hotkey account.
786820 for ( delegate_coldkey_i, stake_i) in
@@ -795,4 +829,96 @@ impl<T: Config> Pallet<T> {
795829 Self :: add_balance_to_coldkey_account ( & delegate_coldkey_i, stake_i) ;
796830 }
797831 }
832+
833+ /// Unstakes all tokens associated with a hotkey and transfers them to a new coldkey.
834+ ///
835+ /// This function performs the following operations:
836+ /// 1. Verifies that the hotkey exists and is owned by the current coldkey.
837+ /// 2. Ensures that the new coldkey is different from the current one.
838+ /// 3. Unstakes all balance if there's any stake.
839+ /// 4. Transfers the entire balance of the hotkey to the new coldkey.
840+ /// 5. Verifies the success of the transfer and handles partial transfers if necessary.
841+ ///
842+ /// # Arguments
843+ ///
844+ /// * `current_coldkey` - The AccountId of the current coldkey.
845+ /// * `hotkey` - The AccountId of the hotkey whose balance is being unstaked and transferred.
846+ /// * `new_coldkey` - The AccountId of the new coldkey to receive the unstaked tokens.
847+ ///
848+ /// # Returns
849+ ///
850+ /// Returns a `DispatchResult` indicating success or failure of the operation.
851+ ///
852+ /// # Errors
853+ ///
854+ /// This function will return an error if:
855+ /// * The hotkey account does not exist.
856+ /// * The current coldkey does not own the hotkey.
857+ /// * The new coldkey is the same as the current coldkey.
858+ /// * There is no balance to transfer.
859+ /// * The transfer fails or is only partially successful.
860+ ///
861+ /// # Events
862+ ///
863+ /// Emits an `AllBalanceUnstakedAndTransferredToNewColdkey` event upon successful execution.
864+ /// Emits a `PartialBalanceTransferredToNewColdkey` event if only a partial transfer is successful.
865+ ///
866+ pub fn do_unstake_all_and_transfer_to_new_coldkey (
867+ current_coldkey : T :: AccountId ,
868+ hotkey : T :: AccountId ,
869+ new_coldkey : T :: AccountId ,
870+ ) -> DispatchResult {
871+ // Ensure the hotkey exists and is owned by the current coldkey
872+ ensure ! (
873+ Self :: hotkey_account_exists( & hotkey) ,
874+ Error :: <T >:: HotKeyAccountNotExists
875+ ) ;
876+ ensure ! (
877+ Self :: coldkey_owns_hotkey( & current_coldkey, & hotkey) ,
878+ Error :: <T >:: NonAssociatedColdKey
879+ ) ;
880+
881+ // Ensure the new coldkey is different from the current one
882+ ensure ! ( current_coldkey != new_coldkey, Error :: <T >:: SameColdkey ) ;
883+
884+ // Get the current stake
885+ let current_stake: u64 = Self :: get_stake_for_coldkey_and_hotkey ( & current_coldkey, & hotkey) ;
886+
887+ // Unstake all balance if there's any stake
888+ if current_stake > 0 {
889+ Self :: do_remove_stake (
890+ RawOrigin :: Signed ( current_coldkey. clone ( ) ) . into ( ) ,
891+ hotkey. clone ( ) ,
892+ current_stake,
893+ ) ?;
894+ }
895+
896+ // Get the total balance of the current coldkey account
897+ // let total_balance: <<T as Config>::Currency as fungible::Inspect<<T as system::Config>::AccountId>>::Balance = T::Currency::total_balance(¤t_coldkey);
898+
899+ let total_balance = Self :: get_coldkey_balance ( & current_coldkey) ;
900+ log:: info!( "Total Bank Balance: {:?}" , total_balance) ;
901+
902+ // Ensure there's a balance to transfer
903+ ensure ! ( !total_balance. is_zero( ) , Error :: <T >:: NoBalanceToTransfer ) ;
904+
905+ // Attempt to transfer the entire total balance to the new coldkey
906+ T :: Currency :: transfer (
907+ & current_coldkey,
908+ & new_coldkey,
909+ total_balance,
910+ Preservation :: Expendable ,
911+ ) ?;
912+
913+ // Emit the event
914+ Self :: deposit_event ( Event :: AllBalanceUnstakedAndTransferredToNewColdkey {
915+ current_coldkey : current_coldkey. clone ( ) ,
916+ new_coldkey : new_coldkey. clone ( ) ,
917+ hotkey : hotkey. clone ( ) ,
918+ current_stake,
919+ total_balance,
920+ } ) ;
921+
922+ Ok ( ( ) )
923+ }
798924}
0 commit comments