diff --git a/pallets/subtensor/src/staking/lock.rs b/pallets/subtensor/src/staking/lock.rs index 866cc91316..bbdb863c0a 100644 --- a/pallets/subtensor/src/staking/lock.rs +++ b/pallets/subtensor/src/staking/lock.rs @@ -1350,24 +1350,33 @@ impl Pallet { Self::ensure_no_active_locks(new_coldkey)?; let mut locks_to_transfer: Vec<(NetUid, T::AccountId, LockState)> = Vec::new(); + let decaying_locks_to_transfer: Vec<(NetUid, bool)> = + DecayingLock::::iter_prefix(old_coldkey).collect(); // Gather locks for old coldkey for ((netuid, hotkey), lock) in Lock::::iter_prefix((old_coldkey,)) { locks_to_transfer.push((netuid, hotkey, lock)); } + for (netuid, decaying) in decaying_locks_to_transfer.iter() { + DecayingLock::::insert(new_coldkey, *netuid, *decaying); + } + // Remove locks for old coldkey and insert for new for (netuid, hotkey, lock) in locks_to_transfer { let now = Self::get_current_block_as_u64(); let unlock_rate = UnlockRate::::get(); let maturity_rate = MaturityRate::::get(); + let perpetual_lock = decaying_locks_to_transfer + .iter() + .any(|(decaying_netuid, decaying)| *decaying_netuid == netuid && !*decaying); let old_lock = ConvictionModel::roll_forward_lock( lock, now, unlock_rate, maturity_rate, Self::is_subnet_owner_hotkey(netuid, &hotkey), - Self::is_perpetual_lock(old_coldkey, netuid), + perpetual_lock, ); let new_lock = ConvictionModel::roll_forward_lock( old_lock.0.clone(), @@ -1375,7 +1384,7 @@ impl Pallet { unlock_rate, maturity_rate, Self::is_subnet_owner_hotkey(netuid, &hotkey), - Self::is_perpetual_lock(new_coldkey, netuid), + perpetual_lock, ) .0; Lock::::remove((old_coldkey.clone(), netuid, hotkey.clone())); @@ -1391,6 +1400,10 @@ impl Pallet { Self::add_aggregate_lock(new_coldkey, &hotkey, netuid, new_lock); } + for (netuid, _) in decaying_locks_to_transfer { + DecayingLock::::remove(old_coldkey, netuid); + } + Ok(()) } diff --git a/pallets/subtensor/src/staking/remove_stake.rs b/pallets/subtensor/src/staking/remove_stake.rs index a305e6ca84..cf640dc661 100644 --- a/pallets/subtensor/src/staking/remove_stake.rs +++ b/pallets/subtensor/src/staking/remove_stake.rs @@ -642,22 +642,8 @@ impl Pallet { } } - // 9) Cleanup all subnet stake locks if any. - let lock_keys: Vec<(T::AccountId, NetUid, T::AccountId)> = Lock::::iter_keys() - .filter(|(_, this_netuid, _)| *this_netuid == netuid) - .collect(); - for (coldkey, netuid, hotkey) in lock_keys { - Lock::::remove((coldkey.clone(), netuid, hotkey.clone())); - Self::maybe_remove_locking_coldkey(&hotkey, netuid, &coldkey); - } - - // 10) Cleanup all subnet hotkey locks if any. - let hotkey_lock_keys: Vec<(NetUid, T::AccountId)> = HotkeyLock::::iter_keys() - .filter(|(this_netuid, _)| *this_netuid == netuid) - .collect(); - for (netuid, hotkey) in hotkey_lock_keys { - HotkeyLock::::remove(netuid, hotkey); - } + // 10) Cleanup all subnet stake locks and lock aggregates if any. + Self::destroy_lock_maps(netuid); Ok(()) } diff --git a/pallets/subtensor/src/tests/locks.rs b/pallets/subtensor/src/tests/locks.rs index 62f78d0710..6aa16fdb70 100644 --- a/pallets/subtensor/src/tests/locks.rs +++ b/pallets/subtensor/src/tests/locks.rs @@ -2989,8 +2989,12 @@ fn test_coldkey_swap_swaps_lock() { .next() .is_none() ); + assert!(!DecayingLock::::contains_key(old_coldkey, netuid)); // New coldkey now has the lock assert!(Lock::::get((new_coldkey, netuid, hotkey)).is_some()); + assert_eq!(DecayingLock::::get(new_coldkey, netuid), Some(false)); + assert!(HotkeyLock::::contains_key(netuid, hotkey)); + assert!(!DecayingHotkeyLock::::contains_key(netuid, hotkey)); }); } diff --git a/pallets/subtensor/src/tests/networks.rs b/pallets/subtensor/src/tests/networks.rs index 65236fed06..4696507e2e 100644 --- a/pallets/subtensor/src/tests/networks.rs +++ b/pallets/subtensor/src/tests/networks.rs @@ -973,6 +973,91 @@ fn destroy_alpha_out_multiple_stakers_pro_rata() { }); } +#[test] +fn destroy_alpha_in_out_stakes_cleans_locking_coldkeys() { + new_test_ext(0).execute_with(|| { + let owner_cold = U256::from(10); + let owner_hot = U256::from(20); + let netuid = add_dynamic_network(&owner_hot, &owner_cold); + remove_owner_registration_stake(netuid); + + let coldkey = U256::from(111); + let hotkey = U256::from(222); + let other_netuid = NetUid::from(u16::from(netuid) + 1); + let lock = LockState { + locked_mass: 10u64.into(), + conviction: U64F64::from_num(1), + last_update: 1, + }; + + Lock::::insert((coldkey, netuid, hotkey), lock.clone()); + LockingColdkeys::::insert((netuid, hotkey, coldkey), ()); + Lock::::insert((coldkey, other_netuid, hotkey), lock); + LockingColdkeys::::insert((other_netuid, hotkey, coldkey), ()); + + assert_ok!(SubtensorModule::destroy_alpha_in_out_stakes(netuid)); + + assert!(!Lock::::contains_key((coldkey, netuid, hotkey))); + assert!(!LockingColdkeys::::contains_key(( + netuid, hotkey, coldkey + ))); + assert!(Lock::::contains_key((coldkey, other_netuid, hotkey))); + assert!(LockingColdkeys::::contains_key(( + other_netuid, + hotkey, + coldkey + ))); + }); +} + +#[test] +fn destroy_alpha_in_out_stakes_cleans_all_lock_aggregates() { + new_test_ext(0).execute_with(|| { + let owner_cold = U256::from(10); + let owner_hot = U256::from(20); + let netuid = add_dynamic_network(&owner_hot, &owner_cold); + remove_owner_registration_stake(netuid); + + let coldkey = U256::from(111); + let hotkey = U256::from(222); + let other_netuid = NetUid::from(u16::from(netuid) + 1); + let lock = LockState { + locked_mass: 10u64.into(), + conviction: U64F64::from_num(1), + last_update: 1, + }; + + HotkeyLock::::insert(netuid, hotkey, lock.clone()); + DecayingHotkeyLock::::insert(netuid, hotkey, lock.clone()); + OwnerLock::::insert(netuid, lock.clone()); + DecayingOwnerLock::::insert(netuid, lock.clone()); + DecayingLock::::insert(coldkey, netuid, false); + + HotkeyLock::::insert(other_netuid, hotkey, lock.clone()); + DecayingHotkeyLock::::insert(other_netuid, hotkey, lock.clone()); + OwnerLock::::insert(other_netuid, lock.clone()); + DecayingOwnerLock::::insert(other_netuid, lock); + DecayingLock::::insert(coldkey, other_netuid, false); + + assert_ok!(SubtensorModule::destroy_alpha_in_out_stakes(netuid)); + + assert!(!HotkeyLock::::contains_key(netuid, hotkey)); + assert!(!DecayingHotkeyLock::::contains_key(netuid, hotkey)); + assert!(!OwnerLock::::contains_key(netuid)); + assert!(!DecayingOwnerLock::::contains_key(netuid)); + assert!(!DecayingLock::::contains_key(coldkey, netuid)); + + assert!(HotkeyLock::::contains_key(other_netuid, hotkey)); + assert!(DecayingHotkeyLock::::contains_key( + other_netuid, + hotkey + )); + assert!(OwnerLock::::contains_key(other_netuid)); + assert!(DecayingOwnerLock::::contains_key(other_netuid)); + assert!(DecayingLock::::contains_key(coldkey, other_netuid)); + }); +} + #[allow(clippy::indexing_slicing)] #[test] fn destroy_alpha_out_many_stakers_complex_distribution() { @@ -2461,10 +2546,13 @@ fn dissolve_clears_all_lock_maps_for_removed_network() { // --- Lock: (coldkey, netuid, hotkey) Lock::::insert((cold_1, net, hot_1), lock_a.clone()); + LockingColdkeys::::insert((net, hot_1, cold_1), ()); Lock::::insert((cold_2, net, hot_2), lock_b.clone()); + LockingColdkeys::::insert((net, hot_2, cold_2), ()); // Same cold/hot on another net should survive. Lock::::insert((cold_1, other_net, hot_1), lock_a.clone()); + LockingColdkeys::::insert((other_net, hot_1, cold_1), ()); // --- HotkeyLock HotkeyLock::::insert(net, hot_1, lock_a.clone()); @@ -2488,6 +2576,8 @@ fn dissolve_clears_all_lock_maps_for_removed_network() { // Sanity checks before dissolve assert!(Lock::::contains_key((cold_1, net, hot_1))); assert!(Lock::::contains_key((cold_2, net, hot_2))); + assert!(LockingColdkeys::::contains_key((net, hot_1, cold_1))); + assert!(LockingColdkeys::::contains_key((net, hot_2, cold_2))); assert!(HotkeyLock::::contains_key(net, hot_1)); assert!(HotkeyLock::::contains_key(net, hot_2)); @@ -2502,6 +2592,9 @@ fn dissolve_clears_all_lock_maps_for_removed_network() { // Sanity: other net keys are present before dissolve. assert!(Lock::::contains_key((cold_1, other_net, hot_1))); + assert!(LockingColdkeys::::contains_key(( + other_net, hot_1, cold_1 + ))); assert!(HotkeyLock::::contains_key(other_net, hot_1)); assert!(DecayingHotkeyLock::::contains_key(other_net, hot_1)); assert!(OwnerLock::::contains_key(other_net)); @@ -2513,6 +2606,8 @@ fn dissolve_clears_all_lock_maps_for_removed_network() { // Ensure removed assert!(!Lock::::contains_key((cold_1, net, hot_1))); assert!(!Lock::::contains_key((cold_2, net, hot_2))); + assert!(!LockingColdkeys::::contains_key((net, hot_1, cold_1))); + assert!(!LockingColdkeys::::contains_key((net, hot_2, cold_2))); assert!(!HotkeyLock::::contains_key(net, hot_1)); assert!(!HotkeyLock::::contains_key(net, hot_2)); @@ -2533,6 +2628,9 @@ fn dissolve_clears_all_lock_maps_for_removed_network() { // Ensure other_net is untouched assert!(Lock::::contains_key((cold_1, other_net, hot_1))); + assert!(LockingColdkeys::::contains_key(( + other_net, hot_1, cold_1 + ))); assert!(HotkeyLock::::contains_key(other_net, hot_1)); assert!(DecayingHotkeyLock::::contains_key(other_net, hot_1)); assert!(OwnerLock::::contains_key(other_net));