Cleanup swap and unstake all code around StakingHotkeys

pallets/subtensor/src/staking.rs

@@ -890,34 +890,6 @@ impl<T: Config> Pallet<T> {
         // Ensure the new coldkey is different from the current one
         ensure!(current_coldkey != new_coldkey, Error::<T>::SameColdkey);
 
-        // Get all the hotkeys associated with this coldkey
-        let hotkeys: Vec<T::AccountId> = OwnedHotkeys::<T>::get(&current_coldkey);
-
-        // iterate over all hotkeys.
-        for next_hotkey in hotkeys {
-            ensure!(
-                Self::hotkey_account_exists(&next_hotkey),
-                Error::<T>::HotKeyAccountNotExists
-            );
-            ensure!(
-                Self::coldkey_owns_hotkey(&current_coldkey, &next_hotkey),
-                Error::<T>::NonAssociatedColdKey
-            );
-
-            // Get the current stake
-            let current_stake: u64 =
-                Self::get_stake_for_coldkey_and_hotkey(&current_coldkey, &next_hotkey);
-
-            // Unstake all balance if there's any stake
-            if current_stake > 0 {
-                Self::do_remove_stake(
-                    RawOrigin::Signed(current_coldkey.clone()).into(),
-                    next_hotkey.clone(),
-                    current_stake,
-                )?;
-            }
-        }
-
         // Unstake all delegate stake make by this coldkey to non-owned hotkeys
         let staking_hotkeys = StakingHotkeys::<T>::get(&current_coldkey);
 

pallets/subtensor/src/swap.rs

@@ -139,6 +139,7 @@ impl<T: Config> Pallet<T> {
         );
         Self::swap_subnet_owner_for_coldkey(old_coldkey, new_coldkey, &mut weight);
         Self::swap_owned_for_coldkey(old_coldkey, new_coldkey, &mut weight);
+        Self::swap_staking_hotkeys_for_coldkey(old_coldkey, new_coldkey, &mut weight);
 
         // Transfer any remaining balance from old_coldkey to new_coldkey
         let remaining_balance = Self::get_coldkey_balance(old_coldkey);
@@ -531,13 +532,36 @@ impl<T: Config> Pallet<T> {
             let stake = Stake::<T>::get(&hotkey, old_coldkey);
             Stake::<T>::remove(&hotkey, old_coldkey);
             Stake::<T>::insert(&hotkey, new_coldkey, stake);
+            weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
+        }
+    }
 
-            // Update StakingHotkeys map
-            let staking_hotkeys = StakingHotkeys::<T>::get(old_coldkey);
-            StakingHotkeys::<T>::insert(new_coldkey.clone(), staking_hotkeys);
+    /// Swaps all staking hotkeys associated with a coldkey from the old coldkey to the new coldkey.
+    ///
+    /// # Arguments
+    ///
+    /// * `old_coldkey` - The AccountId of the old coldkey.
+    /// * `new_coldkey` - The AccountId of the new coldkey.
+    /// * `weight` - Mutable reference to the weight of the transaction.
+    ///
+    /// # Effects
+    ///
+    /// * Removes all stakes associated with the old coldkey.
+    /// * Inserts all stakes for the new coldkey.
+    /// * Updates the transaction weight.
+    pub fn swap_staking_hotkeys_for_coldkey(
+        old_coldkey: &T::AccountId,
+        new_coldkey: &T::AccountId,
+        weight: &mut Weight,
+    ) {
+        // Find all hotkeys for this coldkey
+        let hotkeys = OwnedHotkeys::<T>::get(old_coldkey);
 
-            weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 3));
-        }
+        // Update StakingHotkeys map
+        OwnedHotkeys::<T>::remove(old_coldkey);
+        StakingHotkeys::<T>::insert(new_coldkey.clone(), hotkeys);
+
+        weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
     }
 
     /// Swaps the owner of all hotkeys from the old coldkey to the new coldkey.

pallets/subtensor/tests/swap.rs

@@ -1213,7 +1213,7 @@ fn test_swap_stake_for_coldkey() {
         assert_eq!(TotalIssuance::<Test>::get(), stake_amount1 + stake_amount2);
 
         // Verify weight update
-        let expected_weight = <Test as frame_system::Config>::DbWeight::get().reads_writes(5, 6);
+        let expected_weight = <Test as frame_system::Config>::DbWeight::get().reads_writes(3, 4);
         assert_eq!(weight, expected_weight);
     });
 }