Fix issues in conviction

pallets/subtensor/src/staking/lock.rs

@@ -1350,32 +1350,41 @@ impl<T: Config> Pallet<T> {
         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::<T>::iter_prefix(old_coldkey).collect();
 
         // Gather locks for old coldkey
         for ((netuid, hotkey), lock) in Lock::<T>::iter_prefix((old_coldkey,)) {
             locks_to_transfer.push((netuid, hotkey, lock));
         }
 
+        for (netuid, decaying) in decaying_locks_to_transfer.iter() {
+            DecayingLock::<T>::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::<T>::get();
             let maturity_rate = MaturityRate::<T>::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(),
                 now,
                 unlock_rate,
                 maturity_rate,
                 Self::is_subnet_owner_hotkey(netuid, &hotkey),
-                Self::is_perpetual_lock(new_coldkey, netuid),
+                perpetual_lock,
             )
             .0;
             Lock::<T>::remove((old_coldkey.clone(), netuid, hotkey.clone()));
@@ -1391,6 +1400,10 @@ impl<T: Config> Pallet<T> {
             Self::add_aggregate_lock(new_coldkey, &hotkey, netuid, new_lock);
         }
 
+        for (netuid, _) in decaying_locks_to_transfer {
+            DecayingLock::<T>::remove(old_coldkey, netuid);
+        }
+
         Ok(())
     }
 

pallets/subtensor/src/staking/remove_stake.rs

@@ -642,22 +642,8 @@ impl<T: Config> Pallet<T> {
             }
         }
 
-        // 9) Cleanup all subnet stake locks if any.
-        let lock_keys: Vec<(T::AccountId, NetUid, T::AccountId)> = Lock::<T>::iter_keys()
-            .filter(|(_, this_netuid, _)| *this_netuid == netuid)
-            .collect();
-        for (coldkey, netuid, hotkey) in lock_keys {
-            Lock::<T>::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::<T>::iter_keys()
-            .filter(|(this_netuid, _)| *this_netuid == netuid)
-            .collect();
-        for (netuid, hotkey) in hotkey_lock_keys {
-            HotkeyLock::<T>::remove(netuid, hotkey);
-        }
+        // 10) Cleanup all subnet stake locks and lock aggregates if any.
+        Self::destroy_lock_maps(netuid);
 
         Ok(())
     }

pallets/subtensor/src/tests/locks.rs

@@ -2989,8 +2989,12 @@ fn test_coldkey_swap_swaps_lock() {
                 .next()
                 .is_none()
         );
+        assert!(!DecayingLock::<Test>::contains_key(old_coldkey, netuid));
         // New coldkey now has the lock
         assert!(Lock::<Test>::get((new_coldkey, netuid, hotkey)).is_some());
+        assert_eq!(DecayingLock::<Test>::get(new_coldkey, netuid), Some(false));
+        assert!(HotkeyLock::<Test>::contains_key(netuid, hotkey));
+        assert!(!DecayingHotkeyLock::<Test>::contains_key(netuid, hotkey));
     });
 }
 

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::<Test>::insert((coldkey, netuid, hotkey), lock.clone());
+        LockingColdkeys::<Test>::insert((netuid, hotkey, coldkey), ());
+        Lock::<Test>::insert((coldkey, other_netuid, hotkey), lock);
+        LockingColdkeys::<Test>::insert((other_netuid, hotkey, coldkey), ());
+
+        assert_ok!(SubtensorModule::destroy_alpha_in_out_stakes(netuid));
+
+        assert!(!Lock::<Test>::contains_key((coldkey, netuid, hotkey)));
+        assert!(!LockingColdkeys::<Test>::contains_key((
+            netuid, hotkey, coldkey
+        )));
+        assert!(Lock::<Test>::contains_key((coldkey, other_netuid, hotkey)));
+        assert!(LockingColdkeys::<Test>::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::<Test>::insert(netuid, hotkey, lock.clone());
+        DecayingHotkeyLock::<Test>::insert(netuid, hotkey, lock.clone());
+        OwnerLock::<Test>::insert(netuid, lock.clone());
+        DecayingOwnerLock::<Test>::insert(netuid, lock.clone());
+        DecayingLock::<Test>::insert(coldkey, netuid, false);
+
+        HotkeyLock::<Test>::insert(other_netuid, hotkey, lock.clone());
+        DecayingHotkeyLock::<Test>::insert(other_netuid, hotkey, lock.clone());
+        OwnerLock::<Test>::insert(other_netuid, lock.clone());
+        DecayingOwnerLock::<Test>::insert(other_netuid, lock);
+        DecayingLock::<Test>::insert(coldkey, other_netuid, false);
+
+        assert_ok!(SubtensorModule::destroy_alpha_in_out_stakes(netuid));
+
+        assert!(!HotkeyLock::<Test>::contains_key(netuid, hotkey));
+        assert!(!DecayingHotkeyLock::<Test>::contains_key(netuid, hotkey));
+        assert!(!OwnerLock::<Test>::contains_key(netuid));
+        assert!(!DecayingOwnerLock::<Test>::contains_key(netuid));
+        assert!(!DecayingLock::<Test>::contains_key(coldkey, netuid));
+
+        assert!(HotkeyLock::<Test>::contains_key(other_netuid, hotkey));
+        assert!(DecayingHotkeyLock::<Test>::contains_key(
+            other_netuid,
+            hotkey
+        ));
+        assert!(OwnerLock::<Test>::contains_key(other_netuid));
+        assert!(DecayingOwnerLock::<Test>::contains_key(other_netuid));
+        assert!(DecayingLock::<Test>::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::<Test>::insert((cold_1, net, hot_1), lock_a.clone());
+        LockingColdkeys::<Test>::insert((net, hot_1, cold_1), ());
         Lock::<Test>::insert((cold_2, net, hot_2), lock_b.clone());
+        LockingColdkeys::<Test>::insert((net, hot_2, cold_2), ());
 
         // Same cold/hot on another net should survive.
         Lock::<Test>::insert((cold_1, other_net, hot_1), lock_a.clone());
+        LockingColdkeys::<Test>::insert((other_net, hot_1, cold_1), ());
 
         // --- HotkeyLock
         HotkeyLock::<Test>::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::<Test>::contains_key((cold_1, net, hot_1)));
         assert!(Lock::<Test>::contains_key((cold_2, net, hot_2)));
+        assert!(LockingColdkeys::<Test>::contains_key((net, hot_1, cold_1)));
+        assert!(LockingColdkeys::<Test>::contains_key((net, hot_2, cold_2)));
 
         assert!(HotkeyLock::<Test>::contains_key(net, hot_1));
         assert!(HotkeyLock::<Test>::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::<Test>::contains_key((cold_1, other_net, hot_1)));
+        assert!(LockingColdkeys::<Test>::contains_key((
+            other_net, hot_1, cold_1
+        )));
         assert!(HotkeyLock::<Test>::contains_key(other_net, hot_1));
         assert!(DecayingHotkeyLock::<Test>::contains_key(other_net, hot_1));
         assert!(OwnerLock::<Test>::contains_key(other_net));
@@ -2513,6 +2606,8 @@ fn dissolve_clears_all_lock_maps_for_removed_network() {
         // Ensure removed
         assert!(!Lock::<Test>::contains_key((cold_1, net, hot_1)));
         assert!(!Lock::<Test>::contains_key((cold_2, net, hot_2)));
+        assert!(!LockingColdkeys::<Test>::contains_key((net, hot_1, cold_1)));
+        assert!(!LockingColdkeys::<Test>::contains_key((net, hot_2, cold_2)));
 
         assert!(!HotkeyLock::<Test>::contains_key(net, hot_1));
         assert!(!HotkeyLock::<Test>::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::<Test>::contains_key((cold_1, other_net, hot_1)));
+        assert!(LockingColdkeys::<Test>::contains_key((
+            other_net, hot_1, cold_1
+        )));
         assert!(HotkeyLock::<Test>::contains_key(other_net, hot_1));
         assert!(DecayingHotkeyLock::<Test>::contains_key(other_net, hot_1));
         assert!(OwnerLock::<Test>::contains_key(other_net));