refacto to handle subsubnet when trimming uids

Author: l0r1s

pallets/subtensor/src/subnets/uids.rs

@@ -148,11 +148,11 @@ impl<T: Config> Pallet<T> {
 
         MaxAllowedUids::<T>::insert(netuid, max_n);
 
-        let owner = SubnetOwner::<T>::get(netuid);
-        let owner_uids = BTreeSet::from_iter(Self::get_immune_owner_uids(netuid, &owner));
         let current_n = Self::get_subnetwork_n(netuid);
-
         if current_n > max_n {
+            let owner = SubnetOwner::<T>::get(netuid);
+            let owner_uids = BTreeSet::from_iter(Self::get_immune_owner_uids(netuid, &owner));
+
             // Count the number of immune UIDs
             let mut immune_count: u16 = 0;
             for uid in 0..current_n {
@@ -167,7 +167,8 @@ impl<T: Config> Pallet<T> {
                 immune_percentage < T::MaxImmuneUidsPercentage::get(),
                 Error::<T>::InvalidValue
             );
-
+            
+            
             // Get all emissions with their UIDs and sort by emission (descending)
             // This ensures we keep the highest emitters and remove the lowest ones
             let mut emissions = Emission::<T>::get(netuid)
@@ -176,9 +177,9 @@ impl<T: Config> Pallet<T> {
                 .collect::<Vec<_>>();
             emissions.sort_by_key(|(_, emission)| cmp::Reverse(*emission));
 
-            // Remove uids from the end (lowest emitters) until we reach the new maximum
             let mut removed_uids = BTreeSet::new();
             let mut uids_left_to_process = current_n;
+            let subsubnets_count = SubsubnetCountCurrent::<T>::get(netuid).into();
 
             // Iterate from the end (lowest emitters) to the beginning
             for i in (0..current_n).rev() {
@@ -187,15 +188,17 @@ impl<T: Config> Pallet<T> {
                 }
 
                 if let Some((uid, _)) = emissions.get(i as usize).cloned() {
+                    let neuron_uid = uid as u16;
+
                     // Skip subnet owner's or temporally immune uids
-                    if owner_uids.contains(&(uid as u16))
-                        || Self::get_neuron_is_immune(netuid, uid as u16)
+                    if owner_uids.contains(&neuron_uid)
+                        || Self::get_neuron_is_immune(netuid, neuron_uid)
                     {
                         continue;
                     }
 
                     // Remove hotkey related storage items if hotkey exists
-                    if let Ok(hotkey) = Keys::<T>::try_get(netuid, uid as u16) {
+                    if let Ok(hotkey) = Keys::<T>::try_get(netuid, neuron_uid) {
                         Uids::<T>::remove(netuid, &hotkey);
                         IsNetworkMember::<T>::remove(&hotkey, netuid);
                         LastHotkeyEmissionOnNetuid::<T>::remove(&hotkey, netuid);
@@ -208,10 +211,13 @@ impl<T: Config> Pallet<T> {
 
                     // Remove all storage items associated with this uid
                     #[allow(unknown_lints)]
-                    Keys::<T>::remove(netuid, uid as u16);
-                    BlockAtRegistration::<T>::remove(netuid, uid as u16);
-                    Weights::<T>::remove(netuid, uid as u16);
-                    Bonds::<T>::remove(netuid, uid as u16);
+                    Keys::<T>::remove(netuid, neuron_uid);
+                    BlockAtRegistration::<T>::remove(netuid, neuron_uid);
+                    for subid in 0..subsubnets_count {
+                        let netuid_index = Self::get_subsubnet_storage_index(netuid, subid.into());
+                        Weights::<T>::remove(netuid_index, neuron_uid);
+                        Bonds::<T>::remove(netuid_index, neuron_uid);
+                    }
 
                     // Remove from emissions array and track as removed
                     emissions.remove(i.into());
@@ -233,42 +239,36 @@ impl<T: Config> Pallet<T> {
             let trust = Trust::<T>::get(netuid);
             let active = Active::<T>::get(netuid);
             let consensus = Consensus::<T>::get(netuid);
-            let incentive = Incentive::<T>::get(netuid);
             let dividends = Dividends::<T>::get(netuid);
-            let lastupdate = LastUpdate::<T>::get(netuid);
             let pruning_scores = PruningScores::<T>::get(netuid);
             let vtrust = ValidatorTrust::<T>::get(netuid);
             let vpermit = ValidatorPermit::<T>::get(netuid);
             let stake_weight = StakeWeight::<T>::get(netuid);
 
             // Create trimmed arrays by extracting values for kept uids only
             // Pre-allocate vectors with exact capacity for efficiency
-            let mut trimmed_ranks = Vec::with_capacity(trimmed_uids.len());
-            let mut trimmed_trust = Vec::with_capacity(trimmed_uids.len());
-            let mut trimmed_active = Vec::with_capacity(trimmed_uids.len());
-            let mut trimmed_consensus = Vec::with_capacity(trimmed_uids.len());
-            let mut trimmed_incentive = Vec::with_capacity(trimmed_uids.len());
-            let mut trimmed_dividends = Vec::with_capacity(trimmed_uids.len());
-            let mut trimmed_lastupdate = Vec::with_capacity(trimmed_uids.len());
-            let mut trimmed_pruning_scores = Vec::with_capacity(trimmed_uids.len());
-            let mut trimmed_vtrust = Vec::with_capacity(trimmed_uids.len());
-            let mut trimmed_vpermit = Vec::with_capacity(trimmed_uids.len());
-            let mut trimmed_stake_weight = Vec::with_capacity(trimmed_uids.len());
+            let len = trimmed_uids.len();
+            let mut trimmed_ranks = Vec::with_capacity(len);
+            let mut trimmed_trust = Vec::with_capacity(len);
+            let mut trimmed_active = Vec::with_capacity(len);
+            let mut trimmed_consensus = Vec::with_capacity(len);
+            let mut trimmed_dividends = Vec::with_capacity(len);
+            let mut trimmed_pruning_scores = Vec::with_capacity(len);
+            let mut trimmed_vtrust = Vec::with_capacity(len);
+            let mut trimmed_vpermit = Vec::with_capacity(len);
+            let mut trimmed_stake_weight = Vec::with_capacity(len);
 
             // Single iteration to extract values for all kept uids
-            for &old_uid in &trimmed_uids {
-                trimmed_ranks.push(ranks.get(old_uid).cloned().unwrap_or_default());
-                trimmed_trust.push(trust.get(old_uid).cloned().unwrap_or_default());
-                trimmed_active.push(active.get(old_uid).cloned().unwrap_or_default());
-                trimmed_consensus.push(consensus.get(old_uid).cloned().unwrap_or_default());
-                trimmed_incentive.push(incentive.get(old_uid).cloned().unwrap_or_default());
-                trimmed_dividends.push(dividends.get(old_uid).cloned().unwrap_or_default());
-                trimmed_lastupdate.push(lastupdate.get(old_uid).cloned().unwrap_or_default());
-                trimmed_pruning_scores
-                    .push(pruning_scores.get(old_uid).cloned().unwrap_or_default());
-                trimmed_vtrust.push(vtrust.get(old_uid).cloned().unwrap_or_default());
-                trimmed_vpermit.push(vpermit.get(old_uid).cloned().unwrap_or_default());
-                trimmed_stake_weight.push(stake_weight.get(old_uid).cloned().unwrap_or_default());
+            for &uid in &trimmed_uids {
+                trimmed_ranks.push(ranks.get(uid).cloned().unwrap_or_default());
+                trimmed_trust.push(trust.get(uid).cloned().unwrap_or_default());
+                trimmed_active.push(active.get(uid).cloned().unwrap_or_default());
+                trimmed_consensus.push(consensus.get(uid).cloned().unwrap_or_default());
+                trimmed_dividends.push(dividends.get(uid).cloned().unwrap_or_default());
+                trimmed_pruning_scores.push(pruning_scores.get(uid).cloned().unwrap_or_default());
+                trimmed_vtrust.push(vtrust.get(uid).cloned().unwrap_or_default());
+                trimmed_vpermit.push(vpermit.get(uid).cloned().unwrap_or_default());
+                trimmed_stake_weight.push(stake_weight.get(uid).cloned().unwrap_or_default());
             }
 
             // Update storage with trimmed arrays
@@ -277,14 +277,29 @@ impl<T: Config> Pallet<T> {
             Trust::<T>::insert(netuid, trimmed_trust);
             Active::<T>::insert(netuid, trimmed_active);
             Consensus::<T>::insert(netuid, trimmed_consensus);
-            Incentive::<T>::insert(netuid, trimmed_incentive);
             Dividends::<T>::insert(netuid, trimmed_dividends);
-            LastUpdate::<T>::insert(netuid, trimmed_lastupdate);
             PruningScores::<T>::insert(netuid, trimmed_pruning_scores);
             ValidatorTrust::<T>::insert(netuid, trimmed_vtrust);
             ValidatorPermit::<T>::insert(netuid, trimmed_vpermit);
             StakeWeight::<T>::insert(netuid, trimmed_stake_weight);
 
+            // Update incentives/lastupdates for subsubnets
+            for subid in 0..subsubnets_count {
+                let netuid_index = Self::get_subsubnet_storage_index(netuid, subid.into());
+                let incentive = Incentive::<T>::get(netuid_index);
+                let lastupdate = LastUpdate::<T>::get(netuid_index);
+                let mut trimmed_incentive = Vec::with_capacity(trimmed_uids.len());
+                let mut trimmed_lastupdate = Vec::with_capacity(trimmed_uids.len());
+
+                for uid in &trimmed_uids {
+                    trimmed_incentive.push(incentive.get(*uid).cloned().unwrap_or_default());
+                    trimmed_lastupdate.push(lastupdate.get(*uid).cloned().unwrap_or_default());
+                }
+
+                Incentive::<T>::insert(netuid_index, trimmed_incentive);
+                LastUpdate::<T>::insert(netuid_index, trimmed_lastupdate);
+            }
+
             // Create mapping from old uid to new compressed uid
             // This is needed to update connections (weights and bonds) with correct uid references
             let old_to_new_uid: BTreeMap<usize, usize> = trimmed_uids
@@ -299,35 +314,46 @@ impl<T: Config> Pallet<T> {
             // 2. Update all connections to reference the new compressed uids
             // 3. Clear the connections to the trimmed uids
             for (old_uid, new_uid) in &old_to_new_uid {
-                // Swap uid specific storage items to new compressed positions
-                Keys::<T>::swap(netuid, *old_uid as u16, netuid, *new_uid as u16);
-                BlockAtRegistration::<T>::swap(netuid, *old_uid as u16, netuid, *new_uid as u16);
-
-                // Swap to new position and remap all target uids
-                Weights::<T>::swap(netuid, *old_uid as u16, netuid, *new_uid as u16);
-                Weights::<T>::mutate(netuid, *new_uid as u16, |weights| {
-                    weights.retain_mut(|(target_uid, _weight)| {
-                        if let Some(new_target_uid) = old_to_new_uid.get(&(*target_uid as usize)) {
-                            *target_uid = *new_target_uid as u16;
-                            true
-                        } else {
-                            false
-                        }
-                    })
-                });
+                let old_neuron_uid = *old_uid as u16;
+                let new_neuron_uid = *new_uid as u16;
 
-                // Swap to new position and remap all target uids
-                Bonds::<T>::swap(netuid, *old_uid as u16, netuid, *new_uid as u16);
-                Bonds::<T>::mutate(netuid, *new_uid as u16, |bonds| {
-                    bonds.retain_mut(|(target_uid, _bond)| {
-                        if let Some(new_target_uid) = old_to_new_uid.get(&(*target_uid as usize)) {
-                            *target_uid = *new_target_uid as u16;
-                            true
-                        } else {
-                            false
-                        }
-                    })
-                });
+                // Swap uid specific storage items to new compressed positions
+                Keys::<T>::swap(netuid, old_neuron_uid, netuid, new_neuron_uid);
+                BlockAtRegistration::<T>::swap(netuid, old_neuron_uid, netuid, new_neuron_uid);
+
+                for subid in 0..subsubnets_count {
+                    let netuid_index = Self::get_subsubnet_storage_index(netuid, subid.into());
+
+                    // Swap to new position and remap all target uids
+                    Weights::<T>::swap(netuid_index, old_neuron_uid, netuid_index, new_neuron_uid);
+                    Weights::<T>::mutate(netuid_index, new_neuron_uid, |weights| {
+                        weights.retain_mut(|(target_uid, _weight)| {
+                            if let Some(new_target_uid) =
+                                old_to_new_uid.get(&(*target_uid as usize))
+                            {
+                                *target_uid = *new_target_uid as u16;
+                                true
+                            } else {
+                                false
+                            }
+                        })
+                    });
+
+                    // Swap to new position and remap all target uids
+                    Bonds::<T>::swap(netuid_index, old_neuron_uid, netuid_index, new_neuron_uid);
+                    Bonds::<T>::mutate(netuid_index, new_neuron_uid, |bonds| {
+                        bonds.retain_mut(|(target_uid, _bond)| {
+                            if let Some(new_target_uid) =
+                                old_to_new_uid.get(&(*target_uid as usize))
+                            {
+                                *target_uid = *new_target_uid as u16;
+                                true
+                            } else {
+                                false
+                            }
+                        })
+                    });
+                }
             }
 
             // Update the subnet's uid count to reflect the new maximum