Skip to content

Commit

Permalink
feat: modify total stake coldkey func
Browse files Browse the repository at this point in the history
- gets stake for coldkey on subnet and across all subnets
  • Loading branch information
Shr1ftyy committed Jan 27, 2025
1 parent 573a7b1 commit 3de9c3d
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 6 deletions.
63 changes: 63 additions & 0 deletions pallets/subtensor/src/staking/stake_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
25 changes: 19 additions & 6 deletions runtime/src/precompiles/staking.rs
Original file line number Diff line number Diff line change
Expand Up @@ -118,14 +118,21 @@ impl StakingPrecompile {
Self::dispatch(handle, call)
}

fn get_total_coldkey_stake(data: &[u8]) -> PrecompileResult {
fn get_total_coldkey_stake(data: &[u8]) -> PrecompileResult {
let coldkey: AccountId32 = Self::parse_pub_key(data)?.into();

// get total stake of coldkey
let total_stake = pallet_subtensor::Pallet::<Runtime>::get_total_stake_for_coldkey(&coldkey);
let result_u256 = U256::from(total_stake);
// 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(&result_u256, &mut result);
U256::to_big_endian(&stake_eth, &mut result);

Ok(PrecompileOutput {
exit_status: ExitSucceed::Returned,
Expand All @@ -138,9 +145,15 @@ impl StakingPrecompile {

// get total stake of hotkey
let total_stake = pallet_subtensor::Pallet::<Runtime>::get_total_stake_for_hotkey(&hotkey);
let result_u256 = U256::from(total_stake);
// 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(&result_u256, &mut result);
U256::to_big_endian(&stake_eth, &mut result);

Ok(PrecompileOutput {
exit_status: ExitSucceed::Returned,
Expand Down

0 comments on commit 3de9c3d

Please sign in to comment.