feat: get total delegated stake

Author: Samuel Dare

pallets/subtensor/src/delegate_info.rs

@@ -40,21 +40,19 @@ impl<T: Config> Pallet<T> {
         let mut emissions_per_day: U64F64 = U64F64::from_num(0);
 
         for netuid in registrations.iter() {
-            let _uid = Self::get_uid_for_net_and_hotkey(*netuid, &delegate.clone());
-            if _uid.is_err() {
-                continue; // this should never happen
-            } else {
-                let uid = _uid.expect("Delegate's UID should be ok");
+            if let Ok(uid) = Self::get_uid_for_net_and_hotkey(*netuid, &delegate.clone()) {
                 let validator_permit = Self::get_validator_permit_for_uid(*netuid, uid);
                 if validator_permit {
                     validator_permits.push((*netuid).into());
                 }
 
                 let emission: U64F64 = Self::get_emission_for_uid(*netuid, uid).into();
                 let tempo: U64F64 = Self::get_tempo(*netuid).into();
-                let epochs_per_day: U64F64 = U64F64::from_num(7200).saturating_div(tempo);
-                emissions_per_day =
-                    emissions_per_day.saturating_add(emission.saturating_mul(epochs_per_day));
+                if tempo > U64F64::from_num(0) {
+                    let epochs_per_day: U64F64 = U64F64::from_num(7200).saturating_div(tempo);
+                    emissions_per_day =
+                        emissions_per_day.saturating_add(emission.saturating_mul(epochs_per_day));
+                }
             }
         }
 
@@ -63,15 +61,15 @@ impl<T: Config> Pallet<T> {
 
         let total_stake: U64F64 = Self::get_total_stake_for_hotkey(&delegate.clone()).into();
 
-        let mut return_per_1000: U64F64 = U64F64::from_num(0);
-
-        if total_stake > U64F64::from_num(0) {
-            return_per_1000 = emissions_per_day
+        let return_per_1000: U64F64 = if total_stake > U64F64::from_num(0) {
+            emissions_per_day
                 .saturating_mul(U64F64::from_num(0.82))
-                .saturating_div(total_stake.saturating_div(U64F64::from_num(1000)));
-        }
+                .saturating_div(total_stake.saturating_div(U64F64::from_num(1000)))
+        } else {
+            U64F64::from_num(0)
+        };
 
-        return DelegateInfo {
+        DelegateInfo {
             delegate_ss58: delegate.clone(),
             take,
             nominators,
@@ -80,7 +78,7 @@ impl<T: Config> Pallet<T> {
             validator_permits,
             return_per_1000: U64F64::to_num::<u64>(return_per_1000).into(),
             total_daily_return: U64F64::to_num::<u64>(emissions_per_day).into(),
-        };
+        }
     }
 
     pub fn get_delegate(delegate_account_vec: Vec<u8>) -> Option<DelegateInfo<T>> {
@@ -132,4 +130,40 @@ impl<T: Config> Pallet<T> {
 
         delegates
     }
+
+    /// Returns the total delegated stake for a given delegate, excluding the stake from the delegate's owner.
+    ///
+    /// # Arguments
+    ///
+    /// * `delegate` - A reference to the account ID of the delegate.
+    ///
+    /// # Returns
+    ///
+    /// * `u64` - The total amount of stake delegated to the delegate, excluding the owner's stake.
+    ///
+    ///
+    /// # Notes
+    ///
+    /// This function retrieves the delegate's information and calculates the total stake from all nominators,
+    /// excluding the stake from the delegate's owner.
+    pub fn get_total_delegated_stake(delegate: &T::AccountId) -> u64 {
+        if !<Delegates<T>>::contains_key(delegate) {
+            return 0;
+        }
+
+        // Retrieve the delegate's information
+        let delegate_info: DelegateInfo<T> =
+            Self::get_delegate_by_existing_account(delegate.clone());
+
+        // Retrieve the owner's account ID for the given delegate
+        let owner: T::AccountId = Self::get_owning_coldkey_for_hotkey(delegate);
+
+        // Calculate the total stake from all nominators, excluding the owner's stake
+        delegate_info
+            .nominators
+            .iter()
+            .filter(|(nominator, _)| nominator != &owner) // Exclude the owner's stake
+            .map(|(_, stake)| stake.0 as u64) // Map the stake to u64
+            .sum() // Sum the stakes
+    }
 }

pallets/subtensor/tests/staking.rs

@@ -13,6 +13,7 @@ use frame_support::sp_runtime::DispatchError;
 use mock::*;
 use pallet_balances::Call as BalancesCall;
 use pallet_subtensor::*;
+use serde::de;
 use sp_core::{H256, U256};
 use sp_runtime::traits::SignedExtension;
 
@@ -4195,3 +4196,198 @@ fn test_comprehensive_coldkey_swap_scenarios() {
         assert_eq!(SubtensorModule::get_coldkey_balance(&regular_user), 0);
     });
 }
+
+#[test]
+fn test_get_total_delegated_stake_after_unstaking() {
+    new_test_ext(1).execute_with(|| {
+        let delegate = U256::from(1);
+        let coldkey = U256::from(2);
+        let delegator = U256::from(3);
+        let initial_stake = 2000;
+        let unstake_amount = 500;
+        let netuid = 1u16;
+        let existential_deposit = 1; // Account for the existential deposit
+
+        add_network(netuid, 0, 0);
+
+        register_ok_neuron(netuid, delegate, coldkey, 0);
+
+        // Make the delegate a delegate
+        assert_ok!(SubtensorModule::become_delegate(
+            RuntimeOrigin::signed(coldkey),
+            delegate
+        ));
+
+        // Add balance to delegator
+        SubtensorModule::add_balance_to_coldkey_account(&delegator, initial_stake);
+
+        // Delegate stake
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(delegator),
+            delegate,
+            initial_stake
+        ));
+
+        // Unstake part of the delegation
+        assert_ok!(SubtensorModule::remove_stake(
+            RuntimeOrigin::signed(delegator),
+            delegate,
+            unstake_amount
+        ));
+
+        // Calculate the expected delegated stake
+        let expected_delegated_stake = initial_stake - unstake_amount - existential_deposit;
+
+        // Check the total delegated stake after unstaking
+        assert_eq!(
+            SubtensorModule::get_total_delegated_stake(&delegate),
+            expected_delegated_stake
+        );
+    });
+}
+
+#[test]
+fn test_get_total_delegated_stake_no_delegations() {
+    new_test_ext(1).execute_with(|| {
+        let delegate = U256::from(1);
+        let coldkey = U256::from(2);
+        let netuid = 1u16;
+
+        add_network(netuid, 0, 0);
+        register_ok_neuron(netuid, delegate, coldkey, 0);
+
+        // Make the delegate a delegate
+        assert_ok!(SubtensorModule::become_delegate(
+            RuntimeOrigin::signed(coldkey),
+            delegate
+        ));
+
+        // Check that there's no delegated stake
+        assert_eq!(SubtensorModule::get_total_delegated_stake(&delegate), 0);
+    });
+}
+
+#[test]
+fn test_get_total_delegated_stake_single_delegator() {
+    new_test_ext(1).execute_with(|| {
+        let delegate = U256::from(1);
+        let coldkey = U256::from(2);
+        let delegator = U256::from(3);
+        let stake_amount = 1000;
+        let netuid = 1u16;
+
+        add_network(netuid, 0, 0);
+        register_ok_neuron(netuid, delegate, coldkey, 0);
+
+        // Make the delegate a delegate
+        assert_ok!(SubtensorModule::become_delegate(
+            RuntimeOrigin::signed(coldkey),
+            delegate
+        ));
+
+        // Add balance to delegator
+        SubtensorModule::add_balance_to_coldkey_account(&delegator, stake_amount);
+
+        // Delegate stake
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(delegator),
+            delegate,
+            stake_amount
+        ));
+
+        // Check the total delegated stake
+        assert_eq!(
+            SubtensorModule::get_total_delegated_stake(&delegate),
+            stake_amount - 1 // Subtract 1 for existential deposit
+        );
+    });
+}
+
+#[test]
+fn test_get_total_delegated_stake_multiple_delegators() {
+    new_test_ext(1).execute_with(|| {
+        let delegate = U256::from(1);
+        let coldkey = U256::from(2);
+        let delegator1 = U256::from(3);
+        let delegator2 = U256::from(4);
+        let stake_amount1 = 1000;
+        let stake_amount2 = 2000;
+        let netuid = 1u16;
+
+        add_network(netuid, 0, 0);
+        register_ok_neuron(netuid, delegate, coldkey, 0);
+
+        // Make the delegate a delegate
+        assert_ok!(SubtensorModule::become_delegate(
+            RuntimeOrigin::signed(coldkey),
+            delegate
+        ));
+
+        // Add balance to delegators
+        SubtensorModule::add_balance_to_coldkey_account(&delegator1, stake_amount1);
+        SubtensorModule::add_balance_to_coldkey_account(&delegator2, stake_amount2);
+
+        // Delegate stakes
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(delegator1),
+            delegate,
+            stake_amount1
+        ));
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(delegator2),
+            delegate,
+            stake_amount2
+        ));
+
+        // Check the total delegated stake
+        assert_eq!(
+            SubtensorModule::get_total_delegated_stake(&delegate),
+            stake_amount1 + stake_amount2 - 2 // Subtract 2 for existential deposits
+        );
+    });
+}
+
+#[test]
+fn test_get_total_delegated_stake_exclude_owner_stake() {
+    new_test_ext(1).execute_with(|| {
+        let delegate = U256::from(1);
+        let coldkey = U256::from(2);
+        let delegator = U256::from(3);
+        let owner_stake = 5000;
+        let delegator_stake = 1000;
+        let netuid = 1u16;
+
+        add_network(netuid, 0, 0);
+        register_ok_neuron(netuid, delegate, coldkey, 0);
+
+        // Make the delegate a delegate
+        assert_ok!(SubtensorModule::become_delegate(
+            RuntimeOrigin::signed(coldkey),
+            delegate
+        ));
+
+        // Add balance to owner and delegator
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey, owner_stake);
+        SubtensorModule::add_balance_to_coldkey_account(&delegator, delegator_stake);
+
+        // Owner adds stake
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(coldkey),
+            delegate,
+            owner_stake
+        ));
+
+        // Delegator adds stake
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(delegator),
+            delegate,
+            delegator_stake
+        ));
+
+        // Check the total delegated stake (should exclude owner's stake)
+        assert_eq!(
+            SubtensorModule::get_total_delegated_stake(&delegate),
+            delegator_stake - 1 // Subtract 1 for existential deposit
+        );
+    });
+}