Adding getTotalColdkeyStakeand getTotalHotkeyStake to the staking precompile

pallets/subtensor/src/staking/stake_utils.rs

@@ -379,6 +379,69 @@ impl<T: Config> Pallet<T> {
         TotalHotkeyAlpha::<T>::get(hotkey, netuid)
     }
 
+    /// Retrieves the total stake (alpha) for a given coldkey on a specific sunbet.
+    ///
+    /// This function performs the following steps:
+    /// 1. Retrieves the list of hotkeys associated with the coldkey on the subnet.
+    /// 2. Iterates through the list of hotkeys and retrieves the alpha stake for each hotkey.
+    /// 3. Sums the alpha stakes for all hotkeys to calculate the total stake for the coldkey.
+    /// 4. Returns the total alpha stake for the coldkey on the subnet.
+    ///
+    /// # Arguments
+    /// * `coldkey` - The account ID of the coldkey.
+    /// * `netuid` - The unique identifier of the subnet.
+    ///
+    /// # Returns
+    /// * `u64` - The total alpha value for the coldkey on the specified subnet.
+    ///
+    /// # Note
+    /// This function returns the cumulative stake across all hotkeys associated with this coldkey on the subnet.
+    /// This value represents the sum of stakes from all hotkeys associated with this coldkey.
+    pub fn get_stake_for_coldkey_on_subnet(coldkey: &T::AccountId, netuid: u16) -> u64 {
+        // Retrieve the list of hotkeys associated with the coldkey on the subnet.
+        let hotkeys: Vec<T::AccountId> = StakingHotkeys::<T>::get(coldkey);
+
+        // Calculate the total alpha stake for the coldkey on the subnet.
+        let total_stake: u64 = hotkeys
+            .iter()
+            .map(|hotkey| Self::get_stake_for_hotkey_on_subnet(hotkey, netuid))
+            .sum();
+
+        // Return the total alpha stake for the coldkey on the subnet.
+        total_stake
+    }
+
+    /// Retrieves the total stake (alpha) for a given coldkey across all subnets
+    ///
+    /// This function performs the following steps:
+    /// 1. Retrieves the list of subnets associated with the coldkey.
+    /// 2. Iterates through the list of subnets and retrieves the alpha stake for each subnet.
+    /// 3. Sums the alpha stakes for all subnets to calculate the total stake for the coldkey.
+    /// 4. Returns the total alpha stake for the coldkey across all subnets.
+    ///
+    /// # Arguments
+    /// * `coldkey` - The account ID of the coldkey.
+    ///
+    /// # Returns
+    /// * `u64` - The total alpha value for the coldkey across all subnets.
+    ///
+    /// # Note
+    /// This function returns the cumulative stake across all subnets for the coldkey.
+    /// This value represents the sum of stakes from all subnets associated with this coldkey.
+    /// This function is useful for calculating the total stake for a coldkey across all subnets.
+    pub fn get_stake_for_coldkey(coldkey: &T::AccountId) -> u64 {
+        // get number of subnets
+        let netuids = Self::get_all_subnet_netuids();
+
+        // Calculate the total alpha stake for the coldkey across all subnets.
+        let total_stake: u64 = netuids
+            .iter()
+            .map(|netuid| Self::get_stake_for_coldkey_on_subnet(coldkey, *netuid))
+            .sum();
+        // Return the total alpha stake for the coldkey across all subnets.
+        total_stake
+    }
+
     /// Increase hotkey stake on a subnet.
     ///
     /// The function updates share totals given current prices.

runtime/src/precompiles/solidity/staking.abi

@@ -94,5 +94,43 @@
         "outputs": [],
         "stateMutability": "nonpayable",
         "type": "function"
+    },
+    {
+        "inputs": [
+            {
+                "internalType": "bytes32",
+                "name": "coldkey",
+                "type": "bytes32"
+            }
+        ],
+        "name": "getTotalColdkeyStake",
+        "outputs": [
+            {
+                "internalType": "uint256",
+                "name": "",
+                "type": "uint256"
+            }
+        ],
+        "stateMutability": "view",
+        "type": "function"
+    },
+    {
+        "inputs": [
+            {
+                "internalType": "bytes32",
+                "name": "hotkey",
+                "type": "bytes32"
+            }
+        ],
+        "name": "getTotalHotkeyStake",
+        "outputs": [
+            {
+                "internalType": "uint256",
+                "name": "",
+                "type": "uint256"
+            }
+        ],
+        "stateMutability": "view",
+        "type": "function"
     }
 ]

runtime/src/precompiles/solidity/staking.sol

@@ -42,6 +42,30 @@ interface IStaking {
    */
   function removeStake(bytes32 hotkey, uint256 amount, uint256 netuid) external;
 
+  /**
+   * @dev Returns the amount of RAO staked by the coldkey.
+   *
+   * This function allows external accounts and contracts to query the amount of RAO staked by the coldkey
+   * which effectively calls `get_total_coldkey_stake` on the subtensor pallet with
+   * specified coldkey as a parameter.
+   *
+   * @param coldkey The coldkey public key (32 bytes).
+   * @return The amount of RAO staked by the coldkey.
+   */
+  function getTotalColdkeyStake(bytes32 coldkey) external view returns (uint256);
+
+  /**
+   * @dev Returns the total amount of stake under a hotkey (delegative or otherwise)
+   * 
+   * This function allows external accounts and contracts to query the total amount of RAO staked under a hotkey
+   * which effectively calls `get_total_hotkey_stake` on the subtensor pallet with
+   * specified hotkey as a parameter.
+   * 
+   * @param hotkey The hotkey public key (32 bytes).
+   * @return The total amount of RAO staked under the hotkey.
+   */
+  function getTotalHotkeyStake(bytes32 hotkey) external view returns (uint256);
+
   /**
    * @dev Delegates staking to a proxy account.
    *

runtime/src/precompiles/staking.rs

@@ -61,6 +61,10 @@ impl StakingPrecompile {
             Self::remove_stake(handle, &method_input)
         } else if method_id == get_method_id("getStake(bytes32,bytes32,uint256)") {
             Self::get_stake(&method_input)
+        } else if method_id == get_method_id("getTotalColdkeyStake(bytes32)") {
+            Self::get_total_coldkey_stake(&method_input)
+        } else if method_id == get_method_id("getTotalHotkeyStake(bytes32)") {
+            Self::get_total_hotkey_stake(&method_input)
         } else if method_id == get_method_id("addProxy(bytes32)") {
             Self::add_proxy(handle, &method_input)
         } else if method_id == get_method_id("removeProxy(bytes32)") {
@@ -114,6 +118,49 @@ impl StakingPrecompile {
         Self::dispatch(handle, call)
     }
 
+    fn get_total_coldkey_stake(data: &[u8]) -> PrecompileResult {
+        let coldkey: AccountId32 = Self::parse_pub_key(data)?.into();
+
+        // get total stake of coldkey
+        // TODO: is using the function that was written for this purpose in the pallet the right way to go about this?
+        let total_stake = pallet_subtensor::Pallet::<Runtime>::get_stake_for_coldkey(&coldkey);
+        // Convert to EVM decimals
+        let stake_u256 = U256::from(total_stake);
+        let stake_eth =
+            <Runtime as pallet_evm::Config>::BalanceConverter::into_evm_balance(stake_u256)
+                .ok_or(ExitError::InvalidRange)?;
+
+        // Format output
+        let mut result = [0_u8; 32];
+        U256::to_big_endian(&stake_eth, &mut result);
+
+        Ok(PrecompileOutput {
+            exit_status: ExitSucceed::Returned,
+            output: result.into(),
+        })
+    }
+
+    fn get_total_hotkey_stake(data: &[u8]) -> PrecompileResult {
+        let hotkey: AccountId32 = Self::parse_pub_key(data)?.into();
+
+        // get total stake of hotkey
+        let total_stake = pallet_subtensor::Pallet::<Runtime>::get_total_stake_for_hotkey(&hotkey);
+        // Convert to EVM decimals
+        let stake_u256 = U256::from(total_stake);
+        let stake_eth =
+            <Runtime as pallet_evm::Config>::BalanceConverter::into_evm_balance(stake_u256)
+                .ok_or(ExitError::InvalidRange)?;
+
+        // Format output
+        let mut result = [0_u8; 32];
+        U256::to_big_endian(&stake_eth, &mut result);
+
+        Ok(PrecompileOutput {
+            exit_status: ExitSucceed::Returned,
+            output: result.into(),
+        })
+    }
+
     fn add_proxy(handle: &mut impl PrecompileHandle, data: &[u8]) -> PrecompileResult {
         let delegate = AccountId32::from(Self::parse_pub_key(data)?);
         let delegate = <Runtime as frame_system::Config>::Lookup::unlookup(delegate);
@@ -140,7 +187,7 @@ impl StakingPrecompile {
 
     fn get_stake(data: &[u8]) -> PrecompileResult {
         let (hotkey, coldkey) = Self::parse_hotkey_coldkey(data)?;
-        let netuid = Self::parse_netuid(data, 0x5E)?;
+        let netuid: u16 = Self::parse_netuid(data, 0x5E)?;
 
         let stake = pallet_subtensor::Pallet::<Runtime>::get_stake_for_hotkey_and_coldkey_on_subnet(
             &hotkey.into(),