Add protocol fees to protocol LP proportionally to the price

pallets/swap/src/pallet/impls.rs

@@ -113,6 +113,41 @@ impl<T: Config> Pallet<T> {
         Ok(())
     }
 
+    pub(crate) fn get_proportional_alpha_tao_and_remainders(
+        sqrt_alpha_price: U64F64,
+        amount_tao: TaoCurrency,
+        amount_alpha: AlphaCurrency,
+    ) -> (TaoCurrency, AlphaCurrency, TaoCurrency, AlphaCurrency) {
+        let price = sqrt_alpha_price.saturating_mul(sqrt_alpha_price);
+        let tao_equivalent: u64 = U64F64::saturating_from_num(u64::from(amount_alpha))
+            .saturating_mul(price)
+            .saturating_to_num();
+        let amount_tao_u64 = u64::from(amount_tao);
+
+        if tao_equivalent <= amount_tao_u64 {
+            // Too much or just enough TAO
+            (
+                tao_equivalent.into(),
+                amount_alpha,
+                amount_tao.saturating_sub(TaoCurrency::from(tao_equivalent)),
+                0.into(),
+            )
+        } else {
+            // Too much Alpha
+            let alpha_equivalent: u64 = U64F64::saturating_from_num(u64::from(amount_tao))
+                .safe_div(price)
+                .saturating_to_num();
+            (
+                amount_tao,
+                alpha_equivalent.into(),
+                0.into(),
+                u64::from(amount_alpha)
+                    .saturating_sub(alpha_equivalent)
+                    .into(),
+            )
+        }
+    }
+
     /// Adjusts protocol liquidity with new values of TAO and Alpha reserve
     pub(super) fn adjust_protocol_liquidity(
         netuid: NetUid,
@@ -129,17 +164,31 @@ impl<T: Config> Pallet<T> {
             // Claim protocol fees and add them to liquidity
             let (tao_fees, alpha_fees) = position.collect_fees();
 
-            // Adjust liquidity
+            // Add fee reservoirs and get proportional amounts
             let current_sqrt_price = AlphaSqrtPrice::<T>::get(netuid);
+            let tao_reservoir = ScrapReservoirTao::<T>::get(netuid);
+            let alpha_reservoir = ScrapReservoirAlpha::<T>::get(netuid);
+            let (corrected_tao_delta, corrected_alpha_delta, tao_scrap, alpha_scrap) =
+                Self::get_proportional_alpha_tao_and_remainders(
+                    current_sqrt_price,
+                    tao_delta
+                        .saturating_add(TaoCurrency::from(tao_fees))
+                        .saturating_add(tao_reservoir),
+                    alpha_delta
+                        .saturating_add(AlphaCurrency::from(alpha_fees))
+                        .saturating_add(alpha_reservoir),
+                );
+
+            // Update scrap reservoirs
+            ScrapReservoirTao::<T>::insert(netuid, tao_scrap);
+            ScrapReservoirAlpha::<T>::insert(netuid, alpha_scrap);
+
+            // Adjust liquidity
             let maybe_token_amounts = position.to_token_amounts(current_sqrt_price);
             if let Ok((tao, alpha)) = maybe_token_amounts {
                 // Get updated reserves, calculate liquidity
-                let new_tao_reserve = tao
-                    .saturating_add(tao_delta.to_u64())
-                    .saturating_add(tao_fees);
-                let new_alpha_reserve = alpha
-                    .saturating_add(alpha_delta.to_u64())
-                    .saturating_add(alpha_fees);
+                let new_tao_reserve = tao.saturating_add(corrected_tao_delta.to_u64());
+                let new_alpha_reserve = alpha.saturating_add(corrected_alpha_delta.to_u64());
                 let new_liquidity = helpers_128bit::sqrt(
                     (new_tao_reserve as u128).saturating_mul(new_alpha_reserve as u128),
                 ) as u64;

pallets/swap/src/pallet/mod.rs

@@ -147,6 +147,15 @@ mod pallet {
         ValueQuery,
     >;
 
+    /// TAO reservoir for scraps of protocol claimed fees.
+    #[pallet::storage]
+    pub type ScrapReservoirTao<T> = StorageMap<_, Twox64Concat, NetUid, TaoCurrency, ValueQuery>;
+
+    /// Alpha reservoir for scraps of protocol claimed fees.
+    #[pallet::storage]
+    pub type ScrapReservoirAlpha<T> =
+        StorageMap<_, Twox64Concat, NetUid, AlphaCurrency, ValueQuery>;
+
     #[pallet::event]
     #[pallet::generate_deposit(pub(super) fn deposit_event)]
     pub enum Event<T: Config> {

pallets/swap/src/pallet/tests.rs

@@ -2744,3 +2744,172 @@ fn test_clear_protocol_liquidity_green_path() {
         assert!(!SwapV3Initialized::<Test>::contains_key(netuid));
     });
 }
+
+fn as_tuple(
+    (t_used, a_used, t_rem, a_rem): (TaoCurrency, AlphaCurrency, TaoCurrency, AlphaCurrency),
+) -> (u64, u64, u64, u64) {
+    (
+        u64::from(t_used),
+        u64::from(a_used),
+        u64::from(t_rem),
+        u64::from(a_rem),
+    )
+}
+
+#[test]
+fn proportional_when_price_is_one_and_tao_is_plenty() {
+    // sqrt_price = 1.0  => price = 1.0
+    let sqrt = U64F64::from_num(1u64);
+    let amount_tao: TaoCurrency = 10u64.into();
+    let amount_alpha: AlphaCurrency = 3u64.into();
+
+    // alpha * price = 3 * 1 = 3 <= amount_tao(10)
+    let out =
+        Pallet::<Test>::get_proportional_alpha_tao_and_remainders(sqrt, amount_tao, amount_alpha);
+    assert_eq!(as_tuple(out), (3, 3, 7, 0));
+}
+
+#[test]
+fn proportional_when_price_is_one_and_alpha_is_excess() {
+    // sqrt_price = 1.0  => price = 1.0
+    let sqrt = U64F64::from_num(1u64);
+    let amount_tao: TaoCurrency = 5u64.into();
+    let amount_alpha: AlphaCurrency = 10u64.into();
+
+    // tao is limiting: alpha_equiv = floor(5 / 1) = 5
+    let out =
+        Pallet::<Test>::get_proportional_alpha_tao_and_remainders(sqrt, amount_tao, amount_alpha);
+    assert_eq!(as_tuple(out), (5, 5, 0, 5));
+}
+
+#[test]
+fn proportional_with_higher_price_and_alpha_limiting() {
+    // Choose sqrt_price = 2.0 => price = 4.0 (since implementation squares it)
+    let sqrt = U64F64::from_num(2u64);
+    let amount_tao: TaoCurrency = 85u64.into();
+    let amount_alpha: AlphaCurrency = 20u64.into();
+
+    // tao_equivalent = alpha * price = 20 * 4 = 80 < 85 => alpha limits tao
+    // remainders: tao 5, alpha 0
+    let out =
+        Pallet::<Test>::get_proportional_alpha_tao_and_remainders(sqrt, amount_tao, amount_alpha);
+    assert_eq!(as_tuple(out), (80, 20, 5, 0));
+}
+
+#[test]
+fn proportional_with_higher_price_and_tao_limiting() {
+    // Choose sqrt_price = 2.0 => price = 4.0 (since implementation squares it)
+    let sqrt = U64F64::from_num(2u64);
+    let amount_tao: TaoCurrency = 50u64.into();
+    let amount_alpha: AlphaCurrency = 20u64.into();
+
+    // tao_equivalent = alpha * price = 20 * 4 = 80 > 50 => tao limits alpha
+    // alpha_equivalent = floor(50 / 4) = 12
+    // remainders: tao 0, alpha 20 - 12 = 8
+    let out =
+        Pallet::<Test>::get_proportional_alpha_tao_and_remainders(sqrt, amount_tao, amount_alpha);
+    assert_eq!(as_tuple(out), (50, 12, 0, 8));
+}
+
+#[test]
+fn zero_price_uses_no_tao_and_all_alpha() {
+    // sqrt_price = 0 => price = 0
+    let sqrt = U64F64::from_num(0u64);
+    let amount_tao: TaoCurrency = 42u64.into();
+    let amount_alpha: AlphaCurrency = 17u64.into();
+
+    // tao_equivalent = 17 * 0 = 0 <= 42
+    let out =
+        Pallet::<Test>::get_proportional_alpha_tao_and_remainders(sqrt, amount_tao, amount_alpha);
+    assert_eq!(as_tuple(out), (0, 17, 42, 0));
+}
+
+#[test]
+fn rounding_down_behavior_when_dividing_by_price() {
+    // sqrt_price = 2.0 => price = 4.0
+    let sqrt = U64F64::from_num(2u64);
+    let amount_tao: TaoCurrency = 13u64.into();
+    let amount_alpha: AlphaCurrency = 100u64.into();
+
+    // tao is limiting; alpha_equiv = floor(13 / 4) = 3
+    // remainders: tao 0, alpha 100 - 3 = 97
+    let out =
+        Pallet::<Test>::get_proportional_alpha_tao_and_remainders(sqrt, amount_tao, amount_alpha);
+    assert_eq!(as_tuple(out), (13, 3, 0, 97));
+}
+
+#[test]
+fn exact_fit_when_tao_matches_alpha_times_price() {
+    // sqrt_price = 1.0 => price = 1.0
+    let sqrt = U64F64::from_num(1u64);
+    let amount_tao: TaoCurrency = 9u64.into();
+    let amount_alpha: AlphaCurrency = 9u64.into();
+
+    let out =
+        Pallet::<Test>::get_proportional_alpha_tao_and_remainders(sqrt, amount_tao, amount_alpha);
+    assert_eq!(as_tuple(out), (9, 9, 0, 0));
+}
+
+#[test]
+fn handles_zero_balances() {
+    let sqrt = U64F64::from_num(1u64);
+
+    // Zero TAO, some alpha
+    let out =
+        Pallet::<Test>::get_proportional_alpha_tao_and_remainders(sqrt, 0u64.into(), 7u64.into());
+    // tao limits; alpha_equiv = floor(0 / 1) = 0
+    assert_eq!(as_tuple(out), (0, 0, 0, 7));
+
+    // Some TAO, zero alpha
+    let out =
+        Pallet::<Test>::get_proportional_alpha_tao_and_remainders(sqrt, 7u64.into(), 0u64.into());
+    // tao_equiv = 0 * 1 = 0 <= 7
+    assert_eq!(as_tuple(out), (0, 0, 7, 0));
+
+    // Both zero
+    let out =
+        Pallet::<Test>::get_proportional_alpha_tao_and_remainders(sqrt, 0u64.into(), 0u64.into());
+    assert_eq!(as_tuple(out), (0, 0, 0, 0));
+}
+
+#[test]
+fn adjust_protocol_liquidity_uses_and_sets_scrap_reservoirs() {
+    new_test_ext().execute_with(|| {
+        // --- Arrange
+        let netuid: NetUid = 1u16.into();
+        // Price = 1.0 (since sqrt_price^2 = 1), so proportional match is 1:1
+        AlphaSqrtPrice::<Test>::insert(netuid, U64F64::saturating_from_num(1u64));
+
+        // Start with some non-zero scrap reservoirs
+        ScrapReservoirTao::<Test>::insert(netuid, TaoCurrency::from(7u64));
+        ScrapReservoirAlpha::<Test>::insert(netuid, AlphaCurrency::from(5u64));
+
+        // Create a minimal protocol position so the function’s body executes.
+        let protocol = Pallet::<Test>::protocol_account_id();
+        let position = Position::new(
+            PositionId::from(0),
+            netuid,
+            TickIndex::MIN,
+            TickIndex::MAX,
+            0,
+        );
+        // Ensure collect_fees() returns (0,0) via zeroed fees in `position` (default).
+        Positions::<Test>::insert((netuid, protocol, position.id), position.clone());
+
+        // --- Act
+        // No external deltas or fees; only reservoirs should be considered.
+        // With price=1, the exact proportional pair uses 5 alpha and 5 tao,
+        // leaving tao scrap = 7 - 5 = 2, alpha scrap = 5 - 5 = 0.
+        Pallet::<Test>::adjust_protocol_liquidity(netuid, 0u64.into(), 0u64.into());
+
+        // --- Assert: reservoirs were READ (used in proportional calc) and then SET (updated)
+        assert_eq!(
+            ScrapReservoirTao::<Test>::get(netuid),
+            TaoCurrency::from(2u64)
+        );
+        assert_eq!(
+            ScrapReservoirAlpha::<Test>::get(netuid),
+            AlphaCurrency::from(0u64)
+        );
+    });
+}