Fixing StakingHotkeys index

pallets/subtensor/src/macros/hooks.rs

@@ -162,7 +162,9 @@ mod hooks {
                 // Remove old identity map entries(Identities, SubnetIdentities, SubnetIdentitiesV2)
                 .saturating_add(migrations::migrate_remove_old_identity_maps::migrate_remove_old_identity_maps::<T>())
                 // Remove unknown neuron axon, certificate prom
-                .saturating_add(migrations::migrate_remove_unknown_neuron_axon_cert_prom::migrate_remove_unknown_neuron_axon_cert_prom::<T>());
+                .saturating_add(migrations::migrate_remove_unknown_neuron_axon_cert_prom::migrate_remove_unknown_neuron_axon_cert_prom::<T>())
+                // Fix staking hot keys
+                .saturating_add(migrations::migrate_fix_staking_hot_keys::migrate_fix_staking_hot_keys::<T>());
             weight
         }
 

pallets/subtensor/src/migrations/migrate_fix_staking_hot_keys.rs

@@ -0,0 +1,56 @@
+use super::*;
+use frame_support::{traits::Get, weights::Weight};
+use log;
+use scale_info::prelude::string::String;
+use sp_std::collections::btree_map::BTreeMap;
+
+pub fn migrate_fix_staking_hot_keys<T: Config>() -> Weight {
+    let migration_name = b"migrate_fix_staking_hot_keys".to_vec();
+    let mut weight = T::DbWeight::get().reads(1);
+
+    // Skip if already executed
+    if HasMigrationRun::<T>::get(&migration_name) {
+        log::info!(
+            target: "runtime",
+            "Migration '{}' already run - skipping.",
+            String::from_utf8_lossy(&migration_name)
+        );
+        return weight;
+    }
+
+    let mut cache: BTreeMap<T::AccountId, Vec<T::AccountId>> = BTreeMap::new();
+    let mut storage_reads: u64 = 0;
+    let mut storage_writes: u64 = 0;
+
+    for ((hotkey, coldkey, _netuid), alpha) in Alpha::<T>::iter() {
+        storage_reads = storage_reads.saturating_add(1);
+
+        if alpha == 0 {
+            continue;
+        }
+
+        let staking_hotkeys = cache.entry(coldkey.clone()).or_insert_with(|| {
+            storage_reads = storage_reads.saturating_add(1);
+            StakingHotkeys::<T>::get(&coldkey)
+        });
+
+        if !staking_hotkeys.contains(&hotkey) {
+            staking_hotkeys.push(hotkey.clone());
+            storage_writes = storage_writes.saturating_add(1);
+            StakingHotkeys::<T>::insert(&coldkey, staking_hotkeys.clone());
+        }
+    }
+    weight = weight.saturating_add(T::DbWeight::get().reads_writes(storage_reads, storage_writes));
+
+    // Mark migration done
+    HasMigrationRun::<T>::insert(&migration_name, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    log::info!(
+        target: "runtime",
+        "Migration '{}' completed.",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    weight
+}

pallets/subtensor/src/migrations/mod.rs

@@ -19,6 +19,7 @@ pub mod migrate_fix_childkeys;
 pub mod migrate_fix_is_network_member;
 pub mod migrate_fix_root_subnet_tao;
 pub mod migrate_fix_root_tao_and_alpha_in;
+pub mod migrate_fix_staking_hot_keys;
 pub mod migrate_init_tao_flow;
 pub mod migrate_init_total_issuance;
 pub mod migrate_kappa_map_to_default;

pallets/subtensor/src/tests/migration.rs

@@ -1138,6 +1138,41 @@ fn test_migrate_rate_limit_keys() {
     });
 }
 
+#[test]
+fn test_migrate_fix_staking_hot_keys() {
+    new_test_ext(1).execute_with(|| {
+        const MIGRATION_NAME: &[u8] = b"migrate_fix_staking_hot_keys";
+
+        assert!(
+            !HasMigrationRun::<Test>::get(MIGRATION_NAME.to_vec()),
+            "Migration should not have run yet"
+        );
+
+        // Add some data
+        Alpha::<Test>::insert(
+            (U256::from(1), U256::from(2), NetUid::ROOT),
+            U64F64::from(1_u64),
+        );
+        // Run migration
+        let weight =
+            migrations::migrate_fix_staking_hot_keys::migrate_fix_staking_hot_keys::<Test>();
+
+        assert!(
+            HasMigrationRun::<Test>::get(MIGRATION_NAME.to_vec()),
+            "Migration should be marked as completed"
+        );
+
+        // Check migration has been marked as run
+        assert!(HasMigrationRun::<Test>::get(MIGRATION_NAME.to_vec()));
+
+        // Verify results
+        assert_eq!(
+            StakingHotkeys::<Test>::get(U256::from(2)),
+            vec![U256::from(1)]
+        );
+    });
+}
+
 #[test]
 fn test_migrate_fix_root_subnet_tao() {
     new_test_ext(1).execute_with(|| {