Merge pull request #1142 from opentensor/fix/fix-hotkey-swap-alpha

Fix hotkey swap alpha

Author: sam0x17

pallets/subtensor/src/swap/swap_hotkey.rs

@@ -303,6 +303,36 @@ impl<T: Config> Pallet<T> {
             }
         });
 
+        // 11. Swap Alpha
+        // Alpha( hotkey, coldkey, netuid ) -> alpha
+        let old_alpha_values: Vec<((T::AccountId, u16), U64F64)> =
+            Alpha::<T>::iter_prefix((old_hotkey,)).collect();
+        // Clear the entire old prefix here.
+        let _ = Alpha::<T>::clear_prefix((old_hotkey,), old_alpha_values.len() as u32, None);
+        weight.saturating_accrue(T::DbWeight::get().reads(old_alpha_values.len() as u64));
+        weight.saturating_accrue(T::DbWeight::get().writes(old_alpha_values.len() as u64));
+
+        // Insert the new alpha values.
+        for ((coldkey, netuid), alpha) in old_alpha_values {
+            let new_alpha = Alpha::<T>::get((new_hotkey, &coldkey, netuid));
+            Alpha::<T>::insert(
+                (new_hotkey, &coldkey, netuid),
+                new_alpha.saturating_add(alpha),
+            );
+            weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
+
+            // Swap StakingHotkeys.
+            // StakingHotkeys( coldkey ) --> Vec<hotkey> -- the hotkeys that the coldkey stakes.
+            let mut staking_hotkeys = StakingHotkeys::<T>::get(&coldkey);
+            weight.saturating_accrue(T::DbWeight::get().reads(1));
+            if staking_hotkeys.contains(&old_hotkey) {
+                staking_hotkeys.retain(|hk| *hk != *old_hotkey && *hk != *new_hotkey);
+                staking_hotkeys.push(new_hotkey.clone());
+                StakingHotkeys::<T>::insert(&coldkey, staking_hotkeys);
+                weight.saturating_accrue(T::DbWeight::get().writes(1));
+            }
+        }
+
         // 11. Swap Stake.
         // Stake( hotkey, coldkey ) -> stake -- the stake that the hotkey controls on behalf of the coldkey.
         let stakes: Vec<(T::AccountId, u64)> = Stake::<T>::iter_prefix(old_hotkey).collect();
@@ -322,30 +352,7 @@ impl<T: Config> Pallet<T> {
                 &coldkey,
                 new_stake_value.saturating_add(old_stake_amount),
             );
-            weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
-
-            // 3.1 Swap Alpha
-            for netuid in Self::get_all_subnet_netuids() {
-                // Get the stake on the old (hot,coldkey) account.
-                let old_alpha: U64F64 = Alpha::<T>::get((&old_hotkey, coldkey.clone(), netuid));
-                // Get the stake on the new (hot,coldkey) account.
-                let new_alpha: U64F64 = Alpha::<T>::get((&new_hotkey, coldkey.clone(), netuid));
-                // Add the stake to new account.
-                Alpha::<T>::insert(
-                    (&new_hotkey, coldkey.clone(), netuid),
-                    new_alpha.saturating_add(old_alpha),
-                );
-                // Remove the value from the old account.
-                Alpha::<T>::remove((&old_hotkey, coldkey.clone(), netuid));
-            }
-
-            // Swap StakingHotkeys.
-            // StakingHotkeys( coldkey ) --> Vec<hotkey> -- the hotkeys that the coldkey stakes.
-            let mut staking_hotkeys = StakingHotkeys::<T>::get(coldkey.clone());
-            staking_hotkeys.retain(|hk| *hk != *old_hotkey && *hk != *new_hotkey);
-            staking_hotkeys.push(new_hotkey.clone());
-            StakingHotkeys::<T>::insert(coldkey.clone(), staking_hotkeys);
-            weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
+            weight.saturating_accrue(T::DbWeight::get().writes(1));
         }
 
         // 12. Swap ChildKeys.

pallets/subtensor/src/tests/swap_hotkey.rs

@@ -613,56 +613,132 @@ fn test_swap_hotkey_with_no_stake() {
     });
 }
 
-// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_hotkey_with_multiple_coldkeys_and_subnets --exact --nocapture
+// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --package pallet-subtensor --lib -- tests::swap_hotkey::test_swap_hotkey_with_multiple_coldkeys_and_subnets --exact --show-output
 #[test]
 fn test_swap_hotkey_with_multiple_coldkeys_and_subnets() {
     new_test_ext(1).execute_with(|| {
-        assert!(false);
-
-        // let old_hotkey = U256::from(1);
-        // let new_hotkey = U256::from(2);
-        // let coldkey1 = U256::from(3);
-        // let coldkey2 = U256::from(4);
-        // let netuid1 = 0;
-        // let netuid2 = 1;
-        // let mut weight = Weight::zero();
-
-        // // Set up initial state
-        // add_network(netuid1, 0, 1);
-        // add_network(netuid2, 0, 1);
-        // Owner::<Test>::insert(old_hotkey, coldkey1);
-        // Stake::<Test>::insert(old_hotkey, coldkey1, 100);
-        // Stake::<Test>::insert(old_hotkey, coldkey2, 200);
-        // IsNetworkMember::<Test>::insert(old_hotkey, netuid1, true);
-        // IsNetworkMember::<Test>::insert(old_hotkey, netuid2, true);
-        // TotalHotkeyStake::<Test>::insert(old_hotkey, 300);
-
-        // assert_ok!(SubtensorModule::perform_hotkey_swap(
-        //     &old_hotkey,
-        //     &new_hotkey,
-        //     &coldkey1,
-        //     &mut weight
-        // ));
-
-        // // Check ownership transfer
-        // assert!(!Owner::<Test>::contains_key(old_hotkey));
-        // assert_eq!(Owner::<Test>::get(new_hotkey), coldkey1);
-
-        // // Check stake transfer
-        // assert_eq!(Stake::<Test>::get(new_hotkey, coldkey1), 100);
-        // assert_eq!(Stake::<Test>::get(new_hotkey, coldkey2), 200);
-        // assert!(!Stake::<Test>::contains_key(old_hotkey, coldkey1));
-        // assert!(!Stake::<Test>::contains_key(old_hotkey, coldkey2));
-
-        // // Check subnet membership transfer
-        // assert!(IsNetworkMember::<Test>::get(new_hotkey, netuid1));
-        // assert!(IsNetworkMember::<Test>::get(new_hotkey, netuid2));
-        // assert!(!IsNetworkMember::<Test>::get(old_hotkey, netuid1));
-        // assert!(!IsNetworkMember::<Test>::get(old_hotkey, netuid2));
-
-        // // Check total stake transfer
-        // assert_eq!(TotalHotkeyStake::<Test>::get(new_hotkey), 300);
-        // assert!(!TotalHotkeyStake::<Test>::contains_key(old_hotkey));
+        let old_hotkey = U256::from(1);
+        let new_hotkey = U256::from(2);
+        let coldkey1 = U256::from(3);
+        let coldkey2 = U256::from(4);
+        let netuid1 = 1;
+        let netuid2 = 2;
+        let stake = 1_000_000_u64;
+        let mut weight = Weight::zero();
+
+        // Set up initial state
+        add_network(netuid1, 0, 1);
+        add_network(netuid2, 0, 1);
+        register_ok_neuron(netuid1, old_hotkey, coldkey1, 1234);
+        register_ok_neuron(netuid2, old_hotkey, coldkey1, 1234);
+
+        // Add balance to both coldkeys
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey1, stake + 1_000);
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey2, stake + 1_000);
+
+        // Stake with coldkey1
+        assert_ok!(SubtensorModule::add_stake(
+            <<Test as Config>::RuntimeOrigin>::signed(coldkey1),
+            old_hotkey,
+            netuid1,
+            stake
+        ));
+
+        // Stake with coldkey2 also
+        assert_ok!(SubtensorModule::add_stake(
+            <<Test as Config>::RuntimeOrigin>::signed(coldkey2),
+            old_hotkey,
+            netuid2,
+            stake
+        ));
+
+        let ck1_stake = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+            &old_hotkey,
+            &coldkey1,
+            netuid1,
+        );
+        let ck2_stake = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+            &old_hotkey,
+            &coldkey2,
+            netuid2,
+        );
+        assert!(ck1_stake > 0);
+        assert!(ck2_stake > 0);
+        let total_hk_stake = SubtensorModule::get_total_stake_for_hotkey(&old_hotkey);
+        assert!(total_hk_stake > 0);
+
+        assert_ok!(SubtensorModule::perform_hotkey_swap(
+            &old_hotkey,
+            &new_hotkey,
+            &coldkey1,
+            &mut weight
+        ));
+
+        // Check ownership transfer
+        assert_eq!(
+            SubtensorModule::get_owning_coldkey_for_hotkey(&new_hotkey),
+            coldkey1
+        );
+        assert!(!SubtensorModule::get_owned_hotkeys(&coldkey2).contains(&new_hotkey));
+
+        // Check stake transfer
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &new_hotkey,
+                &coldkey1,
+                netuid1
+            ),
+            ck1_stake
+        );
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &new_hotkey,
+                &coldkey2,
+                netuid2
+            ),
+            ck2_stake
+        );
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &old_hotkey,
+                &coldkey1,
+                netuid1
+            ),
+            0
+        );
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &old_hotkey,
+                &coldkey2,
+                netuid2
+            ),
+            0
+        );
+
+        // Check subnet membership transfer
+        assert!(SubtensorModule::is_hotkey_registered_on_network(
+            netuid1,
+            &new_hotkey
+        ));
+        assert!(SubtensorModule::is_hotkey_registered_on_network(
+            netuid2,
+            &new_hotkey
+        ));
+        assert!(!SubtensorModule::is_hotkey_registered_on_network(
+            netuid1,
+            &old_hotkey
+        ));
+        assert!(!SubtensorModule::is_hotkey_registered_on_network(
+            netuid2,
+            &old_hotkey
+        ));
+
+        // Check total stake transfer
+        assert_eq!(
+            SubtensorModule::get_total_stake_for_hotkey(&new_hotkey),
+            total_hk_stake
+        );
+        assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&old_hotkey), 0);
     });
 }