diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index e33598472b..6d178bb15f 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -1150,6 +1150,16 @@ pub mod pallet { T::AccountId, OptionQuery, >; + #[pallet::storage] // --- DMAP ( hot, netuid )--> Vec | Returns a list of coldkeys that are autostaking to a hotkey. + pub type AutoStakeDestinationColdkeys = StorageDoubleMap< + _, + Blake2_128Concat, + T::AccountId, + Identity, + NetUid, + Vec, + ValueQuery, + >; #[pallet::storage] // --- DMAP ( cold ) --> (block_expected, new_coldkey) | Maps coldkey to the block to swap at and new coldkey. pub type ColdkeySwapScheduled = StorageMap< diff --git a/pallets/subtensor/src/macros/dispatches.rs b/pallets/subtensor/src/macros/dispatches.rs index 2517217647..7c73ea760d 100644 --- a/pallets/subtensor/src/macros/dispatches.rs +++ b/pallets/subtensor/src/macros/dispatches.rs @@ -2313,9 +2313,20 @@ mod dispatches { current_hotkey != hotkey, Error::::SameAutoStakeHotkeyAlreadySet ); + + // Remove the coldkey from the old hotkey (if present) + AutoStakeDestinationColdkeys::::mutate(current_hotkey.clone(), netuid, |v| { + v.retain(|c| c != &coldkey); + }); } + // Add the coldkey to the new hotkey (if not already present) AutoStakeDestination::::insert(coldkey.clone(), netuid, hotkey.clone()); + AutoStakeDestinationColdkeys::::mutate(hotkey.clone(), netuid, |v| { + if !v.contains(&coldkey) { + v.push(coldkey.clone()); + } + }); Self::deposit_event(Event::AutoStakeDestinationSet { coldkey, diff --git a/pallets/subtensor/src/macros/hooks.rs b/pallets/subtensor/src/macros/hooks.rs index a055668f1a..4b4dfcc781 100644 --- a/pallets/subtensor/src/macros/hooks.rs +++ b/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::()) // Migrate subnet burn cost to 2500 - .saturating_add(migrations::migrate_network_lock_cost_2500::migrate_network_lock_cost_2500::()); + .saturating_add(migrations::migrate_network_lock_cost_2500::migrate_network_lock_cost_2500::()) + // Migrate AutoStakeDestinationColdkeys + .saturating_add(migrations::migrate_auto_stake_destination::migrate_auto_stake_destination::()); weight } diff --git a/pallets/subtensor/src/migrations/migrate_auto_stake_destination.rs b/pallets/subtensor/src/migrations/migrate_auto_stake_destination.rs index 13f165db40..e478a3581a 100644 --- a/pallets/subtensor/src/migrations/migrate_auto_stake_destination.rs +++ b/pallets/subtensor/src/migrations/migrate_auto_stake_destination.rs @@ -57,6 +57,11 @@ pub fn migrate_auto_stake_destination() -> Weight { continue; } AutoStakeDestination::::insert(coldkey, netuid, hotkey.clone()); + AutoStakeDestinationColdkeys::::mutate(hotkey.clone(), netuid, |v| { + if !v.contains(coldkey) { + v.push(coldkey.clone()); + } + }); } old::AutoStakeDestination::::remove(coldkey); diff --git a/pallets/subtensor/src/swap/swap_coldkey.rs b/pallets/subtensor/src/swap/swap_coldkey.rs index afee144e5f..06f5d77f57 100644 --- a/pallets/subtensor/src/swap/swap_coldkey.rs +++ b/pallets/subtensor/src/swap/swap_coldkey.rs @@ -159,7 +159,16 @@ impl Pallet { if let Some(old_auto_stake_hotkey) = AutoStakeDestination::::get(old_coldkey, netuid) { AutoStakeDestination::::remove(old_coldkey, netuid); - AutoStakeDestination::::insert(new_coldkey, netuid, old_auto_stake_hotkey); + AutoStakeDestination::::insert( + new_coldkey, + netuid, + old_auto_stake_hotkey.clone(), + ); + AutoStakeDestinationColdkeys::::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()); + }); } } diff --git a/pallets/subtensor/src/swap/swap_hotkey.rs b/pallets/subtensor/src/swap/swap_hotkey.rs index 82a16bc800..dc84f5a4fd 100644 --- a/pallets/subtensor/src/swap/swap_hotkey.rs +++ b/pallets/subtensor/src/swap/swap_hotkey.rs @@ -524,6 +524,18 @@ impl Pallet { } } + // 6.4 Swap AutoStakeDestination + if let Ok(old_auto_stake_coldkeys) = + AutoStakeDestinationColdkeys::::try_get(old_hotkey, netuid) + { + // Move the vector from old hotkey to new hotkey. + for coldkey in &old_auto_stake_coldkeys { + AutoStakeDestination::::insert(coldkey, netuid, new_hotkey); + } + AutoStakeDestinationColdkeys::::remove(old_hotkey, netuid); + AutoStakeDestinationColdkeys::::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::::try_get(netuid) { diff --git a/pallets/subtensor/src/tests/auto_stake_hotkey.rs b/pallets/subtensor/src/tests/auto_stake_hotkey.rs index 242d469c22..ae73d173fa 100644 --- a/pallets/subtensor/src/tests/auto_stake_hotkey.rs +++ b/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::::insert(hotkey, coldkey); + OwnedHotkeys::::insert(coldkey, vec![hotkey]); + + let netuid = add_dynamic_network(&subnet_owner_hk, &subnet_owner_ck); + Uids::::insert(netuid, hotkey, 1); + Uids::::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::::get(coldkey, netuid), + Some(hotkey) + ); + assert_eq!( + AutoStakeDestinationColdkeys::::get(hotkey, netuid), + vec![coldkey] + ); + assert_eq!( + AutoStakeDestinationColdkeys::::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::::get(coldkey, netuid), + Some(new_hotkey) + ); + assert_eq!( + AutoStakeDestinationColdkeys::::get(hotkey, netuid), + vec![] + ); + assert_eq!( + AutoStakeDestinationColdkeys::::get(new_hotkey, netuid), + vec![coldkey] + ); + }); +} diff --git a/pallets/subtensor/src/tests/migration.rs b/pallets/subtensor/src/tests/migration.rs index d6cc596d26..2d5ada4099 100644 --- a/pallets/subtensor/src/tests/migration.rs +++ b/pallets/subtensor/src/tests/migration.rs @@ -1581,8 +1581,17 @@ fn test_migrate_auto_stake_destination() { AutoStakeDestination::::get(coldkey2, *netuid), Some(hotkey2) ); - } + // Verify entry for AutoStakeDestinationColdkeys + assert_eq!( + AutoStakeDestinationColdkeys::::get(hotkey1, *netuid), + vec![coldkey1] + ); + assert_eq!( + AutoStakeDestinationColdkeys::::get(hotkey2, *netuid), + vec![coldkey2] + ); + } } // Verify old format entries are cleared diff --git a/pallets/subtensor/src/tests/swap_coldkey.rs b/pallets/subtensor/src/tests/swap_coldkey.rs index 59792d5602..7fb77c895e 100644 --- a/pallets/subtensor/src/tests/swap_coldkey.rs +++ b/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::::insert(hotkey, netuid, coldkeys.clone()); + AutoStakeDestination::::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::::get(hotkey, netuid); + assert!(new_coldkeys.contains(&new_coldkey)); + assert!(!new_coldkeys.contains(&old_coldkey)); + assert_eq!( + AutoStakeDestination::::try_get(old_coldkey, netuid), + Err(()) + ); + assert_eq!( + AutoStakeDestination::::try_get(new_coldkey, netuid), + Ok(hotkey) + ); + }); +} diff --git a/pallets/subtensor/src/tests/swap_hotkey.rs b/pallets/subtensor/src/tests/swap_hotkey.rs index 5991baf07a..b75d7967f3 100644 --- a/pallets/subtensor/src/tests/swap_hotkey.rs +++ b/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::::insert(old_hotkey, netuid, coldkeys.clone()); + AutoStakeDestination::::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::::get(new_hotkey, netuid), + coldkeys + ); + assert!(AutoStakeDestinationColdkeys::::get(old_hotkey, netuid).is_empty()); + assert_eq!( + AutoStakeDestination::::get(coldkey, netuid), + Some(new_hotkey) + ); + }); +} diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index a05e97cda1..51932583d3 100644 --- a/runtime/src/lib.rs +++ b/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,