Fix/add auto stake dest coldkeys map for hk swap

pallets/subtensor/src/lib.rs

@@ -1150,6 +1150,16 @@ pub mod pallet {
         T::AccountId,
         OptionQuery,
     >;
+    #[pallet::storage] // --- DMAP ( hot, netuid )--> Vec<cold> | Returns a list of coldkeys that are autostaking to a hotkey.
+    pub type AutoStakeDestinationColdkeys<T: Config> = StorageDoubleMap<
+        _,
+        Blake2_128Concat,
+        T::AccountId,
+        Identity,
+        NetUid,
+        Vec<T::AccountId>,
+        ValueQuery,
+    >;
 
     #[pallet::storage] // --- DMAP ( cold ) --> (block_expected, new_coldkey) | Maps coldkey to the block to swap at and new coldkey.
     pub type ColdkeySwapScheduled<T: Config> = StorageMap<

pallets/subtensor/src/macros/dispatches.rs

@@ -2313,9 +2313,20 @@ mod dispatches {
                     current_hotkey != hotkey,
                     Error::<T>::SameAutoStakeHotkeyAlreadySet
                 );
+
+                // Remove the coldkey from the old hotkey (if present)
+                AutoStakeDestinationColdkeys::<T>::mutate(current_hotkey.clone(), netuid, |v| {
+                    v.retain(|c| c != &coldkey);
+                });
             }
 
+            // Add the coldkey to the new hotkey (if not already present)
             AutoStakeDestination::<T>::insert(coldkey.clone(), netuid, hotkey.clone());
+            AutoStakeDestinationColdkeys::<T>::mutate(hotkey.clone(), netuid, |v| {
+                if !v.contains(&coldkey) {
+                    v.push(coldkey.clone());
+                }
+            });
 
             Self::deposit_event(Event::AutoStakeDestinationSet {
                 coldkey,

pallets/subtensor/src/macros/hooks.rs

@@ -149,7 +149,9 @@ mod hooks {
                 // Migrate subnet locked balances
                 .saturating_add(migrations::migrate_subnet_locked::migrate_restore_subnet_locked::<T>())
                 // Migrate subnet burn cost to 2500
-                .saturating_add(migrations::migrate_network_lock_cost_2500::migrate_network_lock_cost_2500::<T>());
+                .saturating_add(migrations::migrate_network_lock_cost_2500::migrate_network_lock_cost_2500::<T>())
+                // Migrate AutoStakeDestinationColdkeys
+                .saturating_add(migrations::migrate_auto_stake_destination::migrate_auto_stake_destination::<T>());
             weight
         }
 

pallets/subtensor/src/migrations/migrate_auto_stake_destination.rs

@@ -57,6 +57,11 @@ pub fn migrate_auto_stake_destination<T: Config>() -> Weight {
                     continue;
                 }
                 AutoStakeDestination::<T>::insert(coldkey, netuid, hotkey.clone());
+                AutoStakeDestinationColdkeys::<T>::mutate(hotkey.clone(), netuid, |v| {
+                    if !v.contains(coldkey) {
+                        v.push(coldkey.clone());
+                    }
+                });
             }
 
             old::AutoStakeDestination::<T>::remove(coldkey);

pallets/subtensor/src/swap/swap_coldkey.rs

@@ -159,7 +159,16 @@ impl<T: Config> Pallet<T> {
             if let Some(old_auto_stake_hotkey) = AutoStakeDestination::<T>::get(old_coldkey, netuid)
             {
                 AutoStakeDestination::<T>::remove(old_coldkey, netuid);
-                AutoStakeDestination::<T>::insert(new_coldkey, netuid, old_auto_stake_hotkey);
+                AutoStakeDestination::<T>::insert(
+                    new_coldkey,
+                    netuid,
+                    old_auto_stake_hotkey.clone(),
+                );
+                AutoStakeDestinationColdkeys::<T>::mutate(old_auto_stake_hotkey, netuid, |v| {
+                    // Remove old/new coldkeys (avoid duplicates), then add the new one.
+                    v.retain(|c| *c != *old_coldkey && *c != *new_coldkey);
+                    v.push(new_coldkey.clone());
+                });
             }
         }
 

pallets/subtensor/src/swap/swap_hotkey.rs

@@ -524,6 +524,18 @@ impl<T: Config> Pallet<T> {
             }
         }
 
+        // 6.4 Swap AutoStakeDestination
+        if let Ok(old_auto_stake_coldkeys) =
+            AutoStakeDestinationColdkeys::<T>::try_get(old_hotkey, netuid)
+        {
+            // Move the vector from old hotkey to new hotkey.
+            for coldkey in &old_auto_stake_coldkeys {
+                AutoStakeDestination::<T>::insert(coldkey, netuid, new_hotkey);
+            }
+            AutoStakeDestinationColdkeys::<T>::remove(old_hotkey, netuid);
+            AutoStakeDestinationColdkeys::<T>::insert(new_hotkey, netuid, old_auto_stake_coldkeys);
+        }
+
         // 7. Swap SubnetOwnerHotkey
         // SubnetOwnerHotkey( netuid ) --> hotkey -- the hotkey that is the owner of the subnet.
         if let Ok(old_subnet_owner_hotkey) = SubnetOwnerHotkey::<T>::try_get(netuid) {

pallets/subtensor/src/tests/auto_stake_hotkey.rs

@@ -110,3 +110,64 @@ fn test_set_coldkey_auto_stake_hotkey_same_hotkey_again() {
         );
     });
 }
+
+#[test]
+fn test_set_coldkey_auto_stake_hotkey_change_hotkey() {
+    new_test_ext(1).execute_with(|| {
+        let subnet_owner_ck = U256::from(0);
+        let subnet_owner_hk = U256::from(1);
+
+        let coldkey = U256::from(10);
+        let hotkey = U256::from(11);
+        let new_hotkey = U256::from(12);
+
+        Owner::<Test>::insert(hotkey, coldkey);
+        OwnedHotkeys::<Test>::insert(coldkey, vec![hotkey]);
+
+        let netuid = add_dynamic_network(&subnet_owner_hk, &subnet_owner_ck);
+        Uids::<Test>::insert(netuid, hotkey, 1);
+        Uids::<Test>::insert(netuid, new_hotkey, 2);
+
+        // First call should succeed
+        assert_ok!(SubtensorModule::set_coldkey_auto_stake_hotkey(
+            RuntimeOrigin::signed(coldkey),
+            netuid,
+            hotkey,
+        ));
+
+        // Check maps
+        assert_eq!(
+            AutoStakeDestination::<Test>::get(coldkey, netuid),
+            Some(hotkey)
+        );
+        assert_eq!(
+            AutoStakeDestinationColdkeys::<Test>::get(hotkey, netuid),
+            vec![coldkey]
+        );
+        assert_eq!(
+            AutoStakeDestinationColdkeys::<Test>::get(new_hotkey, netuid),
+            vec![]
+        );
+
+        // Second call with new hotkey should succeed
+        assert_ok!(SubtensorModule::set_coldkey_auto_stake_hotkey(
+            RuntimeOrigin::signed(coldkey),
+            netuid,
+            new_hotkey,
+        ));
+
+        // Check maps again
+        assert_eq!(
+            AutoStakeDestination::<Test>::get(coldkey, netuid),
+            Some(new_hotkey)
+        );
+        assert_eq!(
+            AutoStakeDestinationColdkeys::<Test>::get(hotkey, netuid),
+            vec![]
+        );
+        assert_eq!(
+            AutoStakeDestinationColdkeys::<Test>::get(new_hotkey, netuid),
+            vec![coldkey]
+        );
+    });
+}

pallets/subtensor/src/tests/migration.rs

@@ -1581,8 +1581,17 @@ fn test_migrate_auto_stake_destination() {
 					AutoStakeDestination::<Test>::get(coldkey2, *netuid),
 					Some(hotkey2)
 				);
-			}
 
+				// Verify entry for AutoStakeDestinationColdkeys
+				assert_eq!(
+					AutoStakeDestinationColdkeys::<Test>::get(hotkey1, *netuid),
+					vec![coldkey1]
+				);
+				assert_eq!(
+					AutoStakeDestinationColdkeys::<Test>::get(hotkey2, *netuid),
+					vec![coldkey2]
+				);
+			}
 		}
 
         // Verify old format entries are cleared

pallets/subtensor/src/tests/swap_coldkey.rs

@@ -2619,3 +2619,37 @@ fn test_coldkey_in_swap_schedule_prevents_critical_calls() {
         );
     });
 }
+
+#[test]
+fn test_swap_auto_stake_destination_coldkeys() {
+    new_test_ext(1).execute_with(|| {
+        let old_coldkey = U256::from(1);
+        let new_coldkey = U256::from(2);
+        let hotkey = U256::from(3);
+        let netuid = NetUid::from(1u16);
+        let coldkeys = vec![U256::from(4), U256::from(5), old_coldkey];
+
+        add_network(netuid, 1, 0);
+        AutoStakeDestinationColdkeys::<Test>::insert(hotkey, netuid, coldkeys.clone());
+        AutoStakeDestination::<Test>::insert(old_coldkey, netuid, hotkey);
+
+        let mut weight = Weight::zero();
+        assert_ok!(SubtensorModule::perform_swap_coldkey(
+            &old_coldkey,
+            &new_coldkey,
+            &mut weight
+        ));
+
+        let new_coldkeys = AutoStakeDestinationColdkeys::<Test>::get(hotkey, netuid);
+        assert!(new_coldkeys.contains(&new_coldkey));
+        assert!(!new_coldkeys.contains(&old_coldkey));
+        assert_eq!(
+            AutoStakeDestination::<Test>::try_get(old_coldkey, netuid),
+            Err(())
+        );
+        assert_eq!(
+            AutoStakeDestination::<Test>::try_get(new_coldkey, netuid),
+            Ok(hotkey)
+        );
+    });
+}

pallets/subtensor/src/tests/swap_hotkey.rs

@@ -1465,3 +1465,39 @@ fn test_swap_hotkey_swap_rate_limits() {
         );
     });
 }
+
+#[test]
+fn test_swap_auto_stake_destination_coldkeys() {
+    new_test_ext(1).execute_with(|| {
+        let old_hotkey = U256::from(1);
+        let new_hotkey = U256::from(2);
+        let coldkey = U256::from(3);
+        let netuid = NetUid::from(2u16); // Can't be root
+        let coldkeys = vec![U256::from(4), U256::from(5), coldkey];
+        let mut weight = Weight::zero();
+
+        // Initialize ChildKeys for old_hotkey
+        add_network(netuid, 1, 0);
+        AutoStakeDestinationColdkeys::<Test>::insert(old_hotkey, netuid, coldkeys.clone());
+        AutoStakeDestination::<Test>::insert(coldkey, netuid, old_hotkey);
+
+        // Perform the swap
+        SubtensorModule::perform_hotkey_swap_on_all_subnets(
+            &old_hotkey,
+            &new_hotkey,
+            &coldkey,
+            &mut weight,
+        );
+
+        // Verify the swap
+        assert_eq!(
+            AutoStakeDestinationColdkeys::<Test>::get(new_hotkey, netuid),
+            coldkeys
+        );
+        assert!(AutoStakeDestinationColdkeys::<Test>::get(old_hotkey, netuid).is_empty());
+        assert_eq!(
+            AutoStakeDestination::<Test>::get(coldkey, netuid),
+            Some(new_hotkey)
+        );
+    });
+}

runtime/src/lib.rs

@@ -223,7 +223,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: 325,
+    spec_version: 326,
     impl_version: 1,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 1,