Merge pull request #1211 from opentensor/fix/k-precision

Change type for pool k to I110F18

Author: sam0x17

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<T: Config> Pallet<T> {
     /// Retrieves the total alpha issuance for a given subnet.
@@ -469,16 +469,15 @@ impl<T: Config> Pallet<T> {
         // Step 2: Initialized vars.
         if mechanism_id == 1 {
             // Step 3.a.1: Dynamic mechanism calculations
-            let tao_reserves: I96F32 = I96F32::saturating_from_num(SubnetTAO::<T>::get(netuid));
-            let alpha_reserves: I96F32 =
-                I96F32::saturating_from_num(SubnetAlphaIn::<T>::get(netuid));
+            let tao_reserves: I110F18 = I110F18::saturating_from_num(SubnetTAO::<T>::get(netuid));
+            let alpha_reserves: I110F18 =
+                I110F18::saturating_from_num(SubnetAlphaIn::<T>::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<T: Config> Pallet<T> {
         // 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::<T>::get(netuid));
-            let alpha_reserves: I96F32 =
-                I96F32::saturating_from_num(SubnetAlphaIn::<T>::get(netuid));
+            let tao_reserves: I110F18 = I110F18::saturating_from_num(SubnetTAO::<T>::get(netuid));
+            let alpha_reserves: I110F18 =
+                I110F18::saturating_from_num(SubnetAlphaIn::<T>::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))

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::<Test>::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::<Test>::insert(netuid, amount);
+        SubnetAlphaIn::<Test>::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,);
+    });
+}