Optimize locks to avoid full iteration of lock map

pallets/subtensor/src/lib.rs

@@ -1532,6 +1532,18 @@ pub mod pallet {
         OptionQuery,
     >;
 
+    /// --- DMAP ( netuid, hotkey ) --> Vec<coldkey> | Coldkeys with non-zero locks targeting this hotkey on this subnet.
+    #[pallet::storage]
+    pub type LockingColdkeys<T: Config> = StorageDoubleMap<
+        _,
+        Identity,
+        NetUid, // subnet
+        Blake2_128Concat,
+        T::AccountId, // hotkey
+        Vec<T::AccountId>,
+        ValueQuery,
+    >;
+
     /// --- DMAP ( netuid, hotkey ) --> LockState | Total lock per hotkey per subnet.
     #[pallet::storage]
     pub type HotkeyLock<T: Config> = StorageDoubleMap<

pallets/subtensor/src/macros/hooks.rs

@@ -180,7 +180,8 @@ mod hooks {
                 // Remove deprecated conviction lock storage.
                 .saturating_add(migrations::migrate_remove_deprecated_conviction_maps::migrate_remove_deprecated_conviction_maps::<T>())
                 // Reset testnet conviction lock storage before deploying the current design.
-                .saturating_add(migrations::migrate_reset_tnet_conviction_locks::migrate_reset_tnet_conviction_locks::<T>());
+                .saturating_add(migrations::migrate_reset_tnet_conviction_locks::migrate_reset_tnet_conviction_locks::<T>())
+                .saturating_add(migrations::migrate_populate_locking_coldkeys::migrate_populate_locking_coldkeys::<T>());
             weight
         }
 

pallets/subtensor/src/migrations/migrate_populate_locking_coldkeys.rs

@@ -0,0 +1,72 @@
+use alloc::string::String;
+use frame_support::{traits::Get, weights::Weight};
+
+use crate::{Config, HasMigrationRun, Lock, Pallet as Subtensor, staking::lock::ConvictionModel};
+
+const MIGRATION_NAME: &[u8] = b"migrate_populate_locking_coldkeys";
+
+pub fn migrate_populate_locking_coldkeys<T: Config>() -> Weight {
+    let mut weight = T::DbWeight::get().reads(1);
+
+    if HasMigrationRun::<T>::get(MIGRATION_NAME) {
+        log::info!(
+            "Migration '{}' already executed - skipping",
+            String::from_utf8_lossy(MIGRATION_NAME)
+        );
+        return weight;
+    }
+
+    log::info!(
+        "Running migration '{}'",
+        String::from_utf8_lossy(MIGRATION_NAME)
+    );
+
+    let now = Subtensor::<T>::get_current_block_as_u64();
+    let unlock_rate = crate::UnlockRate::<T>::get();
+    let maturity_rate = crate::MaturityRate::<T>::get();
+    let mut scanned_count = 0u64;
+    let mut indexed_count = 0u64;
+    let mut removed_count = 0u64;
+    let mut locks_to_remove = sp_std::vec::Vec::new();
+
+    for ((coldkey, netuid, hotkey), lock) in Lock::<T>::iter() {
+        scanned_count = scanned_count.saturating_add(1);
+        let rolled = ConvictionModel::roll_forward_lock(
+            lock,
+            now,
+            unlock_rate,
+            maturity_rate,
+            Subtensor::<T>::is_subnet_owner_hotkey(netuid, &hotkey),
+            Subtensor::<T>::is_perpetual_lock(&coldkey, netuid),
+        );
+
+        if !rolled.is_zero() {
+            Subtensor::<T>::add_locking_coldkey(&hotkey, netuid, &coldkey);
+            indexed_count = indexed_count.saturating_add(1);
+        } else {
+            locks_to_remove.push((coldkey, netuid, hotkey));
+        }
+    }
+
+    for (coldkey, netuid, hotkey) in locks_to_remove {
+        Lock::<T>::remove((coldkey, netuid, hotkey));
+        removed_count = removed_count.saturating_add(1);
+    }
+
+    weight = weight.saturating_add(T::DbWeight::get().reads(scanned_count));
+    weight = weight
+        .saturating_add(T::DbWeight::get().writes(indexed_count.saturating_add(removed_count)));
+
+    HasMigrationRun::<T>::insert(MIGRATION_NAME, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    log::info!(
+        "Migration '{}' completed. scanned_entries={}, indexed_entries={}, removed_zero_entries={}",
+        String::from_utf8_lossy(MIGRATION_NAME),
+        scanned_count,
+        indexed_count,
+        removed_count
+    );
+
+    weight
+}

pallets/subtensor/src/migrations/mod.rs

@@ -32,6 +32,7 @@ pub mod migrate_network_lock_cost_2500;
 pub mod migrate_network_lock_reduction_interval;
 pub mod migrate_orphaned_storage_items;
 pub mod migrate_pending_emissions;
+pub mod migrate_populate_locking_coldkeys;
 pub mod migrate_populate_owned_hotkeys;
 pub mod migrate_rao;
 pub mod migrate_rate_limit_keys;

pallets/subtensor/src/staking/lock.rs

@@ -9,6 +9,7 @@ use substrate_fixed::types::{I64F64, U64F64};
 use subtensor_runtime_common::NetUid;
 
 pub const ONE_YEAR: u64 = 7200 * 365 + 1800;
+pub const LOCK_STATE_ZERO_THRESHOLD: u64 = 100;
 
 /// Exponential lock state for a coldkey on a subnet.
 #[crate::freeze_struct("1f6be20a66128b8d")]
@@ -22,6 +23,13 @@ pub struct LockState {
     pub last_update: u64,
 }
 
+impl LockState {
+    pub fn is_zero(&self) -> bool {
+        self.locked_mass < AlphaBalance::from(LOCK_STATE_ZERO_THRESHOLD)
+            && self.conviction < U64F64::saturating_from_num(LOCK_STATE_ZERO_THRESHOLD)
+    }
+}
+
 /// A struct that incapsulates Lock primitives such as adding, removing,
 /// rolling, and updating aggregates.
 ///
@@ -439,24 +447,55 @@ impl ConvictionModel {
             rolled.conviction = U64F64::saturating_from_num(u64::from(rolled.locked_mass));
         }
 
+        if rolled.is_zero() {
+            rolled.locked_mass = AlphaBalance::ZERO;
+            rolled.conviction = U64F64::saturating_from_num(0);
+        }
+
         rolled
     }
 }
 
 impl<T: Config> Pallet<T> {
+    pub fn add_locking_coldkey(hotkey: &T::AccountId, netuid: NetUid, coldkey: &T::AccountId) {
+        let mut coldkeys = LockingColdkeys::<T>::get(netuid, hotkey);
+        if coldkeys.contains(coldkey) {
+            return;
+        }
+        coldkeys.push(coldkey.clone());
+        LockingColdkeys::<T>::insert(netuid, hotkey, coldkeys);
+    }
+
+    pub fn maybe_remove_locking_coldkey(
+        hotkey: &T::AccountId,
+        netuid: NetUid,
+        coldkey: &T::AccountId,
+    ) {
+        let mut coldkeys = LockingColdkeys::<T>::get(netuid, hotkey);
+        let Some(position) = coldkeys.iter().position(|existing| existing == coldkey) else {
+            return;
+        };
+        coldkeys.remove(position);
+        if coldkeys.is_empty() {
+            LockingColdkeys::<T>::remove(netuid, hotkey);
+        } else {
+            LockingColdkeys::<T>::insert(netuid, hotkey, coldkeys);
+        }
+    }
+
     pub fn insert_lock_state(
         coldkey: &T::AccountId,
         netuid: NetUid,
         hotkey: &T::AccountId,
         lock_state: LockState,
     ) {
-        if !lock_state.locked_mass.is_zero()
-            || lock_state.conviction > U64F64::saturating_from_num(0)
-        {
-            Lock::<T>::insert((coldkey, netuid, hotkey), lock_state);
-        } else {
+        if lock_state.is_zero() {
+            Self::maybe_remove_locking_coldkey(hotkey, netuid, coldkey);
             // If there is no record previously, this is a no-op
             Lock::<T>::remove((coldkey, netuid, hotkey));
+        } else {
+            Self::add_locking_coldkey(hotkey, netuid, coldkey);
+            Lock::<T>::insert((coldkey, netuid, hotkey), lock_state);
         }
     }
 
@@ -504,11 +543,11 @@ impl<T: Config> Pallet<T> {
         }
     }
 
-    fn is_subnet_owner_hotkey(netuid: NetUid, hotkey: &T::AccountId) -> bool {
+    pub(crate) fn is_subnet_owner_hotkey(netuid: NetUid, hotkey: &T::AccountId) -> bool {
         hotkey == &SubnetOwnerHotkey::<T>::get(netuid)
     }
 
-    fn is_perpetual_lock(coldkey: &T::AccountId, netuid: NetUid) -> bool {
+    pub(crate) fn is_perpetual_lock(coldkey: &T::AccountId, netuid: NetUid) -> bool {
         DecayingLock::<T>::get(coldkey, netuid) == Some(false)
     }
 
@@ -1359,6 +1398,7 @@ impl<T: Config> Pallet<T> {
                 Self::is_perpetual_lock(new_coldkey, netuid),
             );
             Lock::<T>::remove((old_coldkey.clone(), netuid, hotkey.clone()));
+            Self::maybe_remove_locking_coldkey(&hotkey, netuid, old_coldkey);
             Self::reduce_aggregate_lock(
                 old_coldkey,
                 &hotkey,
@@ -1417,10 +1457,19 @@ impl<T: Config> Pallet<T> {
             reads = reads.saturating_add(5);
         }
 
-        if !netuids_to_transfer.is_empty() {
-            for ((coldkey, netuid, hotkey), lock) in Lock::<T>::iter() {
-                if hotkey == *old_hotkey {
-                    locks_to_transfer.push((coldkey, netuid, lock));
+        // Build a concrete transfer list from the hotkey-to-coldkey index.
+        // The index can contain stale coldkeys, so only locks that still exist
+        // are carried forward; missing locks are pruned from the index.
+        for (netuid, _, _) in &netuids_to_transfer {
+            let locking_coldkeys = LockingColdkeys::<T>::get(*netuid, old_hotkey);
+            reads = reads.saturating_add(1);
+
+            for coldkey in locking_coldkeys {
+                if let Some(lock) = Lock::<T>::get((coldkey.clone(), *netuid, old_hotkey.clone())) {
+                    locks_to_transfer.push((coldkey, *netuid, lock));
+                } else {
+                    Self::maybe_remove_locking_coldkey(old_hotkey, *netuid, &coldkey);
+                    writes = writes.saturating_add(1);
                 }
                 reads = reads.saturating_add(1);
             }
@@ -1454,6 +1503,7 @@ impl<T: Config> Pallet<T> {
                 perpetual_lock,
             );
             Lock::<T>::remove((coldkey.clone(), netuid, old_hotkey.clone()));
+            Self::maybe_remove_locking_coldkey(old_hotkey, netuid, &coldkey);
             Self::insert_lock_state(&coldkey, netuid, new_hotkey, moved);
             writes = writes.saturating_add(2);
         }
@@ -1612,6 +1662,7 @@ impl<T: Config> Pallet<T> {
                 );
 
                 Lock::<T>::remove((coldkey.clone(), netuid, origin_hotkey.clone()));
+                Self::maybe_remove_locking_coldkey(&origin_hotkey, netuid, coldkey);
                 Self::insert_lock_state(coldkey, netuid, destination_hotkey, lock.clone());
                 Self::reduce_aggregate_lock(
                     coldkey,
@@ -1813,43 +1864,25 @@ impl<T: Config> Pallet<T> {
 
     /// Destroys all lock maps for network dissolution
     pub fn destroy_lock_maps(netuid: NetUid) {
+        // LockingColdkeys: (netuid, hotkey) -> Vec<coldkey>
         // Lock: (coldkey, netuid, hotkey)
         {
-            let to_rm: sp_std::vec::Vec<(T::AccountId, T::AccountId)> = Lock::<T>::iter()
-                .filter_map(
-                    |((cold, n, hot), _)| {
-                        if n == netuid { Some((cold, hot)) } else { None }
-                    },
-                )
-                .collect();
+            let to_rm: sp_std::vec::Vec<(T::AccountId, sp_std::vec::Vec<T::AccountId>)> =
+                LockingColdkeys::<T>::iter_prefix(netuid).collect();
 
-            for (cold, hot) in to_rm {
-                Lock::<T>::remove((cold, netuid, hot));
+            for (hot, coldkeys) in to_rm {
+                for cold in coldkeys {
+                    Lock::<T>::remove((cold, netuid, hot.clone()));
+                }
             }
+            let _ = LockingColdkeys::<T>::clear_prefix(netuid, u32::MAX, None);
         }
 
         // HotkeyLock: (netuid, hotkey) → LockState
-        {
-            let to_rm: sp_std::vec::Vec<T::AccountId> = HotkeyLock::<T>::iter_prefix(netuid)
-                .map(|(hot, _)| hot)
-                .collect();
-
-            for hot in to_rm {
-                HotkeyLock::<T>::remove(netuid, hot);
-            }
-        }
+        let _ = HotkeyLock::<T>::clear_prefix(netuid, u32::MAX, None);
 
         // DecayingHotkeyLock: (netuid, hotkey)
-        {
-            let to_rm: sp_std::vec::Vec<T::AccountId> =
-                DecayingHotkeyLock::<T>::iter_prefix(netuid)
-                    .map(|(hot, _)| hot)
-                    .collect();
-
-            for hot in to_rm {
-                DecayingHotkeyLock::<T>::remove(netuid, hot);
-            }
-        }
+        let _ = DecayingHotkeyLock::<T>::clear_prefix(netuid, u32::MAX, None);
 
         // OwnerLock / DecayingOwnerLock: (netuid)
         OwnerLock::<T>::remove(netuid);

pallets/subtensor/src/staking/remove_stake.rs

@@ -615,7 +615,8 @@ impl<T: Config> Pallet<T> {
             .filter(|(_, this_netuid, _)| *this_netuid == netuid)
             .collect();
         for (coldkey, netuid, hotkey) in lock_keys {
-            Lock::<T>::remove((coldkey, netuid, hotkey));
+            Lock::<T>::remove((coldkey.clone(), netuid, hotkey.clone()));
+            Self::maybe_remove_locking_coldkey(&hotkey, netuid, &coldkey);
         }
 
         // 10) Cleanup all subnet hotkey locks if any.