Merge pull request #2276 from opentensor/improve-loop-alpha-map

improve loop on alpha map

Author: sam0x17

pallets/subtensor/src/staking/remove_stake.rs

@@ -484,31 +484,37 @@ impl<T: Config> Pallet<T> {
         let mut stakers: Vec<(T::AccountId, T::AccountId, u128)> = Vec::new();
         let mut total_alpha_value_u128: u128 = 0;
 
-        for ((hot, cold, this_netuid), share_u64f64) in Alpha::<T>::iter() {
-            if this_netuid != netuid {
-                continue;
-            }
-
-            keys_to_remove.push((hot.clone(), cold.clone()));
-            if !hotkeys_seen.contains(&hot) {
-                hotkeys_seen.push(hot.clone());
-            }
+        let hotkeys_in_subnet: Vec<T::AccountId> = TotalHotkeyAlpha::<T>::iter()
+            .filter(|(_, this_netuid, _)| *this_netuid == netuid)
+            .map(|(hot, _, _)| hot.clone())
+            .collect::<Vec<_>>();
+
+        for hot in hotkeys_in_subnet.iter() {
+            for ((cold, this_netuid), share_u64f64) in Alpha::<T>::iter_prefix((hot,)) {
+                if this_netuid != netuid {
+                    continue;
+                }
+                keys_to_remove.push((hot.clone(), cold.clone()));
+                if !hotkeys_seen.contains(hot) {
+                    hotkeys_seen.push(hot.clone());
+                }
 
-            // Primary: actual α value via share pool.
-            let pool = Self::get_alpha_share_pool(hot.clone(), netuid);
-            let actual_val_u64 = pool.try_get_value(&cold).unwrap_or(0);
+                // Primary: actual α value via share pool.
+                let pool = Self::get_alpha_share_pool(hot.clone(), netuid);
+                let actual_val_u64 = pool.try_get_value(&cold).unwrap_or(0);
 
-            // Fallback: if pool uninitialized, treat raw Alpha share as value.
-            let val_u64 = if actual_val_u64 == 0 {
-                share_u64f64.saturating_to_num::<u64>()
-            } else {
-                actual_val_u64
-            };
+                // Fallback: if pool uninitialized, treat raw Alpha share as value.
+                let val_u64 = if actual_val_u64 == 0 {
+                    share_u64f64.saturating_to_num::<u64>()
+                } else {
+                    actual_val_u64
+                };
 
-            if val_u64 > 0 {
-                let val_u128 = val_u64 as u128;
-                total_alpha_value_u128 = total_alpha_value_u128.saturating_add(val_u128);
-                stakers.push((hot, cold, val_u128));
+                if val_u64 > 0 {
+                    let val_u128 = val_u64 as u128;
+                    total_alpha_value_u128 = total_alpha_value_u128.saturating_add(val_u128);
+                    stakers.push((hot.clone(), cold, val_u128));
+                }
             }
         }
 

pallets/subtensor/src/tests/networks.rs

@@ -106,6 +106,7 @@ fn dissolve_single_alpha_out_staker_gets_all_tao() {
         let pot: u64 = 99_999;
         SubnetTAO::<Test>::insert(net, TaoCurrency::from(pot));
         SubtensorModule::set_subnet_locked_balance(net, 0.into());
+        TotalHotkeyAlpha::<Test>::insert(s_hot, net, AlphaCurrency::from(5_000u64));
 
         // Cold-key balance before
         let before = SubtensorModule::get_coldkey_balance(&s_cold);
@@ -142,6 +143,9 @@ fn dissolve_two_stakers_pro_rata_distribution() {
         Alpha::<Test>::insert((s1_hot, s1_cold, net), U64F64::from_num(a1));
         Alpha::<Test>::insert((s2_hot, s2_cold, net), U64F64::from_num(a2));
 
+        TotalHotkeyAlpha::<Test>::insert(s1_hot, net, AlphaCurrency::from(a1 as u64));
+        TotalHotkeyAlpha::<Test>::insert(s2_hot, net, AlphaCurrency::from(a2 as u64));
+
         let pot: u64 = 10_000;
         SubnetTAO::<Test>::insert(net, TaoCurrency::from(pot));
         SubtensorModule::set_subnet_locked_balance(net, 5_000.into()); // owner refund path present; emission = 0
@@ -627,6 +631,7 @@ fn dissolve_alpha_out_but_zero_tao_no_rewards() {
         SubnetTAO::<Test>::insert(net, TaoCurrency::from(0)); // zero TAO
         SubtensorModule::set_subnet_locked_balance(net, TaoCurrency::from(0));
         Emission::<Test>::insert(net, Vec::<AlphaCurrency>::new());
+        TotalHotkeyAlpha::<Test>::insert(sh, net, AlphaCurrency::from(1_000u64));
 
         let before = SubtensorModule::get_coldkey_balance(&sc);
         assert_ok!(SubtensorModule::do_dissolve_network(net));
@@ -672,6 +677,9 @@ fn dissolve_rounding_remainder_distribution() {
         SubnetTAO::<Test>::insert(net, TaoCurrency::from(1)); // TAO pot = 1
         SubtensorModule::set_subnet_locked_balance(net, TaoCurrency::from(0));
 
+        TotalHotkeyAlpha::<Test>::insert(s1h, net, AlphaCurrency::from(3u64));
+        TotalHotkeyAlpha::<Test>::insert(s2h, net, AlphaCurrency::from(2u64));
+
         // Cold-key balances before
         let c1_before = SubtensorModule::get_coldkey_balance(&s1c);
         let c2_before = SubtensorModule::get_coldkey_balance(&s2c);

runtime/src/lib.rs

@@ -237,7 +237,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: 361,
+    spec_version: 362,
     impl_version: 1,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 1,