Merge pull request #1947 from opentensor/fix/root-reserves

Fix updating root and stable subnets counters on stakes and unstakes

Author: sam0x17

pallets/subtensor/src/staking/stake_utils.rs

@@ -572,53 +572,54 @@ impl<T: Config> Pallet<T> {
     ) -> Result<SwapResult, DispatchError> {
         // Step 1: Get the mechanism type for the subnet (0 for Stable, 1 for Dynamic)
         let mechanism_id: u16 = SubnetMechanism::<T>::get(netuid);
-        if mechanism_id == 1 {
-            let swap_result = T::SwapInterface::swap(
+        let swap_result = if mechanism_id == 1 {
+            T::SwapInterface::swap(
                 netuid.into(),
                 OrderType::Buy,
                 tao.into(),
                 price_limit.into(),
                 drop_fees,
                 false,
-            )?;
-            let alpha_decrease =
-                AlphaCurrency::from(swap_result.alpha_reserve_delta.unsigned_abs());
-
-            // Decrease Alpha reserves.
-            Self::decrease_provided_alpha_reserve(netuid.into(), alpha_decrease);
-
-            // Increase Alpha outstanding.
-            SubnetAlphaOut::<T>::mutate(netuid, |total| {
-                *total = total.saturating_add(swap_result.amount_paid_out.into());
-            });
-
-            // Increase only the protocol TAO reserve. We only use the sum of
-            // (SubnetTAO + SubnetTaoProvided) in tao_reserve(), so it is irrelevant
-            // which one to increase.
-            SubnetTAO::<T>::mutate(netuid, |total| {
-                *total = total.saturating_add((swap_result.tao_reserve_delta as u64).into());
-            });
-
-            // Increase Total Tao reserves.
-            TotalStake::<T>::mutate(|total| *total = total.saturating_add(tao));
-
-            // Increase total subnet TAO volume.
-            SubnetVolume::<T>::mutate(netuid, |total| {
-                *total = total.saturating_add(tao.to_u64() as u128);
-            });
-
-            // Return the alpha received.
-            Ok(swap_result)
+            )?
         } else {
+            let abs_delta: u64 = tao.into();
+
             // Step 3.b.1: Stable mechanism, just return the value 1:1
-            Ok(SwapResult {
+            SwapResult {
                 amount_paid_in: tao.into(),
                 amount_paid_out: tao.into(),
                 fee_paid: 0,
-                tao_reserve_delta: 0,
-                alpha_reserve_delta: 0,
-            })
-        }
+                tao_reserve_delta: abs_delta as i64,
+                alpha_reserve_delta: (abs_delta as i64).neg(),
+            }
+        };
+
+        let alpha_decrease = AlphaCurrency::from(swap_result.alpha_reserve_delta.unsigned_abs());
+
+        // Decrease Alpha reserves.
+        Self::decrease_provided_alpha_reserve(netuid.into(), alpha_decrease);
+
+        // Increase Alpha outstanding.
+        SubnetAlphaOut::<T>::mutate(netuid, |total| {
+            *total = total.saturating_add(swap_result.amount_paid_out.into());
+        });
+
+        // Increase only the protocol TAO reserve. We only use the sum of
+        // (SubnetTAO + SubnetTaoProvided) in tao_reserve(), so it is irrelevant
+        // which one to increase.
+        SubnetTAO::<T>::mutate(netuid, |total| {
+            *total = total.saturating_add((swap_result.tao_reserve_delta as u64).into());
+        });
+
+        // Increase Total Tao reserves.
+        TotalStake::<T>::mutate(|total| *total = total.saturating_add(tao));
+
+        // Increase total subnet TAO volume.
+        SubnetVolume::<T>::mutate(netuid, |total| {
+            *total = total.saturating_add(tao.to_u64() as u128);
+        });
+
+        Ok(swap_result)
     }
 
     /// Swaps a subnet's Alpha token for TAO.
@@ -633,62 +634,64 @@ impl<T: Config> Pallet<T> {
         // Step 1: Get the mechanism type for the subnet (0 for Stable, 1 for Dynamic)
         let mechanism_id: u16 = SubnetMechanism::<T>::get(netuid);
         // Step 2: Swap alpha and attain tao
-        if mechanism_id == 1 {
-            let swap_result = T::SwapInterface::swap(
+        let swap_result = if mechanism_id == 1 {
+            T::SwapInterface::swap(
                 netuid.into(),
                 OrderType::Sell,
                 alpha.into(),
                 price_limit.into(),
                 drop_fees,
                 false,
-            )?;
-
-            // Increase only the protocol Alpha reserve. We only use the sum of
-            // (SubnetAlphaIn + SubnetAlphaInProvided) in alpha_reserve(), so it is irrelevant
-            // which one to increase.
-            SubnetAlphaIn::<T>::mutate(netuid, |total| {
-                *total = total.saturating_add((swap_result.alpha_reserve_delta as u64).into());
-            });
-
-            // Decrease Alpha outstanding.
-            // TODO: Deprecate, not accurate in v3 anymore
-            SubnetAlphaOut::<T>::mutate(netuid, |total| {
-                *total = total.saturating_sub((swap_result.alpha_reserve_delta as u64).into());
-            });
-
-            // Decrease tao reserves.
-            Self::decrease_provided_tao_reserve(
-                netuid.into(),
-                swap_result
-                    .tao_reserve_delta
-                    .abs()
-                    .try_into()
-                    .unwrap_or(0)
-                    .into(),
-            );
-
-            // Reduce total TAO reserves.
-            TotalStake::<T>::mutate(|total| {
-                *total = total.saturating_sub(swap_result.amount_paid_out.into())
-            });
-
-            // Increase total subnet TAO volume.
-            SubnetVolume::<T>::mutate(netuid, |total| {
-                *total = total.saturating_add(swap_result.amount_paid_out.into())
-            });
-
-            // Return the tao received.
-            Ok(swap_result)
+            )?
         } else {
+            let abs_delta: u64 = alpha.into();
+
             // Step 3.b.1: Stable mechanism, just return the value 1:1
-            Ok(SwapResult {
+            SwapResult {
                 amount_paid_in: alpha.into(),
                 amount_paid_out: alpha.into(),
                 fee_paid: 0,
-                tao_reserve_delta: 0,
-                alpha_reserve_delta: 0,
-            })
-        }
+                tao_reserve_delta: (abs_delta as i64).neg(),
+                alpha_reserve_delta: abs_delta as i64,
+            }
+        };
+
+        // Increase only the protocol Alpha reserve. We only use the sum of
+        // (SubnetAlphaIn + SubnetAlphaInProvided) in alpha_reserve(), so it is irrelevant
+        // which one to increase.
+        SubnetAlphaIn::<T>::mutate(netuid, |total| {
+            *total = total.saturating_add((swap_result.alpha_reserve_delta as u64).into());
+        });
+
+        // Decrease Alpha outstanding.
+        // TODO: Deprecate, not accurate in v3 anymore
+        SubnetAlphaOut::<T>::mutate(netuid, |total| {
+            *total = total.saturating_sub((swap_result.alpha_reserve_delta as u64).into());
+        });
+
+        // Decrease tao reserves.
+        Self::decrease_provided_tao_reserve(
+            netuid.into(),
+            swap_result
+                .tao_reserve_delta
+                .abs()
+                .try_into()
+                .unwrap_or(0)
+                .into(),
+        );
+
+        // Reduce total TAO reserves.
+        TotalStake::<T>::mutate(|total| {
+            *total = total.saturating_sub(swap_result.amount_paid_out.into())
+        });
+
+        // Increase total subnet TAO volume.
+        SubnetVolume::<T>::mutate(netuid, |total| {
+            *total = total.saturating_add(swap_result.amount_paid_out.into())
+        });
+
+        // Return the tao received.
+        Ok(swap_result)
     }
 
     /// Unstakes alpha from a subnet for a given hotkey and coldkey pair.

pallets/subtensor/src/tests/staking.rs

@@ -5425,3 +5425,123 @@ fn test_stake_rate_limits() {
         )));
     });
 }
+
+// cargo test --package pallet-subtensor --lib -- tests::staking::test_add_root_updates_counters --exact --show-output
+#[test]
+fn test_add_root_updates_counters() {
+    new_test_ext(0).execute_with(|| {
+        let hotkey_account_id = U256::from(561337);
+        let coldkey_account_id = U256::from(61337);
+        add_network(NetUid::ROOT, 10, 0);
+        assert_ok!(SubtensorModule::root_register(
+            RuntimeOrigin::signed(coldkey_account_id).clone(),
+            hotkey_account_id,
+        ));
+        let stake_amount = 1_000_000_000;
+
+        // Give it some $$$ in his coldkey balance
+        let initial_balance = stake_amount + ExistentialDeposit::get();
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey_account_id, initial_balance);
+
+        // Setup SubnetAlphaIn (because we are going to stake)
+        SubnetAlphaIn::<Test>::insert(NetUid::ROOT, AlphaCurrency::from(stake_amount));
+
+        // Stake to hotkey account, and check if the result is ok
+        assert_ok!(SubtensorModule::add_stake(
+            RuntimeOrigin::signed(coldkey_account_id),
+            hotkey_account_id,
+            NetUid::ROOT,
+            stake_amount.into()
+        ));
+
+        // Check if stake has increased
+        let new_stake = SubtensorModule::get_total_stake_for_hotkey(&hotkey_account_id);
+        assert_eq!(new_stake, stake_amount.into());
+
+        // Check if total stake has increased accordingly.
+        assert_eq!(SubtensorModule::get_total_stake(), stake_amount.into());
+
+        // SubnetTAO updated
+        assert_eq!(SubnetTAO::<Test>::get(NetUid::ROOT), stake_amount.into());
+
+        // SubnetAlphaIn updated
+        assert_eq!(SubnetAlphaIn::<Test>::get(NetUid::ROOT), 0.into());
+
+        // SubnetAlphaOut updated
+        assert_eq!(
+            SubnetAlphaOut::<Test>::get(NetUid::ROOT),
+            stake_amount.into()
+        );
+
+        // SubnetVolume updated
+        assert_eq!(
+            SubnetVolume::<Test>::get(NetUid::ROOT),
+            stake_amount as u128
+        );
+    });
+}
+
+// cargo test --package pallet-subtensor --lib -- tests::staking::test_remove_root_updates_counters --exact --show-output
+#[test]
+fn test_remove_root_updates_counters() {
+    new_test_ext(0).execute_with(|| {
+        let hotkey_account_id = U256::from(561337);
+        let coldkey_account_id = U256::from(61337);
+        add_network(NetUid::ROOT, 10, 0);
+        assert_ok!(SubtensorModule::root_register(
+            RuntimeOrigin::signed(coldkey_account_id).clone(),
+            hotkey_account_id,
+        ));
+        let stake_amount = 1_000_000_000;
+
+        // Give it some $$$ in his coldkey balance
+        let initial_balance = stake_amount + ExistentialDeposit::get();
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey_account_id, initial_balance);
+
+        // Setup existing stake
+        SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey_account_id,
+            &coldkey_account_id,
+            NetUid::ROOT,
+            stake_amount.into(),
+        );
+
+        // Setup TotalStake, SubnetAlphaOut and SubnetTAO (because we are going to unstake)
+        TotalStake::<Test>::set(TaoCurrency::from(stake_amount));
+        SubnetTAO::<Test>::insert(NetUid::ROOT, TaoCurrency::from(stake_amount));
+        SubnetAlphaOut::<Test>::insert(NetUid::ROOT, AlphaCurrency::from(stake_amount));
+
+        // Stake to hotkey account, and check if the result is ok
+        assert_ok!(SubtensorModule::remove_stake(
+            RuntimeOrigin::signed(coldkey_account_id),
+            hotkey_account_id,
+            NetUid::ROOT,
+            stake_amount.into()
+        ));
+
+        // Check if stake has been decreased
+        let new_stake = SubtensorModule::get_total_stake_for_hotkey(&hotkey_account_id);
+        assert_eq!(new_stake, 0.into());
+
+        // Check if total stake has decreased accordingly.
+        assert_eq!(SubtensorModule::get_total_stake(), 0.into());
+
+        // SubnetTAO updated
+        assert_eq!(SubnetTAO::<Test>::get(NetUid::ROOT), 0.into());
+
+        // SubnetAlphaIn updated
+        assert_eq!(
+            SubnetAlphaIn::<Test>::get(NetUid::ROOT),
+            stake_amount.into()
+        );
+
+        // SubnetAlphaOut updated
+        assert_eq!(SubnetAlphaOut::<Test>::get(NetUid::ROOT), 0.into());
+
+        // SubnetVolume updated
+        assert_eq!(
+            SubnetVolume::<Test>::get(NetUid::ROOT),
+            stake_amount as u128
+        );
+    });
+}