Fix block indices to claim

Author: shamil-gadelshin

pallets/subtensor/src/lib.rs

@@ -1865,7 +1865,11 @@ pub mod pallet {
     >;
 
     #[pallet::storage] // --- MAP ( u64 ) --> coldkey | Maps coldkeys that have stake to an index
-    pub type StakingColdkeys<T: Config> = StorageMap<_, Identity, u64, T::AccountId, OptionQuery>;
+    pub type StakingColdkeysByIndex<T: Config> =
+        StorageMap<_, Identity, u64, T::AccountId, OptionQuery>;
+
+    #[pallet::storage] // --- MAP ( coldkey ) --> index | Maps index that have stake to a coldkey
+    pub type StakingColdkeys<T: Config> = StorageMap<_, Identity, T::AccountId, u64, OptionQuery>;
 
     #[pallet::storage] // --- Value --> num_staking_coldkeys
     pub type NumStakingColdkeys<T: Config> = StorageValue<_, u64, ValueQuery, DefaultZeroU64<T>>;

pallets/subtensor/src/staking/claim_root.rs

@@ -1,13 +1,14 @@
 use super::*;
 use frame_support::weights::Weight;
 use sp_core::Get;
+use sp_std::collections::btree_set::BTreeSet;
 use substrate_fixed::types::I96F32;
 use subtensor_swap_interface::SwapHandler;
 
 impl<T: Config> Pallet<T> {
     pub fn block_hash_to_indices(block_hash: T::Hash, k: u64, n: u64) -> Vec<u64> {
         let block_hash_bytes = block_hash.as_ref();
-        let mut indices: Vec<u64> = Vec::new();
+        let mut indices: BTreeSet<u64> = BTreeSet::new();
         // k < n
         let start_index: u64 = u64::from_be_bytes(
             block_hash_bytes
@@ -30,10 +31,10 @@ impl<T: Config> Pallet<T> {
                 .saturating_add(idx_step)
                 .checked_rem(n)
                 .unwrap_or(0);
-            indices.push(idx);
+            indices.insert(idx);
             last_idx = idx;
         }
-        indices
+        indices.into_iter().collect()
     }
 
     pub fn increase_root_claimable_for_hotkey_and_subnet(
@@ -277,6 +278,15 @@ impl<T: Config> Pallet<T> {
         Weight::default()
     }
 
+    pub fn maybe_add_coldkey_index(coldkey: &T::AccountId) {
+        if !StakingColdkeys::<T>::contains_key(coldkey) {
+            let n = NumStakingColdkeys::<T>::get();
+            StakingColdkeysByIndex::<T>::insert(n, coldkey.clone());
+            StakingColdkeys::<T>::insert(coldkey.clone(), n);
+            NumStakingColdkeys::<T>::mutate(|n| *n = n.saturating_add(1));
+        }
+    }
+
     pub fn run_auto_claim_root_divs(last_block_hash: T::Hash) -> Weight {
         let mut weight: Weight = Weight::default();
 
@@ -289,7 +299,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) = StakingColdkeys::<T>::try_get(i) {
+            if let Ok(coldkey) = StakingColdkeysByIndex::<T>::try_get(i) {
                 weight.saturating_accrue(Self::do_root_claim(coldkey.clone()));
             }
 

pallets/subtensor/src/staking/stake_utils.rs

@@ -830,6 +830,7 @@ impl<T: Config> Pallet<T> {
             // Adjust root claimed for this hotkey and coldkey.
             let alpha = swap_result.amount_paid_out.into();
             Self::add_stake_adjust_root_claimed_for_hotkey_and_coldkey(hotkey, coldkey, alpha);
+            Self::maybe_add_coldkey_index(coldkey);
         }
 
         // Deposit and log the staking event.

pallets/subtensor/src/tests/claim_root.rs

@@ -1,12 +1,14 @@
-use crate::tests::mock::{RuntimeOrigin, SubtensorModule, Test, add_dynamic_network, new_test_ext};
+use crate::tests::mock::{
+    RuntimeOrigin, SubtensorModule, Test, add_dynamic_network, new_test_ext, run_to_block,
+};
 use crate::{
     NetworksAdded, RootClaimable, SubnetAlphaIn, SubnetMechanism, SubnetTAO, SubtokenEnabled,
     Tempo, pallet,
 };
 use crate::{RootClaimType, RootClaimTypeEnum, RootClaimed};
 use approx::assert_abs_diff_eq;
 use frame_support::assert_ok;
-use sp_core::U256;
+use sp_core::{H256, U256};
 use substrate_fixed::types::{I96F32, U96F32};
 use subtensor_runtime_common::{AlphaCurrency, Currency, NetUid, TaoCurrency};
 use subtensor_swap_interface::SwapHandler;
@@ -771,3 +773,94 @@ fn test_claim_root_with_run_coinbase() {
         assert!(new_stake > 0);
     });
 }
+
+#[test]
+fn test_claim_root_bloch_hash_indices() {
+    new_test_ext(1).execute_with(|| {
+        let k = 15u64;
+        let n = 15000u64;
+
+        // 1
+        let hash = sp_core::keccak_256(b"some");
+        let mut indices = SubtensorModule::block_hash_to_indices(H256(hash), k, n);
+        indices.sort();
+
+        assert!(indices.len() <= k as usize);
+        assert!(!indices.iter().any(|i| *i >= n));
+        // precomputed values
+        let expected_result = vec![
+            265, 630, 1286, 1558, 4496, 4861, 5517, 5789, 6803, 8096, 9092, 11034, 11399, 12055,
+            12327,
+        ];
+        assert_eq!(indices, expected_result);
+
+        // 2
+        let hash = sp_core::keccak_256(b"some2");
+        let mut indices = SubtensorModule::block_hash_to_indices(H256(hash), k, n);
+        indices.sort();
+
+        assert!(indices.len() <= k as usize);
+        assert!(!indices.iter().any(|i| *i >= n));
+        // precomputed values
+        let expected_result = vec![
+            61, 246, 1440, 2855, 3521, 5236, 6130, 6615, 8511, 9405, 9890, 11786, 11971, 13165,
+            14580,
+        ];
+        assert_eq!(indices, expected_result);
+    });
+}
+
+#[test]
+fn test_claim_root_with_block_emissions() {
+    new_test_ext(0).execute_with(|| {
+        let owner_coldkey = U256::from(1001);
+        let hotkey = U256::from(1002);
+        let coldkey = U256::from(1003);
+        let netuid = add_dynamic_network(&hotkey, &owner_coldkey);
+
+        Tempo::<Test>::insert(netuid, 1);
+        SubtensorModule::set_tao_weight(u64::MAX); // Set TAO weight to 1.0
+
+        let root_stake = 200_000_000u64;
+        SubnetTAO::<Test>::insert(NetUid::ROOT, TaoCurrency::from(root_stake));
+
+        SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey,
+            &coldkey,
+            NetUid::ROOT,
+            root_stake.into(),
+        );
+        SubtensorModule::maybe_add_coldkey_index(&coldkey);
+
+        let initial_total_hotkey_alpha = 10_000_000u64;
+        SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey,
+            &owner_coldkey,
+            netuid,
+            initial_total_hotkey_alpha.into(),
+        );
+
+        assert_ok!(SubtensorModule::set_root_claim_type(
+            RuntimeOrigin::signed(coldkey),
+            RootClaimTypeEnum::Keep
+        ),);
+        assert_eq!(RootClaimType::<Test>::get(coldkey), RootClaimTypeEnum::Keep);
+
+        // Distribute pending root alpha
+
+        let initial_stake: u64 =
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(&hotkey, &coldkey, netuid)
+                .into();
+        assert_eq!(initial_stake, 0u64);
+
+        run_to_block(2);
+
+        // Check stake after block emissions
+
+        let new_stake: u64 =
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(&hotkey, &coldkey, netuid)
+                .into();
+
+        assert!(new_stake > 0);
+    });
+}