Hotfix/propagate hotkey swap conviction

pallets/subtensor/src/migrations/migrate_fix_subnet_hotkey_lock_swaps.rs

@@ -273,30 +273,51 @@ pub fn migrate_fix_subnet_hotkey_lock_swaps<T: Config>() -> Weight {
         };
         let netuid = NetUid::from(fix.netuid);
 
+        let index_reads: u64;
+        let index_writes: u64;
+        let lock_take_reads: u64;
+        let lock_take_writes: u64;
         let locks_to_fix: Vec<(T::AccountId, LockState)> = if let Some(coldkey) = fix.coldkey {
             let Some(coldkey) = decode_account_id32::<T>(coldkey) else {
                 log::error!("Failed to decode coldkey: {}", coldkey);
                 continue;
             };
-            Lock::<T>::take((coldkey.clone(), netuid, old_hotkey.clone()))
-                .map(|lock| vec![(coldkey, lock)])
-                .unwrap_or_default()
+            index_reads = 0;
+            index_writes = 1;
+            lock_take_reads = 1;
+            LockingColdkeys::<T>::remove((netuid, old_hotkey.clone(), coldkey.clone()));
+            if let Some(lock) = Lock::<T>::take((coldkey.clone(), netuid, old_hotkey.clone())) {
+                lock_take_writes = 1;
+                vec![(coldkey, lock)]
+            } else {
+                lock_take_writes = 0;
+                Vec::new()
+            }
         } else {
-            let locks: Vec<(T::AccountId, LockState)> = Lock::<T>::iter()
-                .filter_map(|((coldkey, lock_netuid, hotkey), lock)| {
-                    (lock_netuid == netuid && hotkey == old_hotkey).then_some((coldkey, lock))
+            let coldkeys: Vec<T::AccountId> =
+                LockingColdkeys::<T>::iter_prefix((netuid, old_hotkey.clone()))
+                    .map(|(coldkey, ())| coldkey)
+                    .collect();
+            let indexed_coldkeys = coldkeys.len() as u64;
+            index_reads = indexed_coldkeys;
+            index_writes = indexed_coldkeys;
+            lock_take_reads = indexed_coldkeys;
+
+            let locks: Vec<(T::AccountId, LockState)> = coldkeys
+                .into_iter()
+                .filter_map(|coldkey| {
+                    LockingColdkeys::<T>::remove((netuid, old_hotkey.clone(), coldkey.clone()));
+                    Lock::<T>::take((coldkey.clone(), netuid, old_hotkey.clone()))
+                        .map(|lock| (coldkey, lock))
                 })
                 .collect();
-            for (coldkey, _) in &locks {
-                Lock::<T>::remove((coldkey.clone(), netuid, old_hotkey.clone()));
-            }
+            lock_take_writes = locks.len() as u64;
             locks
         };
-        let locks_to_fix_count = locks_to_fix.len() as u64;
-        weight = weight.saturating_add(
-            T::DbWeight::get()
-                .reads_writes(locks_to_fix_count.saturating_add(1), locks_to_fix_count),
-        );
+        weight = weight.saturating_add(T::DbWeight::get().reads_writes(
+            index_reads.saturating_add(lock_take_reads),
+            index_writes.saturating_add(lock_take_writes),
+        ));
 
         if locks_to_fix.is_empty() {
             missing_locks = missing_locks.saturating_add(1);
@@ -323,7 +344,8 @@ pub fn migrate_fix_subnet_hotkey_lock_swaps<T: Config>() -> Weight {
             }
 
             Lock::<T>::insert((coldkey.clone(), netuid, new_hotkey.clone()), lock.clone());
-            weight = weight.saturating_add(T::DbWeight::get().writes(1));
+            LockingColdkeys::<T>::insert((netuid, new_hotkey.clone(), coldkey.clone()), ());
+            weight = weight.saturating_add(T::DbWeight::get().writes(2));
 
             if !new_hotkey_is_owner {
                 add_to_aggregate::<T>(&coldkey, netuid, &new_hotkey, &lock);

pallets/subtensor/src/swap/swap_hotkey.rs

@@ -46,16 +46,24 @@ impl<T: Config> Pallet<T> {
             Error::<T>::NonAssociatedColdKey
         );
 
-        // 3. Initialize the weight for this operation
+        // 3. If the new hotkey already exists globally, ensure the coldkey owns it
+        if Self::hotkey_account_exists(new_hotkey) {
+            ensure!(
+                Self::coldkey_owns_hotkey(&coldkey, new_hotkey),
+                Error::<T>::NonAssociatedColdKey
+            );
+        }
+
+        // 4. Initialize the weight for this operation
         let mut weight = T::DbWeight::get().reads(2);
 
-        // 4. Ensure the new hotkey is different from the old one
+        // 5. Ensure the new hotkey is different from the old one
         ensure!(old_hotkey != new_hotkey, Error::<T>::NewHotKeyIsSameWithOld);
 
-        // 5. Get the current block number
+        // 6. Get the current block number
         let block: u64 = Self::get_current_block_as_u64();
 
-        // 6. Ensure the transaction rate limit is not exceeded
+        // 7. Ensure the transaction rate limit is not exceeded
         ensure!(
             !Self::exceeds_tx_rate_limit(Self::get_last_tx_block(&coldkey), block),
             Error::<T>::HotKeySetTxRateLimitExceeded
@@ -64,14 +72,14 @@ impl<T: Config> Pallet<T> {
         weight.saturating_accrue(T::DbWeight::get().reads(2));
 
         match netuid {
-            // 7. Ensure the hotkey is not registered on the network before, if netuid is provided
+            // 8. Ensure the hotkey is not registered on the network before, if netuid is provided
             Some(netuid) => {
                 ensure!(
                     !Self::is_hotkey_registered_on_specific_network(new_hotkey, netuid),
                     Error::<T>::HotKeyAlreadyRegisteredInSubNet
                 );
             }
-            // 7.1 Ensure the new hotkey is not already registered on any network, only if netuid is none
+            // 8.1 Ensure the new hotkey is not already registered on any network, only if netuid is none
             None => {
                 ensure!(
                     !Self::is_hotkey_registered_on_any_network(new_hotkey),
@@ -80,7 +88,7 @@ impl<T: Config> Pallet<T> {
             }
         }
 
-        // 7.2 If the swap touches the root subnet, require that new_hotkey is clean
+        // 8.2 If the swap touches the root subnet, require that new_hotkey is clean
         // on root (no outstanding claimable rate and no existing root stake). Merging
         // a non-empty rate-book would either violate total conservation or misallocate
         // dividends across coldkeys that never staked on old_hotkey.
@@ -96,46 +104,46 @@ impl<T: Config> Pallet<T> {
             );
         }
 
-        // 8. Swap LastTxBlock
+        // 9. Swap LastTxBlock
         let last_tx_block: u64 = Self::get_last_tx_block(old_hotkey);
         Self::set_last_tx_block(new_hotkey, last_tx_block);
         weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
 
-        // 9. Swap LastTxBlockDelegateTake
+        // 10. Swap LastTxBlockDelegateTake
         let last_tx_block_delegate_take: u64 = Self::get_last_tx_block_delegate_take(old_hotkey);
         Self::set_last_tx_block_delegate_take(new_hotkey, last_tx_block_delegate_take);
         weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
 
-        // 10. Swap LastTxBlockChildKeyTake
+        // 11. Swap LastTxBlockChildKeyTake
         let last_tx_block_child_key_take: u64 = Self::get_last_tx_block_childkey_take(old_hotkey);
         Self::set_last_tx_block_childkey(new_hotkey, last_tx_block_child_key_take);
         weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
 
-        // 11. fork for swap hotkey on a specific subnet case after do the common check
+        // 12. fork for swap hotkey on a specific subnet case after do the common check
         if let Some(netuid) = netuid {
             return Self::swap_hotkey_on_subnet(
                 &coldkey, old_hotkey, new_hotkey, netuid, weight, keep_stake,
             );
         };
 
         // Start to do everything for swap hotkey on all subnets case
-        // 12. Get the cost for swapping the key
+        // 13. Get the cost for swapping the key
         let swap_cost = Self::get_key_swap_cost();
         log::debug!("Swap cost: {swap_cost:?}");
 
-        // 13. Ensure the coldkey has enough balance to pay for the swap
+        // 14. Ensure the coldkey has enough balance to pay for the swap
         ensure!(
             Self::can_remove_balance_from_coldkey_account(&coldkey, swap_cost.into()),
             Error::<T>::NotEnoughBalanceToPaySwapHotKey
         );
 
         weight.saturating_accrue(T::DbWeight::get().reads_writes(3, 0));
 
-        // 14. Remove the swap cost from the coldkey's account + Recycle the tokens
+        // 15. Remove the swap cost from the coldkey's account + Recycle the tokens
         Self::recycle_tao(&coldkey, swap_cost.into())?;
         weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 2));
 
-        // 19. Perform the hotkey swap
+        // 16. Perform the hotkey swap
         Self::perform_hotkey_swap_on_all_subnets(
             old_hotkey,
             new_hotkey,
@@ -144,18 +152,18 @@ impl<T: Config> Pallet<T> {
             keep_stake,
         )?;
 
-        // 20. Update the last transaction block for the coldkey
+        // 17. Update the last transaction block for the coldkey
         Self::set_last_tx_block(&coldkey, block);
         weight.saturating_accrue(T::DbWeight::get().writes(1));
 
-        // 21. Emit an event for the hotkey swap
+        // 18. Emit an event for the hotkey swap
         Self::deposit_event(Event::HotkeySwapped {
             coldkey,
             old_hotkey: old_hotkey.clone(),
             new_hotkey: new_hotkey.clone(),
         });
 
-        // 22. Return the weight of the operation
+        // 19. Return the weight of the operation
         Ok(Some(weight).into())
     }
 

pallets/subtensor/src/tests/migration.rs

@@ -123,14 +123,17 @@ fn test_migrate_fix_subnet_hotkey_lock_swaps_moves_or_discards_conflicts() {
             (coldkey_to_move, netuid, old_hotkey),
             moved_lock.clone(),
         );
+        LockingColdkeys::<Test>::insert((netuid, old_hotkey, coldkey_to_move), ());
         Lock::<Test>::insert(
             (coldkey_with_conflict, netuid, old_hotkey),
             discarded_lock.clone(),
         );
+        LockingColdkeys::<Test>::insert((netuid, old_hotkey, coldkey_with_conflict), ());
         Lock::<Test>::insert(
             (coldkey_with_conflict, netuid, new_hotkey),
             existing_destination_lock.clone(),
         );
+        LockingColdkeys::<Test>::insert((netuid, new_hotkey, coldkey_with_conflict), ());
         DecayingLock::<Test>::insert(coldkey_to_move, netuid, false);
         DecayingLock::<Test>::insert(coldkey_with_conflict, netuid, false);
         DecayingLock::<Test>::insert(chained_coldkey, chained_netuid, false);
@@ -148,6 +151,10 @@ fn test_migrate_fix_subnet_hotkey_lock_swaps_moves_or_discards_conflicts() {
             (chained_coldkey, chained_netuid, chained_first_hotkey),
             chained_lock.clone(),
         );
+        LockingColdkeys::<Test>::insert(
+            (chained_netuid, chained_first_hotkey, chained_coldkey),
+            (),
+        );
         HotkeyLock::<Test>::insert(chained_netuid, chained_first_hotkey, chained_lock.clone());
 
         let weight =
@@ -157,14 +164,34 @@ fn test_migrate_fix_subnet_hotkey_lock_swaps_moves_or_discards_conflicts() {
         assert!(HasMigrationRun::<Test>::get(&migration_name));
         assert!(Lock::<Test>::get((coldkey_to_move, netuid, old_hotkey)).is_none());
         assert!(Lock::<Test>::get((coldkey_with_conflict, netuid, old_hotkey)).is_none());
+        assert!(!LockingColdkeys::<Test>::contains_key((
+            netuid,
+            old_hotkey,
+            coldkey_to_move
+        )));
+        assert!(!LockingColdkeys::<Test>::contains_key((
+            netuid,
+            old_hotkey,
+            coldkey_with_conflict
+        )));
         assert_eq!(
             Lock::<Test>::get((coldkey_to_move, netuid, new_hotkey)),
             Some(moved_lock.clone())
         );
+        assert!(LockingColdkeys::<Test>::contains_key((
+            netuid,
+            new_hotkey,
+            coldkey_to_move
+        )));
         assert_eq!(
             Lock::<Test>::get((coldkey_with_conflict, netuid, new_hotkey)),
             Some(existing_destination_lock.clone())
         );
+        assert!(LockingColdkeys::<Test>::contains_key((
+            netuid,
+            new_hotkey,
+            coldkey_with_conflict
+        )));
         assert!(HotkeyLock::<Test>::get(netuid, old_hotkey).is_none());
 
         let new_aggregate = HotkeyLock::<Test>::get(netuid, new_hotkey)
@@ -193,10 +220,25 @@ fn test_migrate_fix_subnet_hotkey_lock_swaps_moves_or_discards_conflicts() {
             chained_middle_hotkey
         ))
         .is_none());
+        assert!(!LockingColdkeys::<Test>::contains_key((
+            chained_netuid,
+            chained_first_hotkey,
+            chained_coldkey
+        )));
+        assert!(!LockingColdkeys::<Test>::contains_key((
+            chained_netuid,
+            chained_middle_hotkey,
+            chained_coldkey
+        )));
         assert_eq!(
             Lock::<Test>::get((chained_coldkey, chained_netuid, chained_final_hotkey)),
             Some(chained_lock.clone())
         );
+        assert!(LockingColdkeys::<Test>::contains_key((
+            chained_netuid,
+            chained_final_hotkey,
+            chained_coldkey
+        )));
         assert!(HotkeyLock::<Test>::get(chained_netuid, chained_first_hotkey).is_none());
         assert!(HotkeyLock::<Test>::get(chained_netuid, chained_middle_hotkey).is_none());
         assert_eq!(

pallets/subtensor/src/tests/swap_hotkey_with_subnet.rs

@@ -926,7 +926,7 @@ fn test_swap_owner_new_hotkey_already_exists() {
                 Some(netuid),
                 false
             ),
-            Error::<Test>::HotKeyAlreadyRegisteredInSubNet
+            Error::<Test>::NonAssociatedColdKey
         );
 
         // Verify the swap

runtime/src/lib.rs

@@ -234,7 +234,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: 417,
+    spec_version: 418,
     impl_version: 1,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 1,