Fix balancer performance

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,9 @@ 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 result.min(U64F64::from_num(1));
         }
         U64F64::saturating_from_num(0)
     }
@@ -928,6 +929,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 +995,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

@@ -210,7 +210,7 @@ impl<T: Config> Pallet<T> {
             amount_to_swap,
             limit_price,
             drop_fees,
-        );
+        )?;
 
         let swap_result = swap_step.execute()?;
 

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::*;
 
+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
@@ -45,15 +47,16 @@ where
         amount_remaining: PaidIn,
         limit_price: U64F64,
         drop_fees: bool,
-    ) -> Self {
+    ) -> Result<Self, Error<T>> {
         let fee = Pallet::<T>::calculate_fee_amount(netuid, amount_remaining, drop_fees);
         let requested_delta_in = amount_remaining.saturating_sub(fee);
+        Self::ensure_input_within_reserve_limit(netuid, requested_delta_in)?;
 
         // Target and current prices
         let target_price = Self::price_target(netuid, requested_delta_in);
         let current_price = Pallet::<T>::current_price(netuid);
 
-        Self {
+        Ok(Self {
             netuid,
             drop_fees,
             requested_delta_in,
@@ -64,15 +67,23 @@ where
             final_price: target_price,
             fee,
             _phantom: PhantomData,
-        }
+        })
     }
 
     /// Execute the swap step and return the result
     pub(crate) fn execute(&mut self) -> Result<SwapStepResult<PaidIn, PaidOut>, Error<T>> {
         self.determine_action();
+        Self::ensure_input_within_reserve_limit(self.netuid, self.delta_in)?;
         self.process_swap()
     }
 
+    fn ensure_input_within_reserve_limit(netuid: NetUid, delta_in: PaidIn) -> Result<(), Error<T>> {
+        let input_reserve = Self::input_reserve(netuid);
+        let max_delta_in = input_reserve.saturating_mul(MAX_SWAP_INPUT_RESERVE_MULTIPLIER.into());
+        ensure!(delta_in <= max_delta_in, Error::<T>::SwapInputTooLarge);
+        Ok(())
+    }
+
     /// Determine the appropriate action for this swap step
     fn determine_action(&mut self) {
         let mut recalculate_fee = false;
@@ -176,6 +187,10 @@ impl<T: Config> SwapStep<T, TaoBalance, AlphaBalance>
                 .saturating_to_num::<u64>(),
         )
     }
+
+    fn input_reserve(netuid: NetUid) -> TaoBalance {
+        T::TaoReserve::reserve(netuid.into())
+    }
 }
 
 impl<T: Config> SwapStep<T, AlphaBalance, TaoBalance>
@@ -220,6 +235,10 @@ impl<T: Config> SwapStep<T, AlphaBalance, TaoBalance>
                 .saturating_to_num::<u64>(),
         )
     }
+
+    fn input_reserve(netuid: NetUid) -> AlphaBalance {
+        T::AlphaReserve::reserve(netuid.into())
+    }
 }
 
 pub(crate) trait SwapStep<T, PaidIn, PaidOut>
@@ -244,6 +263,9 @@ where
     /// This is the core method of the swap that tells how much output token is given for an
     /// amount of input token within one price tick.
     fn convert_deltas(netuid: NetUid, delta_in: PaidIn) -> PaidOut;
+
+    /// Return the reserve for the token being paid into this swap step.
+    fn input_reserve(netuid: NetUid) -> PaidIn;
 }
 
 #[derive(Debug, PartialEq)]

pallets/swap/src/pallet/tests.rs

@@ -721,6 +721,60 @@ fn test_rollback_works() {
     })
 }
 
+#[test]
+fn test_swap_rejects_input_over_1000x_input_reserve() {
+    new_test_ext().execute_with(|| {
+        let netuid = NetUid::from(1);
+        TaoReserve::set_mock_reserve(netuid, TaoBalance::from(1_000));
+        AlphaReserve::set_mock_reserve(netuid, AlphaBalance::from(1_000));
+
+        assert_noop!(
+            Pallet::<Test>::do_swap(
+                netuid,
+                GetTaoForAlpha::with_amount(1_000_001),
+                get_min_price(),
+                true,
+                false,
+            ),
+            Error::<Test>::SwapInputTooLarge
+        );
+        assert_noop!(
+            Pallet::<Test>::do_swap(
+                netuid,
+                GetAlphaForTao::with_amount(1_000_001),
+                get_max_price(),
+                true,
+                false,
+            ),
+            Error::<Test>::SwapInputTooLarge
+        );
+    });
+}
+
+#[test]
+fn test_swap_allows_input_at_1000x_input_reserve() {
+    new_test_ext().execute_with(|| {
+        let netuid = NetUid::from(1);
+        TaoReserve::set_mock_reserve(netuid, TaoBalance::from(1_000));
+        AlphaReserve::set_mock_reserve(netuid, AlphaBalance::from(1_000));
+
+        assert_ok!(Pallet::<Test>::do_swap(
+            netuid,
+            GetTaoForAlpha::with_amount(1_000_000),
+            get_min_price(),
+            true,
+            true,
+        ));
+        assert_ok!(Pallet::<Test>::do_swap(
+            netuid,
+            GetAlphaForTao::with_amount(1_000_000),
+            get_max_price(),
+            true,
+            true,
+        ));
+    });
+}
+
 #[allow(dead_code)]
 fn bbox(t: U64F64, a: U64F64, b: U64F64) -> U64F64 {
     if t < a {