Add subnets to claim_root extrinsic.

Author: shamil-gadelshin

pallets/subtensor/src/benchmarks.rs

@@ -14,6 +14,7 @@ use sp_runtime::{
     BoundedVec, Percent,
     traits::{BlakeTwo256, Hash},
 };
+use sp_std::collections::btree_set::BTreeSet;
 use sp_std::vec;
 use subtensor_runtime_common::{AlphaCurrency, NetUid, TaoCurrency};
 
@@ -1630,7 +1631,7 @@ mod pallet_benchmarks {
         ),);
 
         #[extrinsic_call]
-        _(RawOrigin::Signed(coldkey.clone()));
+        _(RawOrigin::Signed(coldkey.clone()), BTreeSet::from([netuid]));
 
         // Verification
         let new_stake =

pallets/subtensor/src/lib.rs

@@ -58,6 +58,8 @@ pub const ALPHA_MAP_BATCH_SIZE: usize = 30;
 
 pub const MAX_NUM_ROOT_CLAIMS: u64 = 50;
 
+pub const MAX_SUBNET_CLAIMS: usize = 5;
+
 pub const MAX_ROOT_CLAIM_THRESHOLD: u64 = 10_000_000;
 
 #[allow(deprecated)]
@@ -89,6 +91,7 @@ pub mod pallet {
     use sp_core::{ConstU32, H160, H256};
     use sp_runtime::traits::{Dispatchable, TrailingZeroInput};
     use sp_std::collections::btree_map::BTreeMap;
+    use sp_std::collections::btree_set::BTreeSet;
     use sp_std::collections::vec_deque::VecDeque;
     use sp_std::vec;
     use sp_std::vec::Vec;

pallets/subtensor/src/macros/dispatches.rs

@@ -15,6 +15,8 @@ mod dispatches {
     use crate::MAX_CRV3_COMMIT_SIZE_BYTES;
     use crate::MAX_NUM_ROOT_CLAIMS;
     use crate::MAX_ROOT_CLAIM_THRESHOLD;
+    use crate::MAX_SUBNET_CLAIMS;
+
     /// Dispatchable functions allow users to interact with the pallet and invoke state changes.
     /// These functions materialize as "extrinsics", which are often compared to transactions.
     /// Dispatchable functions must be annotated with a weight and must return a DispatchResult.
@@ -2331,12 +2333,21 @@ mod dispatches {
             DispatchClass::Normal,
             Pays::Yes
         ))]
-        pub fn claim_root(origin: OriginFor<T>) -> DispatchResultWithPostInfo {
+        pub fn claim_root(
+            origin: OriginFor<T>,
+            subnets: BTreeSet<NetUid>,
+        ) -> DispatchResultWithPostInfo {
             let coldkey: T::AccountId = ensure_signed(origin)?;
 
+            ensure!(!subnets.is_empty(), Error::<T>::InvalidSubnetNumber);
+            ensure!(
+                subnets.len() <= MAX_SUBNET_CLAIMS,
+                Error::<T>::InvalidSubnetNumber
+            );
+
             Self::maybe_add_coldkey_index(&coldkey);
 
-            let weight = Self::do_root_claim(coldkey);
+            let weight = Self::do_root_claim(coldkey, Some(subnets));
             Ok((Some(weight), Pays::Yes).into())
         }
 

pallets/subtensor/src/macros/errors.rs

@@ -264,5 +264,7 @@ mod errors {
         InvalidNumRootClaim,
         /// Invalid value of root claim threshold
         InvalidRootClaimThreshold,
+        /// Exceeded subnet limit number or zero.
+        InvalidSubnetNumber,
     }
 }

pallets/subtensor/src/staking/claim_root.rs

@@ -210,7 +210,11 @@ impl<T: Config> Pallet<T> {
             .saturating_add(T::DbWeight::get().reads(7_u64))
             .saturating_add(T::DbWeight::get().writes(5_u64))
     }
-    pub fn root_claim_all(hotkey: &T::AccountId, coldkey: &T::AccountId) -> Weight {
+    pub fn root_claim_all(
+        hotkey: &T::AccountId,
+        coldkey: &T::AccountId,
+        subnets: Option<BTreeSet<NetUid>>,
+    ) -> Weight {
         let mut weight = Weight::default();
 
         let root_claim_type = RootClaimType::<T>::get(coldkey);
@@ -220,10 +224,19 @@ impl<T: Config> Pallet<T> {
         let root_claimable = RootClaimable::<T>::get(hotkey);
         weight.saturating_accrue(T::DbWeight::get().reads(1));
 
-        root_claimable.iter().for_each(|(netuid, _)| {
+        for (netuid, _) in root_claimable.iter() {
+            let skip = subnets
+                .as_ref()
+                .map(|subnets| !subnets.contains(netuid))
+                .unwrap_or(false);
+
+            if skip {
+                continue;
+            }
+
             Self::root_claim_on_subnet(hotkey, coldkey, *netuid, root_claim_type.clone(), false);
             weight.saturating_accrue(Self::root_claim_on_subnet_weight(root_claim_type.clone()));
-        });
+        }
 
         weight
     }
@@ -278,15 +291,15 @@ impl<T: Config> Pallet<T> {
         }
     }
 
-    pub fn do_root_claim(coldkey: T::AccountId) -> Weight {
+    pub fn do_root_claim(coldkey: T::AccountId, subnets: Option<BTreeSet<NetUid>>) -> Weight {
         let mut weight = Weight::default();
 
         let hotkeys = StakingHotkeys::<T>::get(&coldkey);
         weight.saturating_accrue(T::DbWeight::get().reads(1));
 
         hotkeys.iter().for_each(|hotkey| {
             weight.saturating_accrue(T::DbWeight::get().reads(1));
-            weight.saturating_accrue(Self::root_claim_all(hotkey, &coldkey));
+            weight.saturating_accrue(Self::root_claim_all(hotkey, &coldkey, subnets.clone()));
         });
 
         Self::deposit_event(Event::RootClaimed { coldkey });
@@ -321,7 +334,7 @@ impl<T: Config> Pallet<T> {
         for i in coldkeys_to_claim.iter() {
             weight.saturating_accrue(T::DbWeight::get().reads(1));
             if let Ok(coldkey) = StakingColdkeysByIndex::<T>::try_get(i) {
-                weight.saturating_accrue(Self::do_root_claim(coldkey.clone()));
+                weight.saturating_accrue(Self::do_root_claim(coldkey.clone(), None));
             }
 
             continue;