diff --git a/pallets/subtensor/src/staking/stake_utils.rs b/pallets/subtensor/src/staking/stake_utils.rs index 805b61e50d..e247f7f4f5 100644 --- a/pallets/subtensor/src/staking/stake_utils.rs +++ b/pallets/subtensor/src/staking/stake_utils.rs @@ -2,7 +2,7 @@ use super::*; use safe_math::*; use share_pool::{SharePool, SharePoolDataOperations}; use sp_std::ops::Neg; -use substrate_fixed::types::{I64F64, I96F32, U64F64}; +use substrate_fixed::types::{I110F18, I64F64, I96F32, U64F64}; impl Pallet { /// Retrieves the total alpha issuance for a given subnet. @@ -469,16 +469,15 @@ impl Pallet { // Step 2: Initialized vars. if mechanism_id == 1 { // Step 3.a.1: Dynamic mechanism calculations - let tao_reserves: I96F32 = I96F32::saturating_from_num(SubnetTAO::::get(netuid)); - let alpha_reserves: I96F32 = - I96F32::saturating_from_num(SubnetAlphaIn::::get(netuid)); + let tao_reserves: I110F18 = I110F18::saturating_from_num(SubnetTAO::::get(netuid)); + let alpha_reserves: I110F18 = + I110F18::saturating_from_num(SubnetAlphaIn::::get(netuid)); // Step 3.a.2: Compute constant product k = alpha * tao - let k: I96F32 = alpha_reserves.saturating_mul(tao_reserves); + let k: I110F18 = alpha_reserves.saturating_mul(tao_reserves); // Calculate new alpha reserve - let new_alpha_reserves: I96F32 = k - .checked_div(tao_reserves.saturating_add(I96F32::saturating_from_num(tao))) - .unwrap_or(I96F32::saturating_from_num(0)); + let new_alpha_reserves: I110F18 = + k.safe_div(tao_reserves.saturating_add(I110F18::saturating_from_num(tao))); // Step 3.a.3: Calculate alpha staked using the constant product formula // alpha_stake_recieved = current_alpha - (k / (current_tao + new_tao)) @@ -509,16 +508,16 @@ impl Pallet { // Step 2: Swap alpha and attain tao if mechanism_id == 1 { // Step 3.a.1: Dynamic mechanism calculations - let tao_reserves: I96F32 = I96F32::saturating_from_num(SubnetTAO::::get(netuid)); - let alpha_reserves: I96F32 = - I96F32::saturating_from_num(SubnetAlphaIn::::get(netuid)); + let tao_reserves: I110F18 = I110F18::saturating_from_num(SubnetTAO::::get(netuid)); + let alpha_reserves: I110F18 = + I110F18::saturating_from_num(SubnetAlphaIn::::get(netuid)); // Step 3.a.2: Compute constant product k = alpha * tao - let k: I96F32 = alpha_reserves.saturating_mul(tao_reserves); + let k: I110F18 = alpha_reserves.saturating_mul(tao_reserves); // Calculate new tao reserve - let new_tao_reserves: I96F32 = k - .checked_div(alpha_reserves.saturating_add(I96F32::saturating_from_num(alpha))) - .unwrap_or(I96F32::saturating_from_num(0)); + let new_tao_reserves: I110F18 = k + .checked_div(alpha_reserves.saturating_add(I110F18::saturating_from_num(alpha))) + .unwrap_or(I110F18::saturating_from_num(0)); // Step 3.a.3: Calculate alpha staked using the constant product formula // tao_recieved = tao_reserves - (k / (alpha_reserves + new_tao)) diff --git a/pallets/subtensor/src/tests/staking.rs b/pallets/subtensor/src/tests/staking.rs index a570175c9a..b678968dff 100644 --- a/pallets/subtensor/src/tests/staking.rs +++ b/pallets/subtensor/src/tests/staking.rs @@ -2309,3 +2309,42 @@ fn test_unstake_low_liquidity_validate() { ); }); } + +#[test] +fn test_stake_oveflow() { + new_test_ext(1).execute_with(|| { + let subnet_owner_coldkey = U256::from(1001); + let subnet_owner_hotkey = U256::from(1002); + let coldkey_account_id = U256::from(435445); + let hotkey_account_id = U256::from(54544); + let netuid = add_dynamic_network(&subnet_owner_hotkey, &subnet_owner_coldkey); + let amount = 21_000_000_000_000_000; // Max TAO supply + let fee = DefaultStakingFee::::get(); + register_ok_neuron(netuid, hotkey_account_id, coldkey_account_id, 192213123); + + // Give it some $$$ in his coldkey balance + SubtensorModule::add_balance_to_coldkey_account(&coldkey_account_id, amount); + + // Setup liquidity with 21M TAO values + SubnetTAO::::insert(netuid, amount); + SubnetAlphaIn::::insert(netuid, amount); + + // Stake and check if the result is ok + assert_ok!(SubtensorModule::add_stake( + RuntimeOrigin::signed(coldkey_account_id), + hotkey_account_id, + netuid, + amount + )); + + // Check if stake has increased properly (staking 1:1 to SubnetTAO results in SubnetTAO/2 alpha) + assert_abs_diff_eq!( + SubtensorModule::get_stake_for_hotkey_on_subnet(&hotkey_account_id, netuid), + (amount - fee) / 2, + epsilon = amount / 1_000_000, + ); + + // Check if total stake has increased accordingly. + assert_abs_diff_eq!(SubtensorModule::get_total_stake(), amount, epsilon = 10,); + }); +}