Coldkey migration for root claim.

pallets/subtensor/src/macros/events.rs

@@ -487,5 +487,13 @@ mod events {
             /// The amount of alpha distributed
             alpha: AlphaCurrency,
         },
+
+        /// Root claim coldkey migration event.
+        RootClaimColdkeyMigrated {
+            /// Keys migrated this block
+            processed_keys: u64,
+            /// Added coldkeys
+            added_coldkeys: u64,
+        },
     }
 }

pallets/subtensor/src/macros/hooks.rs

@@ -166,7 +166,10 @@ mod hooks {
                 // Fix staking hot keys
                 .saturating_add(migrations::migrate_fix_staking_hot_keys::migrate_fix_staking_hot_keys::<T>())
                 // Migrate coldkey swap scheduled to announcements
-                .saturating_add(migrations::migrate_coldkey_swap_scheduled_to_announcements::migrate_coldkey_swap_scheduled_to_announcements::<T>());
+                .saturating_add(migrations::migrate_coldkey_swap_scheduled_to_announcements::migrate_coldkey_swap_scheduled_to_announcements::<T>())
+                // Reset coldkey migration
+                .saturating_add(migrations::migrate_reset_coldkey_for_root_claim::migrate_reset_coldkey_for_root_claim::<T>());
+
             weight
         }
 

pallets/subtensor/src/migrations/migrate_reset_coldkey_for_root_claim.rs

@@ -0,0 +1,42 @@
+use super::*;
+use crate::HasMigrationRun;
+use frame_support::{traits::Get, weights::Weight};
+use scale_info::prelude::string::String;
+
+pub fn migrate_reset_coldkey_for_root_claim<T: Config>() -> Weight {
+    let migration_name = b"migrate_reset_coldkey_for_root_claim".to_vec();
+    let mut weight = T::DbWeight::get().reads(1);
+
+    if HasMigrationRun::<T>::get(&migration_name) {
+        log::info!(
+            "Migration '{:?}' has already run. Skipping.",
+            String::from_utf8_lossy(&migration_name)
+        );
+        return weight;
+    }
+
+    log::info!(
+        "Running migration '{}'",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    // ------------------------------
+    // Step 1: Reset AlphaMapLastKey
+    // ------------------------------
+
+    AlphaMapLastKey::<T>::kill();
+
+    // ------------------------------
+    // Step 2: Mark Migration as Completed
+    // ------------------------------
+
+    HasMigrationRun::<T>::insert(&migration_name, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    log::info!(
+        "Migration '{:?}' completed successfully.",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    weight
+}

pallets/subtensor/src/migrations/mod.rs

@@ -43,6 +43,7 @@ pub mod migrate_remove_unknown_neuron_axon_cert_prom;
 pub mod migrate_remove_unused_maps_and_values;
 pub mod migrate_remove_zero_total_hotkey_alpha;
 pub mod migrate_reset_bonds_moving_average;
+pub mod migrate_reset_coldkey_for_root_claim;
 pub mod migrate_reset_max_burn;
 pub mod migrate_reset_unactive_sn;
 pub mod migrate_set_first_emission_block_number;

pallets/subtensor/src/staking/claim_root.rs

@@ -316,13 +316,17 @@ impl<T: Config> Pallet<T> {
             .saturating_add(Weight::from_parts(100_412, 0).saturating_mul(k.into()))
     }
 
-    pub fn maybe_add_coldkey_index(coldkey: &T::AccountId) {
+    pub fn maybe_add_coldkey_index(coldkey: &T::AccountId) -> bool {
         if !StakingColdkeys::<T>::contains_key(coldkey) {
             let n = NumStakingColdkeys::<T>::get();
             StakingColdkeysByIndex::<T>::insert(n, coldkey.clone());
             StakingColdkeys::<T>::insert(coldkey.clone(), n);
             NumStakingColdkeys::<T>::mutate(|n| *n = n.saturating_add(1));
+
+            return true;
         }
+
+        false
     }
 
     pub fn run_auto_claim_root_divs(last_block_hash: T::Hash) -> Weight {

pallets/subtensor/src/staking/helpers.rs

@@ -367,6 +367,8 @@ impl<T: Config> Pallet<T> {
     /// alphas. It keeps the alpha value stored when it's >= than MIN_ALPHA.
     /// The function uses AlphaMapLastKey as a storage for key iterator between runs.
     pub fn populate_root_coldkey_staking_maps() {
+        let mut processed_keys = 0u64;
+        let mut added_coldkeys = 0u64;
         // Get starting key for the batch. Get the first key if we restart the process.
         let mut new_starting_raw_key = AlphaMapLastKey::<T>::get();
         let mut starting_key = None;
@@ -381,6 +383,8 @@ impl<T: Config> Pallet<T> {
                 .take(ALPHA_MAP_BATCH_SIZE)
                 .collect::<Vec<_>>();
 
+            processed_keys = processed_keys.saturating_add(keys.len() as u64);
+
             // New iteration: insert the starting key in the batch if it's a new iteration
             // iter_keys_from() skips the starting key
             if let Some(starting_key) = starting_key {
@@ -397,8 +401,8 @@ impl<T: Config> Pallet<T> {
             for key in keys {
                 let (_, coldkey, netuid) = key.clone();
 
-                if netuid == NetUid::ROOT {
-                    Self::maybe_add_coldkey_index(&coldkey);
+                if netuid == NetUid::ROOT && Self::maybe_add_coldkey_index(&coldkey) {
+                    added_coldkeys = added_coldkeys.saturating_add(1);
                 }
 
                 new_starting_key = Some(Alpha::<T>::hashed_key_for(key));
@@ -411,6 +415,11 @@ impl<T: Config> Pallet<T> {
 
             AlphaMapLastKey::<T>::put(new_starting_key);
         }
+
+        Self::deposit_event(Event::<T>::RootClaimColdkeyMigrated {
+            processed_keys,
+            added_coldkeys,
+        })
     }
 
     pub fn burn_subnet_alpha(_netuid: NetUid, _amount: AlphaCurrency) {

pallets/subtensor/src/tests/migration.rs

@@ -3039,3 +3039,70 @@ fn test_migrate_coldkey_swap_scheduled_to_announcements() {
         );
     });
 }
+
+#[test]
+fn test_migrate_reset_coldkey_for_root_claim() {
+    new_test_ext(1).execute_with(|| {
+        const MIGRATION_NAME: &[u8] = b"migrate_reset_coldkey_for_root_claim";
+
+        // Step 1: Set up initial state - put a value in AlphaMapLastKey
+        let test_key: Option<Vec<u8>> = Some(vec![1, 2, 3, 4, 5]);
+        AlphaMapLastKey::<Test>::put(test_key.clone());
+
+        // Verify the value was set
+        assert_eq!(
+            AlphaMapLastKey::<Test>::get(),
+            test_key,
+            "AlphaMapLastKey should have the test value."
+        );
+
+        // Step 2: Verify migration hasn't run yet
+        assert!(
+            !HasMigrationRun::<Test>::get(MIGRATION_NAME.to_vec()),
+            "Migration should not have run yet."
+        );
+
+        // Step 3: Run the migration
+        let weight = crate::migrations::migrate_reset_coldkey_for_root_claim::migrate_reset_coldkey_for_root_claim::<Test>();
+
+        // Step 4: Verify AlphaMapLastKey was killed (reset to default None)
+        assert_eq!(
+            AlphaMapLastKey::<Test>::get(),
+            None,
+            "AlphaMapLastKey should be reset to None after migration."
+        );
+
+        // Step 5: Verify migration is marked as completed
+        assert!(
+            HasMigrationRun::<Test>::get(MIGRATION_NAME.to_vec()),
+            "Migration should be marked as run."
+        );
+
+        // Step 6: Verify weight is non-zero
+        assert!(
+            !weight.is_zero(),
+            "Migration weight should be non-zero."
+        );
+
+        // Step 7: Run migration again - should be idempotent
+        // Set the value again to verify migration doesn't run twice
+        let test_key_2: Option<Vec<u8>> = Some(vec![6, 7, 8, 9, 10]);
+        AlphaMapLastKey::<Test>::put(test_key_2.clone());
+
+        let weight_second_run = crate::migrations::migrate_reset_coldkey_for_root_claim::migrate_reset_coldkey_for_root_claim::<Test>();
+
+        // Value should NOT be reset since migration already ran
+        assert_eq!(
+            AlphaMapLastKey::<Test>::get(),
+            test_key_2,
+            "AlphaMapLastKey should not be changed on second migration run."
+        );
+
+        // Second run should only have read weight (checking if migration ran)
+        assert_eq!(
+            weight_second_run,
+            <Test as frame_system::Config>::DbWeight::get().reads(1),
+            "Second migration run should only have read weight."
+        );
+    });
+}