Fix balancer performance

chain-extensions/src/tests.rs

@@ -1297,15 +1297,26 @@ fn add_stake_recycle_rollback_on_recycle_failure() {
 
         let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
 
-        // Set up very low reserves so recycle will fail with InsufficientLiquidity
+        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
+        pallet_subtensor::Pallet::<mock::Test>::insert_lock_state(
+            &coldkey,
+            netuid,
+            &hotkey,
+            pallet_subtensor::staking::lock::LockState {
+                locked_mass: AlphaBalance::from(u64::MAX / 4),
+                conviction: U64F64::saturating_from_num(0),
+                last_update: pallet_subtensor::Pallet::<mock::Test>::get_current_block_as_u64(),
+            },
+        );
+
+        // Leave enough input-side liquidity for add_stake to pass the 1000x swap input cap.
+        // The lock above makes the recycle leg fail, exercising atomic rollback.
         mock::setup_reserves(
             netuid,
-            TaoBalance::from(1_000_u64),
+            TaoBalance::from(tao_amount_raw / 1000 + 1),
             AlphaBalance::from(1_000_u64),
         );
 
-        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
-
         add_balance_to_coldkey_account(
             &coldkey,
             TaoBalance::from(tao_amount_raw.saturating_add(1_000_000_000)),
@@ -1368,15 +1379,26 @@ fn add_stake_burn_rollback_on_burn_failure() {
 
         let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
 
-        // Set up very low reserves so burn will fail with InsufficientLiquidity
+        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
+        pallet_subtensor::Pallet::<mock::Test>::insert_lock_state(
+            &coldkey,
+            netuid,
+            &hotkey,
+            pallet_subtensor::staking::lock::LockState {
+                locked_mass: AlphaBalance::from(u64::MAX / 4),
+                conviction: U64F64::saturating_from_num(0),
+                last_update: pallet_subtensor::Pallet::<mock::Test>::get_current_block_as_u64(),
+            },
+        );
+
+        // Leave enough input-side liquidity for add_stake to pass the 1000x swap input cap.
+        // The lock above makes the burn leg fail, exercising atomic rollback.
         mock::setup_reserves(
             netuid,
-            TaoBalance::from(1_000_u64),
+            TaoBalance::from(tao_amount_raw / 1000 + 1),
             AlphaBalance::from(1_000_u64),
         );
 
-        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
-
         add_balance_to_coldkey_account(
             &coldkey,
             TaoBalance::from(tao_amount_raw.saturating_add(1_000_000_000)),

pallets/subtensor/src/staking/recycle_alpha.rs

@@ -1,5 +1,6 @@
 use super::*;
 use crate::{Error, system::ensure_signed};
+use frame_support::storage::{TransactionOutcome, with_transaction};
 use subtensor_runtime_common::{AlphaBalance, NetUid};
 
 impl<T: Config> Pallet<T> {
@@ -132,22 +133,38 @@ impl<T: Config> Pallet<T> {
         amount: TaoBalance,
         limit: Option<TaoBalance>,
     ) -> DispatchResult {
-        let alpha = if let Some(limit) = limit {
-            Self::do_add_stake_limit(origin.clone(), hotkey.clone(), netuid, amount, limit, false)?
-        } else {
-            Self::do_add_stake(origin.clone(), hotkey.clone(), netuid, amount)?
-        };
-
-        Self::do_burn_alpha(origin, hotkey.clone(), alpha, netuid)?;
-
-        Self::deposit_event(Event::AddStakeBurn {
-            netuid,
-            hotkey,
-            amount,
-            alpha,
-        });
-
-        Ok(())
+        with_transaction(|| {
+            let result = (|| {
+                let alpha = if let Some(limit) = limit {
+                    Self::do_add_stake_limit(
+                        origin.clone(),
+                        hotkey.clone(),
+                        netuid,
+                        amount,
+                        limit,
+                        false,
+                    )?
+                } else {
+                    Self::do_add_stake(origin.clone(), hotkey.clone(), netuid, amount)?
+                };
+
+                Self::do_burn_alpha(origin, hotkey.clone(), alpha, netuid)?;
+
+                Self::deposit_event(Event::AddStakeBurn {
+                    netuid,
+                    hotkey,
+                    amount,
+                    alpha,
+                });
+
+                Ok(())
+            })();
+
+            match result {
+                Ok(()) => TransactionOutcome::Commit(Ok(())),
+                Err(err) => TransactionOutcome::Rollback(Err(err)),
+            }
+        })
     }
 
     /// Atomically stakes TAO and recycles the resulting alpha.
@@ -160,8 +177,17 @@ impl<T: Config> Pallet<T> {
         netuid: NetUid,
         amount: TaoBalance,
     ) -> Result<AlphaBalance, DispatchError> {
-        let alpha = Self::do_add_stake(origin.clone(), hotkey.clone(), netuid, amount)?;
-        Self::do_recycle_alpha(origin, hotkey, alpha, netuid)
+        with_transaction(|| {
+            let result = (|| {
+                let alpha = Self::do_add_stake(origin.clone(), hotkey.clone(), netuid, amount)?;
+                Self::do_recycle_alpha(origin, hotkey, alpha, netuid)
+            })();
+
+            match result {
+                Ok(alpha) => TransactionOutcome::Commit(Ok(alpha)),
+                Err(err) => TransactionOutcome::Rollback(Err(err)),
+            }
+        })
     }
 
     /// Atomically stakes TAO and burns the resulting alpha. Permissionless
@@ -173,7 +199,16 @@ impl<T: Config> Pallet<T> {
         netuid: NetUid,
         amount: TaoBalance,
     ) -> Result<AlphaBalance, DispatchError> {
-        let alpha = Self::do_add_stake(origin.clone(), hotkey.clone(), netuid, amount)?;
-        Self::do_burn_alpha(origin, hotkey, alpha, netuid)
+        with_transaction(|| {
+            let result = (|| {
+                let alpha = Self::do_add_stake(origin.clone(), hotkey.clone(), netuid, amount)?;
+                Self::do_burn_alpha(origin, hotkey, alpha, netuid)
+            })();
+
+            match result {
+                Ok(alpha) => TransactionOutcome::Commit(Ok(alpha)),
+                Err(err) => TransactionOutcome::Rollback(Err(err)),
+            }
+        })
     }
 }

pallets/subtensor/src/tests/staking.rs

@@ -776,9 +776,9 @@ fn test_add_stake_insufficient_liquidity() {
     });
 }
 
-/// cargo test --package pallet-subtensor --lib -- tests::staking::test_add_stake_insufficient_liquidity_one_side_ok --exact --show-output
+/// cargo test --package pallet-subtensor --lib -- tests::staking::test_add_stake_input_reserve_too_low_fails --exact --show-output
 #[test]
-fn test_add_stake_insufficient_liquidity_one_side_ok() {
+fn test_add_stake_input_reserve_too_low_fails() {
     new_test_ext(1).execute_with(|| {
         let subnet_owner_coldkey = U256::from(1001);
         let subnet_owner_hotkey = U256::from(1002);
@@ -795,13 +795,17 @@ fn test_add_stake_insufficient_liquidity_one_side_ok() {
         let reserve_tao = u64::from(mock::SwapMinimumReserve::get()) - 1;
         mock::setup_reserves(netuid, reserve_tao.into(), reserve_alpha.into());
 
-        // Check the error
-        assert_ok!(SubtensorModule::add_stake(
-            RuntimeOrigin::signed(coldkey),
-            hotkey,
-            netuid,
-            amount_staked.into()
-        ));
+        // The output-side reserve is sufficient, but the input-side reserve is too small for the
+        // requested swap under the 1000x input-reserve cap.
+        assert_noop!(
+            SubtensorModule::add_stake(
+                RuntimeOrigin::signed(coldkey),
+                hotkey,
+                netuid,
+                amount_staked.into()
+            ),
+            pallet_subtensor_swap::Error::<Test>::SwapInputTooLarge
+        );
     });
 }
 
@@ -876,7 +880,7 @@ fn test_remove_stake_insufficient_liquidity() {
 
         // Mock more liquidity - remove becomes successful
         SubnetTAO::<Test>::insert(netuid, TaoBalance::from(amount_staked + 1));
-        SubnetAlphaIn::<Test>::insert(netuid, AlphaBalance::from(1));
+        SubnetAlphaIn::<Test>::insert(netuid, AlphaBalance::from(alpha.to_u64() / 1000 + 1));
         assert_ok!(SubtensorModule::remove_stake(
             RuntimeOrigin::signed(coldkey),
             hotkey,
@@ -5248,15 +5252,15 @@ fn test_large_swap() {
         // add network
         let netuid = add_dynamic_network(&owner_hotkey, &owner_coldkey);
         add_balance_to_coldkey_account(&coldkey, 1_000_000_000_000_000_u64.into());
-        let tao = TaoBalance::from(100_000_000u64);
+        let swap_amount = TaoBalance::from(100_000_000_000_000_u64);
+        let tao = TaoBalance::from(swap_amount.to_u64() / 1000);
         let alpha = AlphaBalance::from(1_000_000_000_000_000_u64);
         SubnetTAO::<Test>::insert(netuid, tao);
         SubnetAlphaIn::<Test>::insert(netuid, alpha);
 
         // Force the swap to initialize
         <Test as pallet::Config>::SwapInterface::init_swap(netuid, None);
 
-        let swap_amount = TaoBalance::from(100_000_000_000_000_u64);
         assert_ok!(SubtensorModule::add_stake(
             RuntimeOrigin::signed(coldkey),
             owner_hotkey,

pallets/swap/src/pallet/balancer.rs

@@ -149,7 +149,7 @@ impl Balancer {
         let w1: u128 = self.get_base_weight().deconstruct() as u128;
         let w2: u128 = self.get_quote_weight().deconstruct() as u128;
 
-        let precision = 1024;
+        let precision = 256;
         let x_safe = SafeInt::from(x);
         let w1_safe = SafeInt::from(w1);
         let w2_safe = SafeInt::from(w2);
@@ -187,8 +187,13 @@ impl Balancer {
         if let Some(result_safe_int) = maybe_result_safe_int
             && let Some(result_u64) = result_safe_int.to_u64()
         {
-            return U64F64::saturating_from_num(result_u64)
+            let result = U64F64::saturating_from_num(result_u64)
                 .safe_div(U64F64::saturating_from_num(ACCURACY));
+            return if dx >= 0 {
+                result.min(U64F64::from_num(1))
+            } else {
+                result
+            };
         }
         U64F64::saturating_from_num(0)
     }
@@ -791,6 +796,12 @@ mod tests {
                 let dy1 = y_fixed * (one - e1);
                 let dy2 = y_fixed * (one - e2);
 
+                if dx > x.saturating_mul(1_000) {
+                    assert!(e1 <= one);
+                    assert!(e2 <= one);
+                    return;
+                }
+
                 let w1 = perquintill_to_f64(bal.get_base_weight());
                 let w2 = perquintill_to_f64(bal.get_quote_weight());
                 let e1_expected = (x as f64 / (x as f64 + dx as f64)).powf(w1 / w2);
@@ -928,6 +939,7 @@ mod tests {
         }
     }
 
+    // cargo test --package pallet-subtensor-swap --lib -- pallet::balancer::tests::test_exp_quote_fuzzy --include-ignored --exact --nocapture
     #[ignore]
     #[test]
     fn test_exp_quote_fuzzy() {
@@ -993,7 +1005,7 @@ mod tests {
 
             // Print progress
             let done = counter.fetch_add(1, Ordering::Relaxed) + 1;
-            if done % 100_000_000 == 0 {
+            if done % 10_000_000 == 0 {
                 let progress = done as f64 / ITERATIONS as f64 * 100.0;
                 // Replace with println for real-time progress
                 log::debug!("progress = {progress:.4}%");

pallets/swap/src/pallet/impls.rs

@@ -14,7 +14,7 @@ use subtensor_swap_interface::{
 };
 
 use super::pallet::*;
-use super::swap_step::{BasicSwapStep, SwapStep};
+use super::swap_step::{BasicSwapStep, MAX_SWAP_INPUT_RESERVE_MULTIPLIER, SwapStep};
 use crate::{pallet::Balancer, pallet::balancer::BalancerError};
 
 impl<T: Config> Pallet<T> {
@@ -154,8 +154,17 @@ impl<T: Config> Pallet<T> {
         transactional::with_transaction(|| {
             let reserve = Order::ReserveOut::reserve(netuid.into());
 
-            let result = Self::swap_inner::<Order>(netuid, order, limit_price, drop_fees)
-                .map_err(Into::into);
+            let result = if simulate {
+                Ok(())
+            } else {
+                Self::ensure_swap_input_within_reserve_limit::<Order>(
+                    netuid,
+                    order.amount(),
+                    drop_fees,
+                )
+            }
+            .and_then(|_| Self::swap_inner::<Order>(netuid, order, limit_price, drop_fees))
+            .map_err(Into::into);
 
             if simulate || result.is_err() {
                 // Simulation only
@@ -177,6 +186,23 @@ impl<T: Config> Pallet<T> {
         })
     }
 
+    fn ensure_swap_input_within_reserve_limit<Order>(
+        netuid: NetUid,
+        amount: Order::PaidIn,
+        drop_fees: bool,
+    ) -> Result<(), Error<T>>
+    where
+        Order: OrderT,
+    {
+        let fee = Self::calculate_fee_amount(netuid, amount, drop_fees);
+        let net_amount = amount.saturating_sub(fee);
+        let input_reserve = Order::ReserveIn::reserve(netuid);
+        let max_amount = input_reserve.saturating_mul(MAX_SWAP_INPUT_RESERVE_MULTIPLIER.into());
+
+        ensure!(net_amount <= max_amount, Error::<T>::SwapInputTooLarge);
+        Ok(())
+    }
+
     fn swap_inner<Order>(
         netuid: NetUid,
         order: Order,

pallets/swap/src/pallet/mod.rs

@@ -165,6 +165,9 @@ mod pallet {
         /// Swap reserves are too imbalanced
         ReservesOutOfBalance,
 
+        /// Swap input is too large relative to input-side liquidity
+        SwapInputTooLarge,
+
         /// The extrinsic is deprecated
         Deprecated,
     }

pallets/swap/src/pallet/swap_step.rs

@@ -7,6 +7,8 @@ use subtensor_runtime_common::{AlphaBalance, NetUid, TaoBalance, Token, TokenRes
 
 use super::pallet::*;
 
+pub(crate) const MAX_SWAP_INPUT_RESERVE_MULTIPLIER: u64 = 1_000;
+
 /// A struct representing a single swap step with all its parameters and state
 pub(crate) struct BasicSwapStep<T, PaidIn, PaidOut>
 where