Fix coldkey_swap and its tests

gztensor · gztensor · commit bda0c5a31b63 · 2024-07-03T19:16:47.000-04:00
diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs
@@ -1209,11 +1209,13 @@ pub mod pallet {
 
                     // Update Owned map
                     let mut hotkeys = Owned::<T>::get(&coldkey);
-                    hotkeys.push(hotkey.clone());
-                    Owned::<T>::insert(
-                        &coldkey,
-                        hotkeys,
-                    );
+                    if !hotkeys.contains(&hotkey) {
+                        hotkeys.push(hotkey.clone());
+                        Owned::<T>::insert(
+                            &coldkey,
+                            hotkeys,
+                        );
+                    }
         
                     TotalHotkeyStake::<T>::insert(hotkey.clone(), stake);
                     TotalColdkeyStake::<T>::insert(
diff --git a/pallets/subtensor/src/migration.rs b/pallets/subtensor/src/migration.rs
@@ -494,10 +494,12 @@ pub fn migrate_populate_owned<T: Config>() -> Weight {
         let mut longest_coldkey: Option<T::AccountId> = None;
         Owner::<T>::iter().for_each(|(hotkey, coldkey)| {
             let mut hotkeys = Owned::<T>::get(&coldkey);
-            hotkeys.push(hotkey);
-            if longest_hotkey_vector < hotkeys.len() {
-                longest_hotkey_vector = hotkeys.len();
-                longest_coldkey = Some(coldkey.clone());
+            if !hotkeys.contains(&hotkey) {
+                hotkeys.push(hotkey);
+                if longest_hotkey_vector < hotkeys.len() {
+                    longest_hotkey_vector = hotkeys.len();
+                    longest_coldkey = Some(coldkey.clone());
+                }
             }
 
             Owned::<T>::insert(
diff --git a/pallets/subtensor/src/staking.rs b/pallets/subtensor/src/staking.rs
@@ -563,11 +563,13 @@ impl<T: Config> Pallet<T> {
 
             // Update Owned map
             let mut hotkeys = Owned::<T>::get(&coldkey);
-            hotkeys.push(hotkey.clone());
-            Owned::<T>::insert(
-                &coldkey,
-                hotkeys,
-            );
+            if !hotkeys.contains(&hotkey) {
+                hotkeys.push(hotkey.clone());
+                Owned::<T>::insert(
+                    &coldkey,
+                    hotkeys,
+                );
+            }
         }
     }
 
@@ -789,6 +791,31 @@ impl<T: Config> Pallet<T> {
         Ok(credit)
     }
 
+    pub fn kill_coldkey_account(
+        coldkey: &T::AccountId,
+        amount: <<T as Config>::Currency as fungible::Inspect<<T as system::Config>::AccountId>>::Balance,
+    ) -> Result<u64, DispatchError> {
+        if amount == 0 {
+            return Ok(0);
+        }
+
+        let credit = T::Currency::withdraw(
+            coldkey,
+            amount,
+            Precision::Exact,
+            Preservation::Expendable,
+            Fortitude::Force,
+        )
+        .map_err(|_| Error::<T>::BalanceWithdrawalError)?
+        .peek();
+
+        if credit == 0 {
+            return Err(Error::<T>::ZeroBalanceAfterWithdrawn.into());
+        }
+
+        Ok(credit)
+    }
+
     pub fn unstake_all_coldkeys_from_hotkey_account(hotkey: &T::AccountId) {
         // Iterate through all coldkeys that have a stake on this hotkey account.
         for (delegate_coldkey_i, stake_i) in
diff --git a/pallets/subtensor/src/swap.rs b/pallets/subtensor/src/swap.rs
@@ -167,13 +167,13 @@ impl<T: Config> Pallet<T> {
             new_coldkey,
             &mut weight,
         );
-        Self::swap_keys_for_coldkey(old_coldkey, new_coldkey, &mut weight);
         Self::swap_subnet_owner_for_coldkey(old_coldkey, new_coldkey, &mut weight);
+        Self::swap_owned_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);
         if remaining_balance > 0 {
-            Self::remove_balance_from_coldkey_account(old_coldkey, remaining_balance)?;
+            Self::kill_coldkey_account(old_coldkey, remaining_balance)?;
             Self::add_balance_to_coldkey_account(new_coldkey, remaining_balance);
         }
 
@@ -226,7 +226,10 @@ impl<T: Config> Pallet<T> {
 
         // Update Owned map
         let mut hotkeys = Owned::<T>::get(&coldkey);
-        hotkeys.push(new_hotkey.clone());
+        if !hotkeys.contains(&new_hotkey) {
+            hotkeys.push(new_hotkey.clone());
+        }
+        hotkeys.retain(|hk| *hk != *old_hotkey);
         Owned::<T>::insert(
             &coldkey,
             hotkeys,
@@ -578,11 +581,6 @@ impl<T: Config> Pallet<T> {
             Owner::<T>::insert(&hotkey, new_coldkey);
             weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 1));
         }
-
-        // Update Owned map with new coldkey
-        Owned::<T>::remove(old_coldkey);
-        Owned::<T>::insert(new_coldkey, hotkeys);
-        weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 2));
     }
 
     /// Swaps the total hotkey-coldkey stakes for the current interval from the old coldkey to the new coldkey.
@@ -603,48 +601,17 @@ impl<T: Config> Pallet<T> {
         new_coldkey: &T::AccountId,
         weight: &mut Weight,
     ) {
-        for (hotkey, (stakes, block)) in
-            TotalHotkeyColdkeyStakesThisInterval::<T>::iter_prefix(old_coldkey)
-        {
-            TotalHotkeyColdkeyStakesThisInterval::<T>::remove(old_coldkey, &hotkey);
+        weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0));
+        for hotkey in Owned::<T>::get(old_coldkey).iter() {
+            let (stake, block) = TotalHotkeyColdkeyStakesThisInterval::<T>::get(&hotkey, old_coldkey);
+            TotalHotkeyColdkeyStakesThisInterval::<T>::remove(&hotkey, old_coldkey);
             TotalHotkeyColdkeyStakesThisInterval::<T>::insert(
-                new_coldkey,
                 &hotkey,
-                (stakes, block),
+                new_coldkey,
+                (stake, block),
             );
+            weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
         }
-        weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
-    }
-
-    /// Swaps all keys associated with a coldkey from the old coldkey to the new coldkey across all networks.
-    ///
-    /// # 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
-    ///
-    /// * Updates all keys associated with the old coldkey to be associated with the new coldkey.
-    /// * Updates the transaction weight.
-    pub fn swap_keys_for_coldkey(
-        old_coldkey: &T::AccountId,
-        new_coldkey: &T::AccountId,
-        weight: &mut Weight,
-    ) {
-        for netuid in 0..TotalNetworks::<T>::get() {
-            for (uid, hotkey) in Keys::<T>::iter_prefix(netuid) {
-                if Owner::<T>::get(&hotkey) == *old_coldkey {
-                    Keys::<T>::remove(netuid, uid);
-                    Keys::<T>::insert(netuid, uid, new_coldkey);
-                }
-            }
-        }
-        weight.saturating_accrue(T::DbWeight::get().reads_writes(
-            TotalNetworks::<T>::get() as u64,
-            TotalNetworks::<T>::get() as u64,
-        ));
     }
 
     /// Checks if a coldkey has any associated hotkeys.
@@ -677,7 +644,7 @@ impl<T: Config> Pallet<T> {
         new_coldkey: &T::AccountId,
         weight: &mut Weight,
     ) {
-        for netuid in 0..TotalNetworks::<T>::get() {
+        for netuid in 0..=TotalNetworks::<T>::get() {
             let subnet_owner = SubnetOwner::<T>::get(netuid);
             if subnet_owner == *old_coldkey {
                 SubnetOwner::<T>::insert(netuid, new_coldkey.clone());
@@ -687,6 +654,30 @@ impl<T: Config> Pallet<T> {
         weight.saturating_accrue(T::DbWeight::get().reads(TotalNetworks::<T>::get() as u64));
     }
 
+    /// Swaps the owned hotkeys for the 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
+    ///
+    /// * Updates the subnet owner to the new coldkey for all networks where the old coldkey was the owner.
+    /// * Updates the transaction weight.
+    pub fn swap_owned_for_coldkey(
+        old_coldkey: &T::AccountId,
+        new_coldkey: &T::AccountId,
+        weight: &mut Weight,
+    ) {
+        // Update Owned map with new coldkey
+        let hotkeys = Owned::<T>::get(&old_coldkey);
+        Owned::<T>::remove(old_coldkey);
+        Owned::<T>::insert(new_coldkey, hotkeys);
+        weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 2));
+    }
+    
     /// Returns the cost of swapping a coldkey.
     ///
     /// # Returns
diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs
@@ -1058,11 +1058,12 @@ fn test_do_swap_coldkey_success() {
         let netuid = 1u16;
         let stake_amount = 1000u64;
         let swap_cost = SubtensorModule::get_coldkey_swap_cost();
+        let free_balance = 12345;
 
         // Setup initial state
         add_network(netuid, 13, 0);
         register_ok_neuron(netuid, hotkey, old_coldkey, 0);
-        SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + swap_cost);
+        SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + swap_cost + free_balance);
 
         // Add stake to the neuron
         assert_ok!(SubtensorModule::add_stake(
@@ -1072,12 +1073,22 @@ fn test_do_swap_coldkey_success() {
         ));
 
         log::info!("TotalColdkeyStake::<Test>::get(old_coldkey): {:?}", TotalColdkeyStake::<Test>::get(old_coldkey));
-        log::info!("Stake::<Test>::get(old_coldkey, hotkey): {:?}", Stake::<Test>::get(old_coldkey, hotkey));
+        log::info!("Stake::<Test>::get(old_coldkey, hotkey): {:?}", Stake::<Test>::get(hotkey, old_coldkey));
 
         // Verify initial stake
         assert_eq!(TotalColdkeyStake::<Test>::get(old_coldkey), stake_amount);
-        // assert_eq!(Stake::<Test>::get(old_coldkey, hotkey), stake_amount);
+        assert_eq!(Stake::<Test>::get(hotkey, old_coldkey), stake_amount);
+
+
+        assert_eq!(Owned::<Test>::get(old_coldkey), vec![hotkey]);
+        assert!(!Owned::<Test>::contains_key(new_coldkey));
+
+
+
 
+        // Get coldkey free balance before swap
+        let balance = SubtensorModule::get_coldkey_balance(&old_coldkey);
+        assert_eq!(balance, free_balance + swap_cost);
 
         // Perform the swap
         assert_ok!(SubtensorModule::do_swap_coldkey(
@@ -1086,19 +1097,19 @@ fn test_do_swap_coldkey_success() {
             &new_coldkey
         ));
 
-               // Verify the swap
+        // Verify the swap
         assert_eq!(Owner::<Test>::get(hotkey), new_coldkey);
         assert_eq!(TotalColdkeyStake::<Test>::get(new_coldkey), stake_amount);
         assert!(!TotalColdkeyStake::<Test>::contains_key(old_coldkey));
-        assert_eq!(Stake::<Test>::get(new_coldkey, hotkey), stake_amount);
-        assert!(!Stake::<Test>::contains_key(old_coldkey, hotkey));
+        assert_eq!(Stake::<Test>::get(hotkey, new_coldkey), stake_amount);
+        assert!(!Stake::<Test>::contains_key(hotkey, old_coldkey));
+        assert_eq!(Owned::<Test>::get(new_coldkey), vec![hotkey]);
+        assert!(!Owned::<Test>::contains_key(old_coldkey));
 
         // Verify balance transfer
-        assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), stake_amount);
+        assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), balance - swap_cost);
         assert_eq!(SubtensorModule::get_coldkey_balance(&old_coldkey), 0);
 
-
-
         // Verify event emission
         System::assert_last_event(Event::ColdkeySwapped { 
             old_coldkey: old_coldkey,
@@ -1244,6 +1255,9 @@ fn test_swap_owner_for_coldkey() {
         Owner::<Test>::insert(hotkey1, old_coldkey);
         Owner::<Test>::insert(hotkey2, old_coldkey);
 
+        // Initialize Owned map 
+        Owned::<Test>::insert(old_coldkey, vec![hotkey1, hotkey2]);
+
         // Perform the swap
         SubtensorModule::swap_owner_for_coldkey(&old_coldkey, &new_coldkey, &mut weight);
 
@@ -1269,8 +1283,11 @@ fn test_swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey() {
         let mut weight = Weight::zero();
 
         // Initialize TotalHotkeyColdkeyStakesThisInterval for old_coldkey
-        TotalHotkeyColdkeyStakesThisInterval::<Test>::insert(old_coldkey, hotkey1, stake1);
-        TotalHotkeyColdkeyStakesThisInterval::<Test>::insert(old_coldkey, hotkey2, stake2);
+        TotalHotkeyColdkeyStakesThisInterval::<Test>::insert(hotkey1, old_coldkey, stake1);
+        TotalHotkeyColdkeyStakesThisInterval::<Test>::insert(hotkey2, old_coldkey, stake2);
+
+        // Populate Owned map
+        Owned::<Test>::insert(old_coldkey, vec![hotkey1, hotkey2]);
 
         // Perform the swap
         SubtensorModule::swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey(
@@ -1281,11 +1298,11 @@ fn test_swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey() {
 
         // Verify the swap
         assert_eq!(
-            TotalHotkeyColdkeyStakesThisInterval::<Test>::get(new_coldkey, hotkey1),
+            TotalHotkeyColdkeyStakesThisInterval::<Test>::get(hotkey1, new_coldkey),
             stake1
         );
         assert_eq!(
-            TotalHotkeyColdkeyStakesThisInterval::<Test>::get(new_coldkey, hotkey2),
+            TotalHotkeyColdkeyStakesThisInterval::<Test>::get(hotkey2, new_coldkey),
             stake2
         );
         assert!(!TotalHotkeyColdkeyStakesThisInterval::<Test>::contains_key(
@@ -1298,42 +1315,7 @@ fn test_swap_total_hotkey_coldkey_stakes_this_interval_for_coldkey() {
         ));
 
         // Verify weight update
-        let expected_weight = <Test as frame_system::Config>::DbWeight::get().reads_writes(1, 4);
-        assert_eq!(weight, expected_weight);
-    });
-}
-
-#[test]
-fn test_swap_keys_for_coldkey() {
-    new_test_ext(1).execute_with(|| {
-        let old_coldkey = U256::from(1);
-        let new_coldkey = U256::from(2);
-        let hotkey1 = U256::from(3);
-        let hotkey2 = U256::from(4);
-        let netuid1 = 1u16;
-        let netuid2 = 2u16;
-        let uid1 = 10u16;
-        let uid2 = 20u16;
-        let mut weight = Weight::zero();
-
-        // Initialize Keys and Owner for old_coldkey
-        Keys::<Test>::insert(netuid1, uid1, hotkey1);
-        Keys::<Test>::insert(netuid2, uid2, hotkey2);
-        Owner::<Test>::insert(hotkey1, old_coldkey);
-        Owner::<Test>::insert(hotkey2, old_coldkey);
-
-        // Set up TotalNetworks
-        TotalNetworks::<Test>::put(3);
-
-        // Perform the swap
-        SubtensorModule::swap_keys_for_coldkey(&old_coldkey, &new_coldkey, &mut weight);
-
-        // Verify the swap
-        assert_eq!(Keys::<Test>::get(netuid1, uid1), hotkey1);
-        assert_eq!(Keys::<Test>::get(netuid2, uid2), hotkey2);
-
-        // Verify weight update
-        let expected_weight = <Test as frame_system::Config>::DbWeight::get().reads_writes(3, 2);
+        let expected_weight = <Test as frame_system::Config>::DbWeight::get().reads_writes(5, 4);
         assert_eq!(weight, expected_weight);
     });
 }
@@ -1380,9 +1362,16 @@ fn test_do_swap_coldkey_with_subnet_ownership() {
         // Setup initial state
         add_network(netuid, 13, 0);
         register_ok_neuron(netuid, hotkey, old_coldkey, 0);
+
+        // Set TotalNetworks because swap relies on it
+        pallet_subtensor::TotalNetworks::<Test>::set(1);
+
         SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + swap_cost);
         SubnetOwner::<Test>::insert(netuid, old_coldkey);
 
+        // Populate Owned map
+        Owned::<Test>::insert(old_coldkey, vec![hotkey]);
+
         // Perform the swap
         assert_ok!(SubtensorModule::do_swap_coldkey(
             <<Test as Config>::RuntimeOrigin>::signed(old_coldkey),
@@ -1402,8 +1391,12 @@ fn test_do_swap_coldkey_tx_rate_limit() {
         let new_coldkey = U256::from(2);
         let swap_cost = SubtensorModule::get_coldkey_swap_cost();
 
+        // Set non-zero tx rate limit
+        SubtensorModule::set_tx_rate_limit(1);
+
         // Setup initial state
         SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, swap_cost * 2);
+        SubtensorModule::add_balance_to_coldkey_account(&new_coldkey, swap_cost * 2);
 
         // Perform first swap
         assert_ok!(SubtensorModule::do_swap_coldkey(
