Cleanup coldkey stake swap

pallets/subtensor/src/swap/swap_coldkey.rs

@@ -1,5 +1,4 @@
 use super::*;
-use share_pool::SafeFloat;
 
 impl<T: Config> Pallet<T> {
     /// Transfer all assets, stakes, subnet ownerships, and hotkey associations from `old_coldkey` to
@@ -98,38 +97,23 @@ impl<T: Config> Pallet<T> {
         new_coldkey: &T::AccountId,
     ) {
         for hotkey in StakingHotkeys::<T>::get(old_coldkey) {
-            // Swap and lazy-migrate Alpha to AlphaV2
-            // TotalHotkeyShares does not have to be migrated here, these migrations can be independent
-
-            // Get the v1 alpha shares on the old (hot,coldkey) account.
-            let orig_alpha_v1: SafeFloat =
-                SafeFloat::from(Alpha::<T>::get((&hotkey, old_coldkey, netuid)));
-            // Get the v1 alpha shares on the new (hot,coldkey) account.
-            let dest_alpha_v1: SafeFloat =
-                SafeFloat::from(Alpha::<T>::get((&hotkey, new_coldkey, netuid)));
-            // Get the v2 alpha shares on the old (hot,coldkey) account.
-            let orig_alpha_v2: SafeFloat = AlphaV2::<T>::get((&hotkey, old_coldkey, netuid));
-            // Get the v2 alpha shares on the new (hot,coldkey) account.
-            let dest_alpha_v2: SafeFloat = AlphaV2::<T>::get((&hotkey, new_coldkey, netuid));
-
-            // Calculate and save new alpha shares on the destination new_coldkey
-            let new_dest_alpha = orig_alpha_v1
-                .add(&dest_alpha_v1)
-                .unwrap_or_default()
-                .add(&orig_alpha_v2)
-                .unwrap_or_default()
-                .add(&dest_alpha_v2)
-                .unwrap_or_default();
-            if !new_dest_alpha.is_zero() {
-                AlphaV2::<T>::insert((&hotkey, new_coldkey, netuid), new_dest_alpha.clone());
-            }
-
-            // Remove shares on the origin old_coldkey in both Alpha and AlphaV2 maps
-            Alpha::<T>::remove((&hotkey, old_coldkey, netuid));
-            AlphaV2::<T>::remove((&hotkey, old_coldkey, netuid));
-
-            // Remove shares on the destination new_coldkey in Alpha map
-            Alpha::<T>::remove((&hotkey, new_coldkey, netuid));
+            // Swap
+            let alpha_old =
+                Self::get_stake_for_hotkey_and_coldkey_on_subnet(&hotkey, old_coldkey, netuid);
+            Self::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+                &hotkey,
+                old_coldkey,
+                netuid,
+                alpha_old,
+            );
+            Self::increase_stake_for_hotkey_and_coldkey_on_subnet(
+                &hotkey,
+                new_coldkey,
+                netuid,
+                alpha_old,
+            );
+            let new_dest_alpha =
+                Self::get_stake_for_hotkey_and_coldkey_on_subnet(&hotkey, new_coldkey, netuid);
 
             if !new_dest_alpha.is_zero() {
                 Self::transfer_root_claimed_for_new_keys(

pallets/subtensor/src/tests/swap_coldkey.rs

@@ -496,6 +496,7 @@ fn test_swap_coldkey_works() {
     });
 }
 
+// cargo test --package pallet-subtensor --lib -- tests::swap_coldkey::test_swap_coldkey_works_with_zero_cost --exact --nocapture
 #[test]
 fn test_swap_coldkey_works_with_zero_cost() {
     new_test_ext(1).execute_with(|| {
@@ -1763,9 +1764,13 @@ macro_rules! comprehensive_checks {
             $total_ck_stake,
         );
 
-        // Ensure the staking hotkeys are correctly swapped
+        // Ensure the staking hotkeys are correctly swapped (order-incensitive)
         assert!(StakingHotkeys::<Test>::get($who).is_empty());
-        assert_eq!(StakingHotkeys::<Test>::get($new_coldkey), $hotkeys);
+        let mut actual_st_hots = StakingHotkeys::<Test>::get($new_coldkey);
+        let mut expected_st_hots = $hotkeys.clone();
+        actual_st_hots.sort();
+        expected_st_hots.sort();
+        assert_eq!(actual_st_hots, expected_st_hots);
 
         // Ensure the hotkey ownership is correctly swapped
         assert!(OwnedHotkeys::<Test>::get($who).is_empty());