add migration and tests for total coldkey hotfix

Author: const

pallets/subtensor/src/lib.rs

@@ -1419,7 +1419,9 @@ pub mod pallet {
                 // Populate OwnedHotkeys map for coldkey swap. Doesn't update storage vesion.
                 .saturating_add(migration::migrate_populate_owned::<T>())
                 // Populate StakingHotkeys map for coldkey swap. Doesn't update storage vesion.
-                .saturating_add(migration::migrate_populate_staking_hotkeys::<T>());
+                .saturating_add(migration::migrate_populate_staking_hotkeys::<T>())
+                // Fix total coldkey stake.
+                .saturating_add(migration::migrate_fix_total_coldkey_stake::<T>());
 
             weight
         }
@@ -2258,12 +2260,10 @@ pub mod pallet {
         #[pallet::call_index(64)]
         #[pallet::weight((0, DispatchClass::Operational, Pays::No))]
         pub fn sudo_hotfix_swap_coldkey_delegates(
-            origin: OriginFor<T>,
-            old_coldkey: T::AccountId,
-            new_coldkey: T::AccountId,
+            _origin: OriginFor<T>,
+            _old_coldkey: T::AccountId,
+            _new_coldkey: T::AccountId,
         ) -> DispatchResult {
-            ensure_root(origin)?;
-            Self::swap_hotfix(&old_coldkey, &new_coldkey);
             Ok(())
         }
     }

pallets/subtensor/src/migration.rs

@@ -24,6 +24,63 @@ pub mod deprecated_loaded_emission_format {
         StorageMap<Pallet<T>, Identity, u16, Vec<(AccountIdOf<T>, u64)>, OptionQuery>;
 }
 
+
+/// Migrates and fixes the total coldkey stake.
+///
+/// This function iterates through all staking hotkeys, calculates the total stake for each coldkey,
+/// and updates the `TotalColdkeyStake` storage accordingly. The migration is only performed if the
+/// on-chain storage version is 6.
+///
+/// # Returns
+/// The weight of the migration process.
+pub fn do_migrate_fix_total_coldkey_stake<T: Config>() -> Weight{
+    // Initialize the weight with one read operation.
+    let mut weight = T::DbWeight::get().reads(1);
+
+    // Iterate through all staking hotkeys.
+    for (coldkey, hotkey_vec) in StakingHotkeys::<T>::iter() {
+        // Init the zero value.
+        let mut coldkey_stake_sum: u64 = 0;
+        weight = weight.saturating_add(T::DbWeight::get().reads(1));
+
+        // Calculate the total stake for the current coldkey.
+        for hotkey in hotkey_vec {
+            // Cant fail on retrieval.
+            coldkey_stake_sum = coldkey_stake_sum.saturating_add(Stake::<T>::get(hotkey, coldkey.clone()));
+            weight = weight.saturating_add(T::DbWeight::get().reads(1));
+        }
+        // Update the `TotalColdkeyStake` storage with the calculated stake sum.
+        // Cant fail on insert.
+        TotalColdkeyStake::<T>::insert( coldkey.clone(), coldkey_stake_sum );
+        weight = weight.saturating_add(T::DbWeight::get().writes(1));
+    }
+    weight
+}
+// Public migrate function to be called by Lib.rs on upgrade.
+pub fn migrate_fix_total_coldkey_stake<T: Config>() -> Weight {
+
+    let current_storage_version: u16 = 6;
+    let next_storage_version: u16 = 7;
+
+    // Initialize the weight with one read operation.
+    let mut weight = T::DbWeight::get().reads(1);
+
+    // Grab the current on-chain storage version.
+    // Cant fail on retrieval.
+    let onchain_version = Pallet::<T>::on_chain_storage_version(); 
+
+    // Only run this migration on storage version 6.
+    if onchain_version == current_storage_version {
+        weight = weight.saturating_add( do_migrate_fix_total_coldkey_stake::<T>() );
+        // Cant fail on insert.
+        StorageVersion::new( next_storage_version ).put::<Pallet<T>>();
+        weight.saturating_accrue(T::DbWeight::get().writes(1));
+    } 
+
+    // Return the migration weight.
+    weight
+}
+
 /// Performs migration to update the total issuance based on the sum of stakes and total balances.
 /// This migration is applicable only if the current storage version is 5, after which it updates the storage version to 6.
 ///

pallets/subtensor/src/swap.rs

@@ -972,35 +972,6 @@ impl<T: Config> Pallet<T> {
         );
     }
 
-    pub fn swap_hotfix(old_coldkey: &T::AccountId, new_coldkey: &T::AccountId) {
-        let weight = T::DbWeight::get().reads_writes(2, 1);
-        let staking_hotkeys = StakingHotkeys::<T>::get(old_coldkey);
-        for staking_hotkey in staking_hotkeys {
-            if Stake::<T>::contains_key(staking_hotkey.clone(), old_coldkey) {
-                let hotkey = &staking_hotkey;
-                // Retrieve and remove the stake associated with the hotkey and old coldkey
-                let stake: u64 = Stake::<T>::get(hotkey, old_coldkey);
-                Stake::<T>::remove(hotkey, old_coldkey);
-                if stake > 0 {
-                    // Insert the stake for the hotkey and new coldkey
-                    let old_stake = Stake::<T>::get(hotkey, new_coldkey);
-                    Stake::<T>::insert(hotkey, new_coldkey, stake.saturating_add(old_stake));
-                }
-            }
-        }
-
-        let mut existing_staking_hotkeys = StakingHotkeys::<T>::get(new_coldkey);
-
-        let staking_hotkeys = StakingHotkeys::<T>::get(old_coldkey);
-        for hotkey in staking_hotkeys {
-            if !existing_staking_hotkeys.contains(&hotkey) {
-                existing_staking_hotkeys.push(hotkey);
-            }
-        }
-        StakingHotkeys::<T>::insert(new_coldkey, existing_staking_hotkeys);
-        StakingHotkeys::<T>::remove(old_coldkey);
-    }
-
     /// Swaps the total hotkey-coldkey stakes for the current interval from the old coldkey to the new coldkey.
     ///
     /// # Arguments

pallets/subtensor/tests/migration.rs

@@ -3,6 +3,7 @@
 mod mock;
 use frame_support::assert_ok;
 use frame_system::Config;
+use pallet_subtensor::*;
 use mock::*;
 use sp_core::U256;
 
@@ -276,3 +277,62 @@ fn test_migration_delete_subnet_21() {
         assert!(!SubtensorModule::if_subnet_exist(21));
     })
 }
+
+
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test migration -- test_migrate_fix_total_coldkey_stake --exact --nocapture
+#[test]
+fn test_migrate_fix_total_coldkey_stake() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(0); 
+        TotalColdkeyStake::<Test>::insert(coldkey, 0);
+        StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
+        Stake::<Test>::insert(U256::from(1), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(2), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(3), U256::from(0), 10000);
+        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+        assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 30000);
+    })
+}
+
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test migration -- test_migrate_fix_total_coldkey_stake_value_already_in_total --exact --nocapture
+#[test]
+fn test_migrate_fix_total_coldkey_stake_value_already_in_total() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(0); 
+        TotalColdkeyStake::<Test>::insert(coldkey, 100000000);
+        StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
+        Stake::<Test>::insert(U256::from(1), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(2), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(3), U256::from(0), 10000);
+        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+        assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 30000);
+    })
+}
+
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test migration -- test_migrate_fix_total_coldkey_stake_no_entry --exact --nocapture
+#[test]
+fn test_migrate_fix_total_coldkey_stake_no_entry() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(0); 
+        StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
+        Stake::<Test>::insert(U256::from(1), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(2), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(3), U256::from(0), 10000);
+        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+        assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 30000);
+    })
+}
+
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test migration -- test_migrate_fix_total_coldkey_stake_no_entry_in_hotkeys --exact --nocapture
+#[test]
+fn test_migrate_fix_total_coldkey_stake_no_entry_in_hotkeys() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(0); 
+        TotalColdkeyStake::<Test>::insert(coldkey, 100000000);
+        StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
+        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+        assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 0);
+    })
+}
+
+

pallets/subtensor/tests/swap.rs

@@ -1886,94 +1886,4 @@ fn test_coldkey_delegations() {
         assert_eq!(Stake::<Test>::get(delegate, new_coldkey), 100);
         assert_eq!(Stake::<Test>::get(delegate, coldkey), 0);
     });
-}
-
-// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test swap -- test_sudo_hotfix_swap_coldkey_delegates --exact --nocapture
-// #[test]
-// fn test_sudo_hotfix_swap_coldkey_delegates() {
-//     new_test_ext(1).execute_with(|| {
-//         let new_coldkey = U256::from(0);
-//         let owner = U256::from(1);
-//         let coldkey = U256::from(4);
-//         let delegate = U256::from(2);
-//         let netuid = 1u16;
-//         add_network(netuid, 13, 0);
-//         register_ok_neuron(netuid, delegate, owner, 0);
-//         SubtensorModule::add_balance_to_coldkey_account(&coldkey, 1000);
-//         assert_ok!(SubtensorModule::do_become_delegate(
-//             <<Test as Config>::RuntimeOrigin>::signed(owner),
-//             delegate,
-//             u16::MAX / 10
-//         ));
-//         assert_ok!(SubtensorModule::add_stake(
-//             <<Test as Config>::RuntimeOrigin>::signed(coldkey),
-//             delegate,
-//             100
-//         ));
-
-//         assert_ok!(SubtensorModule::perform_swap_coldkey(
-//             &coldkey,
-//             &new_coldkey
-//         ));
-
-// 		assert_ok!(AdminUtils::sudo_hotfix_swap_coldkey_delegates(
-//             <<Test as Config>::RuntimeOrigin>::root(),
-//             to_be_set
-//         ));
-
-//         assert_eq!( SubtensorModule::get_total_stake_for_hotkey( &delegate), 100 );
-//         assert_eq!( SubtensorModule::get_total_stake_for_coldkey( &coldkey), 0 );
-//         assert_eq!( SubtensorModule::get_total_stake_for_coldkey( &new_coldkey), 100 );
-//         assert_eq!( Stake::<Test>::get( delegate, new_coldkey ), 100 );
-//         assert_eq!( Stake::<Test>::get( delegate, coldkey ), 0 );
-
-//     });
-// }
-
-// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test swap -- test_sudo_hotfix_swap_coldkey_delegates --exact --nocapture
-#[test]
-fn test_sudo_hotfix_swap_coldkey_delegates() {
-    new_test_ext(1).execute_with(|| {
-        let new_coldkey = U256::from(0);
-        let coldkey = U256::from(4);
-        assert_ok!(SubtensorModule::sudo_hotfix_swap_coldkey_delegates(
-            <<Test as Config>::RuntimeOrigin>::root(),
-            coldkey,
-            new_coldkey
-        ));
-    });
-}
-
-// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test swap -- test_sudo_hotfix_swap_coldkey_delegates_with_broken_stake --exact --nocapture
-#[test]
-fn test_sudo_hotfix_swap_coldkey_delegates_with_broken_stake() {
-    new_test_ext(1).execute_with(|| {
-        let new_coldkey = U256::from(0);
-        let old_coldkey = U256::from(4);
-        let h1 = U256::from(5);
-        let h2 = U256::from(6);
-        let h3 = U256::from(7);
-        Stake::<Test>::insert(h3, old_coldkey, 100);
-        Stake::<Test>::insert(h2, old_coldkey, 100);
-        assert_eq!(Stake::<Test>::get(h3, old_coldkey), 100);
-        assert_eq!(Stake::<Test>::get(h2, old_coldkey), 100);
-        StakingHotkeys::<Test>::insert(new_coldkey, vec![h1, h2]);
-        StakingHotkeys::<Test>::insert(old_coldkey, vec![h3, h2]);
-        assert_ok!(SubtensorModule::sudo_hotfix_swap_coldkey_delegates(
-            <<Test as Config>::RuntimeOrigin>::root(),
-            old_coldkey,
-            new_coldkey
-        ));
-        let hotkeys = StakingHotkeys::<Test>::get(new_coldkey);
-        assert_eq!(hotkeys.len(), 3);
-        assert_eq!(hotkeys[0], h1);
-        assert_eq!(hotkeys[1], h2);
-        assert_eq!(hotkeys[2], h3);
-        let hotkeys_old = StakingHotkeys::<Test>::get(old_coldkey);
-        assert_eq!(hotkeys_old.len(), 0);
-        assert_eq!(Stake::<Test>::get(h3, old_coldkey), 0);
-        assert_eq!(Stake::<Test>::get(h2, old_coldkey), 0);
-        assert_eq!(Stake::<Test>::get(h3, new_coldkey), 100);
-        assert_eq!(Stake::<Test>::get(h2, new_coldkey), 100);
-    });
-}
+}

runtime/src/lib.rs

@@ -139,7 +139,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
     //   `spec_version`, and `authoring_version` are the same between Wasm and native.
     // This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use
     //   the compatible custom types.
-    spec_version: 164,
+    spec_version: 165,
     impl_version: 1,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 1,