Network dissolution optimization for root claim

pallets/subtensor/src/lib.rs

@@ -1962,9 +1962,9 @@ pub mod pallet {
     pub type RootClaimed<T: Config> = StorageNMap<
         _,
         (
+            NMapKey<Identity, NetUid>,               // subnet
             NMapKey<Blake2_128Concat, T::AccountId>, // hot
             NMapKey<Blake2_128Concat, T::AccountId>, // cold
-            NMapKey<Identity, NetUid>,               // subnet
         ),
         u128,
         ValueQuery,

pallets/subtensor/src/macros/dispatches.rs

@@ -712,8 +712,8 @@ mod dispatches {
         ///
         #[pallet::call_index(2)]
         #[pallet::weight((Weight::from_parts(340_800_000, 0)
-		.saturating_add(T::DbWeight::get().reads(24_u64))
-		.saturating_add(T::DbWeight::get().writes(15)), DispatchClass::Normal, Pays::Yes))]
+		.saturating_add(T::DbWeight::get().reads(25_u64))
+		.saturating_add(T::DbWeight::get().writes(16_u64)), DispatchClass::Normal, Pays::Yes))]
         pub fn add_stake(
             origin: OriginFor<T>,
             hotkey: T::AccountId,
@@ -1587,8 +1587,8 @@ mod dispatches {
         ///     - Thrown if key has hit transaction rate limit
         #[pallet::call_index(84)]
         #[pallet::weight((Weight::from_parts(358_500_000, 0)
-        .saturating_add(T::DbWeight::get().reads(39_u64))
-        .saturating_add(T::DbWeight::get().writes(24_u64)), DispatchClass::Normal, Pays::Yes))]
+        .saturating_add(T::DbWeight::get().reads(41_u64))
+        .saturating_add(T::DbWeight::get().writes(26_u64)), DispatchClass::Normal, Pays::Yes))]
         pub fn unstake_all_alpha(origin: OriginFor<T>, hotkey: T::AccountId) -> DispatchResult {
             Self::do_unstake_all_alpha(origin, hotkey)
         }
@@ -1701,8 +1701,8 @@ mod dispatches {
         #[pallet::call_index(87)]
         #[pallet::weight((
             Weight::from_parts(351_300_000, 0)
-            .saturating_add(T::DbWeight::get().reads(35_u64))
-            .saturating_add(T::DbWeight::get().writes(22_u64)),
+            .saturating_add(T::DbWeight::get().reads(37_u64))
+            .saturating_add(T::DbWeight::get().writes(24_u64)),
             DispatchClass::Normal,
             Pays::Yes
         ))]
@@ -1766,8 +1766,8 @@ mod dispatches {
         ///
         #[pallet::call_index(88)]
         #[pallet::weight((Weight::from_parts(402_900_000, 0)
-		.saturating_add(T::DbWeight::get().reads(24_u64))
-		.saturating_add(T::DbWeight::get().writes(15)), DispatchClass::Normal, Pays::Yes))]
+		.saturating_add(T::DbWeight::get().reads(25_u64))
+		.saturating_add(T::DbWeight::get().writes(16_u64)), DispatchClass::Normal, Pays::Yes))]
         pub fn add_stake_limit(
             origin: OriginFor<T>,
             hotkey: T::AccountId,
@@ -1830,8 +1830,8 @@ mod dispatches {
         ///
         #[pallet::call_index(89)]
         #[pallet::weight((Weight::from_parts(377_400_000, 0)
-		.saturating_add(T::DbWeight::get().reads(28_u64))
-		.saturating_add(T::DbWeight::get().writes(14)), DispatchClass::Normal, Pays::Yes))]
+		.saturating_add(T::DbWeight::get().reads(29_u64))
+		.saturating_add(T::DbWeight::get().writes(15_u64)), DispatchClass::Normal, Pays::Yes))]
         pub fn remove_stake_limit(
             origin: OriginFor<T>,
             hotkey: T::AccountId,
@@ -1874,8 +1874,8 @@ mod dispatches {
         #[pallet::call_index(90)]
         #[pallet::weight((
             Weight::from_parts(411_500_000, 0)
-            .saturating_add(T::DbWeight::get().reads(35_u64))
-            .saturating_add(T::DbWeight::get().writes(22_u64)),
+            .saturating_add(T::DbWeight::get().reads(37_u64))
+            .saturating_add(T::DbWeight::get().writes(24_u64)),
             DispatchClass::Normal,
             Pays::Yes
         ))]
@@ -2052,8 +2052,8 @@ mod dispatches {
         /// Without limit_price it remove all the stake similar to `remove_stake` extrinsic
         #[pallet::call_index(103)]
         #[pallet::weight((Weight::from_parts(395_300_000, 10142)
-			.saturating_add(T::DbWeight::get().reads(28_u64))
-			.saturating_add(T::DbWeight::get().writes(14_u64)), DispatchClass::Normal, Pays::Yes))]
+			.saturating_add(T::DbWeight::get().reads(29_u64))
+			.saturating_add(T::DbWeight::get().writes(15_u64)), DispatchClass::Normal, Pays::Yes))]
         pub fn remove_stake_full_limit(
             origin: T::RuntimeOrigin,
             hotkey: T::AccountId,

pallets/subtensor/src/staking/claim_root.rs

@@ -2,7 +2,7 @@ use super::*;
 use frame_support::weights::Weight;
 use sp_core::Get;
 use sp_std::collections::btree_set::BTreeSet;
-use substrate_fixed::types::{I96F32, U64F64};
+use substrate_fixed::types::I96F32;
 use subtensor_swap_interface::SwapHandler;
 
 impl<T: Config> Pallet<T> {
@@ -99,7 +99,7 @@ impl<T: Config> Pallet<T> {
 
         // Attain the root claimed to avoid overclaiming.
         let root_claimed: I96F32 =
-            I96F32::saturating_from_num(RootClaimed::<T>::get((hotkey, coldkey, netuid)));
+            I96F32::saturating_from_num(RootClaimed::<T>::get((netuid, hotkey, coldkey)));
 
         // Subtract the already claimed alpha.
         let owed: I96F32 = claimable.saturating_sub(root_claimed);
@@ -200,7 +200,7 @@ impl<T: Config> Pallet<T> {
         };
 
         // Increase root claimed by owed amount.
-        RootClaimed::<T>::mutate((hotkey, coldkey, netuid), |root_claimed| {
+        RootClaimed::<T>::mutate((netuid, hotkey, coldkey), |root_claimed| {
             *root_claimed = root_claimed.saturating_add(owed_u64.into());
         });
     }
@@ -250,7 +250,7 @@ impl<T: Config> Pallet<T> {
         let root_claimable = RootClaimable::<T>::get(hotkey);
         for (netuid, claimable_rate) in root_claimable.iter() {
             // Get current staker root claimed value.
-            let root_claimed: u128 = RootClaimed::<T>::get((hotkey, coldkey, netuid));
+            let root_claimed: u128 = RootClaimed::<T>::get((netuid, hotkey, coldkey));
 
             // Increase root claimed based on the claimable rate.
             let new_root_claimed = root_claimed.saturating_add(
@@ -260,7 +260,7 @@ impl<T: Config> Pallet<T> {
             );
 
             // Set the new root claimed value.
-            RootClaimed::<T>::insert((hotkey, coldkey, netuid), new_root_claimed);
+            RootClaimed::<T>::insert((netuid, hotkey, coldkey), new_root_claimed);
         }
     }
 
@@ -277,7 +277,7 @@ impl<T: Config> Pallet<T> {
             }
 
             // Get current staker root claimed value.
-            let root_claimed: u128 = RootClaimed::<T>::get((hotkey, coldkey, netuid));
+            let root_claimed: u128 = RootClaimed::<T>::get((netuid, hotkey, coldkey));
 
             // Decrease root claimed based on the claimable rate.
             let new_root_claimed = root_claimed.saturating_sub(
@@ -287,7 +287,7 @@ impl<T: Config> Pallet<T> {
             );
 
             // Set the new root_claimed value.
-            RootClaimed::<T>::insert((hotkey, coldkey, netuid), new_root_claimed);
+            RootClaimed::<T>::insert((netuid, hotkey, coldkey), new_root_claimed);
         }
     }
 
@@ -359,10 +359,10 @@ impl<T: Config> Pallet<T> {
         old_coldkey: &T::AccountId,
         new_coldkey: &T::AccountId,
     ) {
-        let old_root_claimed = RootClaimed::<T>::get((old_hotkey, old_coldkey, netuid));
-        RootClaimed::<T>::remove((old_hotkey, old_coldkey, netuid));
+        let old_root_claimed = RootClaimed::<T>::get((netuid, old_hotkey, old_coldkey));
+        RootClaimed::<T>::remove((netuid, old_hotkey, old_coldkey));
 
-        RootClaimed::<T>::mutate((new_hotkey, new_coldkey, netuid), |new_root_claimed| {
+        RootClaimed::<T>::mutate((netuid, new_hotkey, new_coldkey), |new_root_claimed| {
             *new_root_claimed = old_root_claimed.saturating_add(*new_root_claimed);
         });
     }
@@ -386,35 +386,14 @@ impl<T: Config> Pallet<T> {
 
     /// Claim all root dividends for subnet and remove all associated data.
     pub fn finalize_all_subnet_root_dividends(netuid: NetUid) {
-        let mut hotkeys_to_clear = BTreeSet::new();
-        for (hotkey, root_claimable) in RootClaimable::<T>::iter() {
-            if root_claimable.contains_key(&netuid) {
-                let alpha_values: Vec<((T::AccountId, NetUid), U64F64)> =
-                    Alpha::<T>::iter_prefix((&hotkey,)).collect();
-
-                for ((coldkey, alpha_netuid), _) in alpha_values.into_iter() {
-                    if alpha_netuid == NetUid::ROOT {
-                        Self::root_claim_on_subnet(
-                            &hotkey,
-                            &coldkey,
-                            netuid,
-                            RootClaimTypeEnum::Swap,
-                            true,
-                        );
-
-                        RootClaimed::<T>::remove((hotkey.clone(), coldkey, netuid));
-                        if !hotkeys_to_clear.contains(&hotkey) {
-                            hotkeys_to_clear.insert(hotkey.clone());
-                        }
-                    }
-                }
-            }
-        }
+        let hotkeys = RootClaimable::<T>::iter_keys().collect::<Vec<_>>();
 
-        for hotkey in hotkeys_to_clear.into_iter() {
-            RootClaimable::<T>::mutate(&hotkey, |claimable| {
+        for hotkey in hotkeys.iter() {
+            RootClaimable::<T>::mutate(hotkey, |claimable| {
                 claimable.remove(&netuid);
             });
         }
+
+        let _ = RootClaimed::<T>::clear_prefix((netuid,), u32::MAX, None);
     }
 }

pallets/subtensor/src/tests/claim_root.rs

@@ -135,7 +135,7 @@ fn test_claim_root_with_drain_emissions() {
 
         // Check root claimed value saved
 
-        let claimed = RootClaimed::<Test>::get((&hotkey, &coldkey, netuid));
+        let claimed = RootClaimed::<Test>::get((netuid, &hotkey, &coldkey));
         assert_eq!(u128::from(new_stake), claimed);
 
         // Distribute pending root alpha (round 2)
@@ -180,7 +180,7 @@ fn test_claim_root_with_drain_emissions() {
 
         // Check root claimed value saved (round 2)
 
-        let claimed = RootClaimed::<Test>::get((&hotkey, &coldkey, netuid));
+        let claimed = RootClaimed::<Test>::get((netuid, &hotkey, &coldkey));
         assert_eq!(u128::from(u64::from(new_stake2)), claimed);
     });
 }
@@ -1125,11 +1125,11 @@ fn test_claim_root_with_swap_coldkey() {
 
         assert_eq!(
             u128::from(new_stake),
-            RootClaimed::<Test>::get((&hotkey, &coldkey, netuid))
+            RootClaimed::<Test>::get((netuid, &hotkey, &coldkey))
         );
         assert_eq!(
             0u128,
-            RootClaimed::<Test>::get((&hotkey, &new_coldkey, netuid))
+            RootClaimed::<Test>::get((netuid, &hotkey, &new_coldkey))
         );
 
         // Swap coldkey
@@ -1143,10 +1143,10 @@ fn test_claim_root_with_swap_coldkey() {
 
         // Check swapped keys claimed values
 
-        assert_eq!(0u128, RootClaimed::<Test>::get((&hotkey, &coldkey, netuid)));
+        assert_eq!(0u128, RootClaimed::<Test>::get((netuid, &hotkey, &coldkey)));
         assert_eq!(
             u128::from(new_stake),
-            RootClaimed::<Test>::get((&hotkey, &new_coldkey, netuid))
+            RootClaimed::<Test>::get((netuid, &hotkey, &new_coldkey,))
         );
     });
 }
@@ -1215,11 +1215,11 @@ fn test_claim_root_with_swap_hotkey() {
 
         assert_eq!(
             u128::from(new_stake),
-            RootClaimed::<Test>::get((&hotkey, &coldkey, netuid))
+            RootClaimed::<Test>::get((netuid, &hotkey, &coldkey,))
         );
         assert_eq!(
             0u128,
-            RootClaimed::<Test>::get((&new_hotkey, &coldkey, netuid))
+            RootClaimed::<Test>::get((netuid, &new_hotkey, &coldkey,))
         );
 
         let _old_claimable = *RootClaimable::<Test>::get(hotkey)
@@ -1239,10 +1239,13 @@ fn test_claim_root_with_swap_hotkey() {
 
         // Check swapped keys claimed values
 
-        assert_eq!(0u128, RootClaimed::<Test>::get((&hotkey, &coldkey, netuid)));
+        assert_eq!(
+            0u128,
+            RootClaimed::<Test>::get((netuid, &hotkey, &coldkey,))
+        );
         assert_eq!(
             u128::from(new_stake),
-            RootClaimed::<Test>::get((&new_hotkey, &coldkey, netuid))
+            RootClaimed::<Test>::get((netuid, &new_hotkey, &coldkey,))
         );
 
         assert!(!RootClaimable::<Test>::get(hotkey).contains_key(&netuid));
@@ -1281,7 +1284,6 @@ fn test_claim_root_on_network_deregistration() {
             NetUid::ROOT,
             root_stake.into(),
         );
-        let root_stake_rate = 0.1f64;
         SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
             &hotkey,
             &other_coldkey,
@@ -1307,33 +1309,23 @@ fn test_claim_root_on_network_deregistration() {
             AlphaCurrency::ZERO,
         );
 
-        // Claim root via network deregistration
-
-        assert_ok!(SubtensorModule::do_dissolve_network(netuid));
+        assert_ok!(SubtensorModule::claim_root(
+            RuntimeOrigin::signed(coldkey),
+            BTreeSet::from([netuid])
+        ));
 
-        // Check new stake
-        let validator_take_percent = 0.18f64;
+        assert!(RootClaimable::<Test>::get(hotkey).contains_key(&netuid));
 
-        let new_stake: u64 = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
-            &hotkey,
-            &coldkey,
-            NetUid::ROOT,
-        )
-        .into();
+        assert!(RootClaimed::<Test>::contains_key((
+            netuid, &hotkey, &coldkey,
+        )));
 
-        let estimated_stake_increment = (pending_root_alpha as f64)
-            * (1f64 - validator_take_percent)
-            * current_price
-            * root_stake_rate;
+        // Claim root via network deregistration
 
-        assert_abs_diff_eq!(
-            new_stake,
-            root_stake + estimated_stake_increment as u64,
-            epsilon = 10000u64,
-        );
+        assert_ok!(SubtensorModule::do_dissolve_network(netuid));
 
         assert!(!RootClaimed::<Test>::contains_key((
-            &hotkey, &coldkey, netuid
+            netuid, &hotkey, &coldkey,
         )));
         assert!(!RootClaimable::<Test>::get(hotkey).contains_key(&netuid));
     });
@@ -1505,7 +1497,7 @@ fn test_claim_root_with_unrelated_subnets() {
 
         // Check root claimed value saved
 
-        let claimed = RootClaimed::<Test>::get((&hotkey, &coldkey, netuid));
+        let claimed = RootClaimed::<Test>::get((netuid, &hotkey, &coldkey));
         assert_eq!(u128::from(new_stake), claimed);
     });
 }