balanced pair optimization

script/DeployMock.sol

@@ -36,7 +36,7 @@ contract DeployMock is Script {
         int128 _targetSlippage = ABDKMath64x64.divu(1,10000);
         uint256 _feePpm = 100;
 
-        IPartyPool pool = new PartyPool(name, symbol, tokens, _bases, _tradeFrac, _targetSlippage, _feePpm, _feePpm);
+        IPartyPool pool = new PartyPool(name, symbol, tokens, _bases, _tradeFrac, _targetSlippage, _feePpm, _feePpm, false);
 
         // initial mint
         mintAll(address(pool), 10_000);

src/LMSRStabilized.sol

@@ -855,8 +855,7 @@ library LMSRStabilized {
     ) internal pure returns (int128) {
         // Quick sanity checks that decide whether the heterogeneous formula is applicable.
         // If not, fall back to the closed-form equal-asset formula for stability.
-        int128 one = _one();
+        int128 onePlusS = ONE.add(targetSlippage);
-        int128 onePlusS = one.add(targetSlippage);
 
         int128 n64 = ABDKMath64x64.fromUInt(nAssets);
         int128 nMinus1_64 = ABDKMath64x64.fromUInt(nAssets - 1);
@@ -896,8 +895,6 @@ library LMSRStabilized {
         require(f > int128(0), "LMSR: f=0");
         require(f < ONE, "LMSR: f>=1");
 
-        int128 one = _one();
-
         // Top-level input debug
         console2.log("computeBFromSlippageCore: inputs");
         console2.log("q (64.64)");
@@ -921,8 +918,8 @@ library LMSRStabilized {
             //   b = q / y = q * f / (-ln(E))
 
             int128 nMinus1 = ABDKMath64x64.fromUInt(nAssets - 1);
-            int128 numerator = one.sub(targetSlippage.mul(nMinus1)); // 1 - s*(n-1)
+            int128 numerator = ONE.sub(targetSlippage.mul(nMinus1)); // 1 - s*(n-1)
-            int128 denominator = one.add(targetSlippage);             // 1 + s
+            int128 denominator = ONE.add(targetSlippage);             // 1 + s
 
             console2.log("equal-case intermediates:");
             console2.log("numerator = 1 - s*(n-1)");
@@ -937,7 +934,7 @@ library LMSRStabilized {
             console2.logInt(ratio);
 
             // E must be strictly between 0 and 1 for a positive y
-            require(ratio > int128(0) && ratio < one, "LMSR: bad E ratio");
+            require(ratio > int128(0) && ratio < ONE, "LMSR: bad E ratio");
 
             int128 lnE = _ln(ratio); // ln(E) < 0
             console2.log("ln(E)");
@@ -959,7 +956,7 @@ library LMSRStabilized {
             int128 E_sim = _exp(expArg);
             int128 n64 = ABDKMath64x64.fromUInt(nAssets);
             int128 nMinus1_64 = ABDKMath64x64.fromUInt(nAssets - 1);
-            int128 simulatedSlippage = n64.div(nMinus1_64.add(E_sim)).sub(_one());
+            int128 simulatedSlippage = n64.div(nMinus1_64.add(E_sim)).sub(ONE);
             console2.log("simulatedSlippage (using computed b)");
             console2.logInt(simulatedSlippage);
 
@@ -969,7 +966,7 @@ library LMSRStabilized {
             // E = exp(-y * f) where y = q / b
             // and E = (1+s) * (n-1) / (n - (1+s))
             // so y = -ln(E) / f and b = q / y.
-            int128 onePlusS = one.add(targetSlippage);
+            int128 onePlusS = ONE.add(targetSlippage);
 
             console2.log("heterogeneous intermediates:");
             console2.log("onePlusS = 1 + s");
@@ -1007,7 +1004,7 @@ library LMSRStabilized {
 
             // Correct E candidate for the slippage relation:
             // E = (1 - s*(n-1)) / (1 + s)
-            int128 E_candidate = (one.sub(targetSlippage.mul(nMinus1_64))).div(onePlusS);
+            int128 E_candidate = (ONE.sub(targetSlippage.mul(nMinus1_64))).div(onePlusS);
             console2.log("E candidate ((1 - s*(n-1)) / (1+s))");
             console2.logInt(E_candidate);
 
@@ -1031,7 +1028,7 @@ library LMSRStabilized {
             // Simulate slippage using this b to verify
             int128 expArg = y.mul(f).neg();
             int128 E_sim = _exp(expArg);
-            int128 simulatedSlippage = n64.div(nMinus1_64.add(E_sim)).sub(_one());
+            int128 simulatedSlippage = n64.div(nMinus1_64.add(E_sim)).sub(ONE);
             console2.log("simulatedSlippage (heterogeneous)");
             console2.logInt(simulatedSlippage);
 
@@ -1055,7 +1052,7 @@ library LMSRStabilized {
     }
 
     /// @notice Compute M (shift) and Z (sum of exponentials) dynamically
-    function _computeMAndZ(int128 b, int128[] memory qInternal) private pure returns (int128 M, int128 Z) {
+    function _computeMAndZ(int128 b, int128[] memory qInternal) internal pure returns (int128 M, int128 Z) {
         require(qInternal.length > 0, "LMSR: no assets");
 
         // Precompute reciprocal of b to replace divisions with multiplications in the loop
@@ -1083,7 +1080,7 @@ library LMSRStabilized {
     }
 
     /// @notice Compute all e[i] = exp(z[i]) values dynamically
-    function _computeE(int128 b, int128[] memory qInternal, int128 M) private pure returns (int128[] memory e) {
+    function _computeE(int128 b, int128[] memory qInternal, int128 M) internal pure returns (int128[] memory e) {
         uint len = qInternal.length;
         e = new int128[](len);
 
@@ -1100,7 +1097,7 @@ library LMSRStabilized {
 
     /// @notice Compute r0 = e_i / e_j directly as exp((q_i - q_j) / b)
     /// This avoids computing two separate exponentials and a division
-    function _computeR0(int128 b, int128[] memory qInternal, uint256 i, uint256 j) private pure returns (int128) {
+    function _computeR0(int128 b, int128[] memory qInternal, uint256 i, uint256 j) internal pure returns (int128) {
         return _exp(qInternal[i].sub(qInternal[j]).div(b));
     }
 
@@ -1110,16 +1107,15 @@ library LMSRStabilized {
        -------------------- */
 
     // Precomputed Q64.64 representation of 1.0 (1 << 64).
-    int128 private constant ONE = 0x10000000000000000;
+    int128 internal constant ONE = 0x10000000000000000;
     // Precomputed Q64.64 representation of 32.0 for exp guard
-    int128 private constant EXP_LIMIT = 0x200000000000000000;
+    int128 internal constant EXP_LIMIT = 0x200000000000000000;
 
-    function _exp(int128 x) private pure returns (int128) { return ABDKMath64x64.exp(x); }
+    function _exp(int128 x) internal pure returns (int128) { return ABDKMath64x64.exp(x); }
-    function _ln(int128 x) private pure returns (int128)  { return ABDKMath64x64.ln(x); }
+    function _ln(int128 x) internal pure returns (int128)  { return ABDKMath64x64.ln(x); }
-    function _one() private pure returns (int128) { return ONE; }
 
     /// @notice Compute size metric S(q) = sum of all asset quantities
-    function _computeSizeMetric(int128[] memory qInternal) private pure returns (int128) {
+    function _computeSizeMetric(int128[] memory qInternal) internal pure returns (int128) {
         int128 total = int128(0);
         for (uint i = 0; i < qInternal.length; ) {
             total = total.add(qInternal[i]);
@@ -1129,7 +1125,7 @@ library LMSRStabilized {
     }
 
     /// @notice Compute b from kappa and current asset quantities
-    function _computeB(State storage s) private view returns (int128) {
+    function _computeB(State storage s) internal view returns (int128) {
         int128 sizeMetric = _computeSizeMetric(s.qInternal);
         require(sizeMetric > int128(0), "LMSR: size metric zero");
         return s.kappa.mul(sizeMetric);

src/LMSRStabilizedBalancedPair.sol

@@ -0,0 +1,261 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.30;
+
+import "forge-std/console2.sol";
+import "@abdk/ABDKMath64x64.sol";
+import "./LMSRStabilized.sol";
+
+/// @notice Specialized functions for the 2-asset stablecoin case
+library LMSRStabilizedBalancedPair {
+    using ABDKMath64x64 for int128;
+
+    // Precomputed Q64.64 representation of 1.0 (1 << 64).
+    int128 private constant ONE = 0x10000000000000000;
+
+    /// @notice Specialized 2-asset balanced approximation of swapAmountsForExactInput.
+    /// - Assumes exactly two assets and that the two assets' internal balances are within ~1% of parity.
+    /// - Implements a gas-optimized two-tier Taylor approximation to avoid most exp()/ln() calls:
+    ///     * Tier 1 (quadratic, cheapest): for small u = a/b (u <= 0.1) we compute
+    ///         X = u*(1 + δ) - u^2/2
+    ///         ln(1+X) ≈ X - X^2/2
+    ///       and return amountOut ≈ b * lnApprox. This Horner-style form minimizes multiplies/divides
+    ///       and temporaries compared to the earlier a^2/a^3 expansion.
+    ///     * Tier 2 (cubic correction): for moderate u (0.1 < u <= 0.5) we add the X^3/3 term:
+    ///         ln(1+X) ≈ X - X^2/2 + X^3/3
+    ///       which improves accuracy while still being significantly cheaper than full exp/ln.
+    /// - For cases where |δ| (the per-asset imbalance scaled by b) or u are outside the safe ranges,
+    ///   or when limitPrice handling cannot be reliably approximated, the function falls back to the
+    ///   numerically-exact swapAmountsForExactInput(...) implementation to preserve correctness.
+    /// - The goal is to keep relative error well below 0.001% in the intended small-u, near-parity regime,
+    ///   while substantially reducing gas in the common fast path.
+    function swapAmountsForExactInput(
+        LMSRStabilized.State storage s,
+        uint256 i,
+        uint256 j,
+        int128 a,
+        int128 limitPrice
+    ) internal view returns (int128 amountIn, int128 amountOut) {
+        // Quick index check
+        require(i < s.nAssets && j < s.nAssets, "LMSR: idx");
+
+        // If not exactly a two-asset pool, fall back to the general routine.
+        if (s.nAssets != 2) {
+            console2.log('balanced2: fallback nAssets!=n2');
+            return LMSRStabilized.swapAmountsForExactInput(s, i, j, a, limitPrice);
+        }
+
+        // Compute b and inverse early (needed to evaluate delta and limit-price)
+        int128 b = LMSRStabilized._computeB(s);
+        // Guard: if b not positive, fallback to exact implementation (will revert there if necessary)
+        if (!(b > int128(0))) {
+            console2.log("balanced2: fallback b<=0");
+            return LMSRStabilized.swapAmountsForExactInput(s, i, j, a, limitPrice);
+        }
+        int128 invB = ABDKMath64x64.div(ONE, b);
+
+        // Small-signal delta = (q_i - q_j) / b (used to approximate r0 = exp(delta))
+        int128 delta = s.qInternal[i].sub(s.qInternal[j]).mul(invB);
+
+        // If a positive limitPrice is given, attempt a 2-asset near-parity polynomial solution
+        if (limitPrice > int128(0)) {
+            console2.log("balanced2: handling limitPrice via small-delta approx");
+
+            // Approximate r0 = exp(delta) using Taylor: 1 + δ + δ^2/2 + δ^3/6
+            int128 delta_sq = delta.mul(delta);
+            int128 delta_cu = delta_sq.mul(delta);
+            int128 r0_approx = ONE
+                .add(delta)
+                .add(delta_sq.div(ABDKMath64x64.fromUInt(2)))
+                .add(delta_cu.div(ABDKMath64x64.fromUInt(6)));
+
+            console2.log("r0_approx:");
+            console2.logInt(r0_approx);
+
+            // If limitPrice <= r0 (current price) we must revert (same semantic as original)
+            if (limitPrice <= r0_approx) {
+                console2.log("balanced2: limitPrice <= r0_approx -> revert");
+                revert("LMSR: limitPrice <= current price");
+            }
+
+            // Ratio = limitPrice / r0_approx
+            int128 ratio = limitPrice.div(r0_approx);
+            console2.log("limitPrice/r0_approx:");
+            console2.logInt(ratio);
+
+            // x = ratio - 1; use Taylor for ln(1+x) when |x| is small
+            int128 x = ratio.sub(ONE);
+            int128 absX = x >= int128(0) ? x : x.neg();
+
+            // Acceptable range for ln Taylor approx: |x| <= 0.1 (conservative)
+            int128 X_MAX = ABDKMath64x64.divu(1, 10); // 0.1
+            if (absX > X_MAX) {
+                // Too large to safely approximate; fall back to exact computation
+                console2.log("balanced2: fallback limitPrice ratio too far from 1");
+                return LMSRStabilized.swapAmountsForExactInput(s, i, j, a, limitPrice);
+            }
+
+            // ln(1+x) ≈ x - x^2/2 + x^3/3
+            int128 x_sq = x.mul(x);
+            int128 x_cu = x_sq.mul(x);
+            int128 lnRatioApprox = x
+                .sub(x_sq.div(ABDKMath64x64.fromUInt(2)))
+                .add(x_cu.div(ABDKMath64x64.fromUInt(3)));
+
+            console2.log("lnRatioApprox (64x64):");
+            console2.logInt(lnRatioApprox);
+
+            // aLimitOverB = ln(limitPrice / r0) approximated
+            int128 aLimitOverB = lnRatioApprox;
+
+            // Must be > 0; otherwise fall back
+            if (!(aLimitOverB > int128(0))) {
+                console2.log("balanced2: fallback non-positive aLimitOverB");
+                return LMSRStabilized.swapAmountsForExactInput(s, i, j, a, limitPrice);
+            }
+
+            // aLimit = b * aLimitOverB (in Q64.64)
+            int128 aLimit64 = b.mul(aLimitOverB);
+            console2.log("aLimit64 (64x64):");
+            console2.logInt(aLimit64);
+
+            // If computed aLimit is less than requested a, use the truncated value.
+            if (aLimit64 < a) {
+                console2.log("balanced2: truncating input a to aLimit64 due to limitPrice");
+                console2.log("original a:");
+                console2.logInt(a);
+                console2.log("truncated aLimit64:");
+                console2.logInt(aLimit64);
+                a = aLimit64;
+            } else {
+                console2.log("balanced2: limitPrice does not truncate input");
+            }
+
+            // Note: after potential truncation we continue with the polynomial approximation below
+        }
+
+        // Debug: entry trace
+        console2.log("balanced2: enter");
+        console2.log("i", i);
+        console2.log("j", j);
+        console2.log("nAssets", s.nAssets);
+        console2.log("a (64x64):");
+        console2.logInt(a);
+        console2.log("b (64x64):");
+        console2.logInt(b);
+        console2.log("invB (64x64):");
+        console2.logInt(invB);
+
+        // Small-signal delta already computed above; reuse it
+        int128 absDelta = delta >= int128(0) ? delta : delta.neg();
+
+        console2.log("delta (q_i - q_j)/b:");
+        console2.logInt(delta);
+        console2.log("absDelta:");
+        console2.logInt(absDelta);
+
+        // Allow balanced pools only: require |delta| <= 1% (approx ln(1.01) ~ 0.00995; we use conservative 0.01)
+        int128 DELTA_MAX = ABDKMath64x64.divu(1, 100); // 0.01
+        if (absDelta > DELTA_MAX) {
+            // Not balanced within 1% -> use exact routine
+            console2.log("balanced2: fallback delta too large");
+            return LMSRStabilized.swapAmountsForExactInput(s, i, j, a, limitPrice);
+        }
+
+        // Scaled input u = a / b (Q64.64). For polynomial approximation we require moderate u.
+        int128 u = a.mul(invB);
+        if (u <= int128(0)) {
+            // Non-positive input -> behave like exact implementation (will revert if invalid)
+            console2.log("balanced2: fallback u<=0");
+            return LMSRStabilized.swapAmountsForExactInput(s, i, j, a, limitPrice);
+        }
+
+        console2.log("u = a/b (64x64):");
+        console2.logInt(u);
+
+        // Restrict to a conservative polynomial radius for accuracy; fallback otherwise.
+        // We choose u <= 0.5 (0.5 in Q64.64) as safe for cubic approximation in typical parameters.
+        int128 U_MAX = ABDKMath64x64.divu(1, 2); // 0.5
+        if (u > U_MAX) {
+            console2.log("balanced2: fallback u too large");
+            return LMSRStabilized.swapAmountsForExactInput(s, i, j, a, limitPrice);
+        }
+
+        // Now compute a two-tier approximation using Horner-style evaluation to reduce mul/divs.
+        // Primary tier (cheap quadratic): accurate for small u = a/b.
+        // Secondary tier (cubic correction): used when u is moderate but still within U_MAX.
+        int128 one = ONE;
+        int128 HALF = ABDKMath64x64.divu(1, 2); // 0.5
+        int128 THIRD = ABDKMath64x64.divu(1, 3); // ~0.333...
+
+        // Precomputed thresholds
+        int128 U_TIER1 = ABDKMath64x64.divu(1, 10); // 0.1 -> cheap quadratic tier
+        int128 U_MAX_LOCAL = ABDKMath64x64.divu(1, 2); // 0.5 -> still allowed cubic tier
+
+        // u is already computed above
+        // Compute X = u*(1 + delta) - u^2/2
+        int128 u2 = u.mul(u);
+        int128 X = u.mul(one.add(delta)).sub(u2.div(ABDKMath64x64.fromUInt(2)));
+
+        // Compute X^2 once
+        int128 X2 = X.mul(X);
+
+        int128 lnApprox;
+        if (u <= U_TIER1) {
+            // Cheap quadratic ln(1+X) ≈ X - X^2/2
+            lnApprox = X.sub(X2.div(ABDKMath64x64.fromUInt(2)));
+            console2.log("balanced2: using tier1 quadratic approx");
+        } else if (u <= U_MAX_LOCAL) {
+            // Secondary cubic correction: ln(1+X) ≈ X - X^2/2 + X^3/3
+            int128 X3 = X2.mul(X);
+            lnApprox = X.sub(X2.div(ABDKMath64x64.fromUInt(2))).add(X3.div(ABDKMath64x64.fromUInt(3)));
+            console2.log("balanced2: using tier2 cubic approx");
+        } else {
+            // u beyond allowed range - fallback
+            console2.log("balanced2: fallback u too large for approximation");
+            return LMSRStabilized.swapAmountsForExactInput(s, i, j, a, limitPrice);
+        }
+
+        console2.log("lnApprox (64x64):");
+        console2.logInt(lnApprox);
+
+        int128 approxOut = b.mul(lnApprox);
+
+        console2.log("approxOut (64x64):");
+        console2.logInt(approxOut);
+
+        // Safety sanity: approximation must be > 0
+        if (approxOut <= int128(0)) {
+            console2.log("balanced2: fallback approxOut <= 0");
+            return LMSRStabilized.swapAmountsForExactInput(s, i, j, a, limitPrice);
+        }
+
+        // Cap to available j balance: if approximated output exceeds q_j, it's likely approximation break;
+        // fall back to the exact solver to handle capping/edge cases.
+        int128 qj64 = s.qInternal[j];
+        console2.log("qj64 (64x64):");
+        console2.logInt(qj64);
+        if (approxOut >= qj64) {
+            console2.log("balanced2: fallback approxOut >= qj");
+            return LMSRStabilized.swapAmountsForExactInput(s, i, j, a, limitPrice);
+        }
+
+        // Everything looks fine; return approximated amountOut and used amountIn (a)
+        amountIn = a;
+        amountOut = approxOut;
+
+        console2.log("balanced2: returning approx results");
+        console2.log("amountIn (64x64):");
+        console2.logInt(amountIn);
+        console2.log("amountOut (64x64):");
+        console2.logInt(amountOut);
+
+        // Final guard: ensure output is sensible and not NaN-like (rely on positivity checks above)
+        if (amountOut < int128(0)) {
+            console2.log("balanced2: fallback final guard amountOut<0");
+            return LMSRStabilized.swapAmountsForExactInput(s, i, j, a, limitPrice);
+        }
+
+        return (amountIn, amountOut);
+    }
+
+}

src/PartyPool.sol

@@ -8,6 +8,7 @@ import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@openzeppelin/contracts/utils/ReentrancyGuard.sol";
 import "./LMSRStabilized.sol";
+import "./LMSRStabilizedBalancedPair.sol";
 import "./IPartyPool.sol";
 import "./IPartyFlashCallback.sol";
 
@@ -46,6 +47,7 @@ contract PartyPool is IPartyPool, ERC20, ReentrancyGuard {
     //
 
     LMSRStabilized.State internal lmsr;
+    bool immutable private _stablePair; // if true, the optimized LMSRStabilizedBalancedPair optimization path is enabled
 
     // Cached on-chain balances (uint) and internal 64.64 representation
     // balance / base = internal
@@ -65,6 +67,7 @@ contract PartyPool is IPartyPool, ERC20, ReentrancyGuard {
     /// @param _targetSlippage target slippage in 64.64 fixed-point (as used by LMSR)
     /// @param _swapFeePpm fee in parts-per-million, taken from swap input amounts before LMSR calculations
     /// @param _flashFeePpm fee in parts-per-million, taken for flash loans
+    /// @param _stable if true and assets.length==2, then the optimization for 2-asset stablecoin pools is activated.
     constructor(
         string memory name_,
         string memory symbol_,
@@ -73,7 +76,8 @@ contract PartyPool is IPartyPool, ERC20, ReentrancyGuard {
         int128 _tradeFrac,
         int128 _targetSlippage,
         uint256 _swapFeePpm,
-        uint256 _flashFeePpm
+        uint256 _flashFeePpm,
+        bool _stable
     ) ERC20(name_, symbol_) {
         require(_tokens.length > 1, "Pool: need >1 asset");
         require(_tokens.length == _bases.length, "Pool: lengths mismatch");
@@ -85,6 +89,7 @@ contract PartyPool is IPartyPool, ERC20, ReentrancyGuard {
         swapFeePpm = _swapFeePpm;
         require(_flashFeePpm < 1_000_000, "Pool: flash fee >= ppm");
         flashFeePpm = _flashFeePpm;
+        _stablePair = _stable && _tokens.length == 2;
 
         uint256 n = _tokens.length;
 
@@ -369,7 +374,11 @@ contract PartyPool is IPartyPool, ERC20, ReentrancyGuard {
         require(deltaInternalI > int128(0), "swap: input too small after fee");
 
         // Compute internal amounts using LMSR (exact-input with price limit)
-        (amountInInternalUsed, amountOutInternal) = lmsr.swapAmountsForExactInput(i, j, deltaInternalI, limitPrice);
+        // if _stablePair is true, use the optimized path
+        console2.log('stablepair optimization?', _stablePair);
+        (amountInInternalUsed, amountOutInternal) =
+            _stablePair ? LMSRStabilizedBalancedPair.swapAmountsForExactInput(lmsr, i, j, deltaInternalI, limitPrice)
+                        : lmsr.swapAmountsForExactInput(i, j, deltaInternalI, limitPrice);
 
         // Convert actual used input internal -> uint (ceil)
         amountInUintNoFee = _internalToUintCeil(amountInInternalUsed, bases[i]);

test/GasTest.sol

@@ -162,7 +162,42 @@ contract GasTest is Test {
         // Deploy pool with a small fee to test fee-handling paths (use 1000 ppm = 0.1%)
         uint256 feePpm = 1000;
         string memory poolName = string(abi.encodePacked("LP", vm.toString(numTokens)));
-        PartyPool newPool = new PartyPool(poolName, poolName, tokens, bases, tradeFrac, targetSlippage, feePpm, feePpm);
+        PartyPool newPool = new PartyPool(poolName, poolName, tokens, bases, tradeFrac, targetSlippage, feePpm, feePpm, false);
+
+        // Transfer initial deposit amounts into pool before initial mint
+        for (uint256 i = 0; i < numTokens; i++) {
+            TestERC20(tokens[i]).transfer(address(newPool), INIT_BAL);
+        }
+
+        // Perform initial mint (initial deposit); receiver is this contract
+        newPool.mint(address(0), address(this), 0, 0);
+
+        return newPool;
+    }
+
+    /// @notice Helper to create a pool with the stable-pair optimization enabled
+    function createPoolStable(uint256 numTokens) internal returns (PartyPool) {
+        // Deploy tokens dynamically
+        address[] memory tokens = new address[](numTokens);
+        uint256[] memory bases = new uint256[](numTokens);
+
+        for (uint256 i = 0; i < numTokens; i++) {
+            string memory name = string(abi.encodePacked("T", vm.toString(i)));
+            TestERC20 token = new TestERC20(name, name, 0);
+            tokens[i] = address(token);
+            bases[i] = BASE;
+
+            // Mint initial balances for pool initialization and test users
+            token.mint(address(this), INIT_BAL);
+            token.mint(alice, INIT_BAL);
+            token.mint(bob, INIT_BAL);
+        }
+
+        // Deploy pool with a small fee to test fee-handling paths (use 1000 ppm = 0.1%)
+        uint256 feePpm = 1000;
+        string memory poolName = string(abi.encodePacked("LPs", vm.toString(numTokens)));
+        // Note the final 'true' arg to activate stable-pair optimization path
+        PartyPool newPool = new PartyPool(poolName, poolName, tokens, bases, tradeFrac, targetSlippage, feePpm, feePpm, true);
 
         // Transfer initial deposit amounts into pool before initial mint
         for (uint256 i = 0; i < numTokens; i++) {
@@ -253,6 +288,24 @@ contract GasTest is Test {
         _performSwapGasTest(pool50);
     }
 
+    /// @notice Gas measurement: perform 10 swaps back-and-forth on a 2-token stable pair (stable-path enabled)
+    function testSwapGasStablePair() public {
+        PartyPool stablePair = createPoolStable(2);
+        _performSwapGasTest(stablePair);
+    }
+
+    /// @notice Gas-style test: alternate swapMint then burnSwap on a 2-token stable pair
+    function testSwapMintBurnSwapGasStablePair() public {
+        PartyPool stablePair = createPoolStable(2);
+        _performSwapMintBurnSwapGasTest(stablePair);
+    }
+
+    /// @notice Combined gas test (mint then burn) on 2-token stable pair using mint() and burn().
+    function testMintBurnGasStablePair() public {
+        PartyPool stablePair = createPoolStable(2);
+        _performMintBurnGasTest(stablePair);
+    }
+
     /// @notice Helper function: alternate swapMint then burnSwap to keep pool size roughly stable.
     function _performSwapMintBurnSwapGasTest(PartyPool testPool) internal {
         uint256 iterations = 10;

test/LMSRStabilized.t.sol

@@ -5,6 +5,7 @@ import "forge-std/Test.sol";
 import "forge-std/console.sol";
 import "@openzeppelin/contracts/interfaces/IERC20Metadata.sol";
 import "../src/LMSRStabilized.sol";
+import "../src/LMSRStabilizedBalancedPair.sol";
 
 
 /// @notice Forge tests for LMSRStabilized
@@ -253,6 +254,96 @@ contract LMSRStabilizedTest is Test {
         assertLt(relativeError, ABDKMath64x64.divu(1, 100), "Slippage target not met after deposit");
     }
 
+    /// @notice Test balanced2 handling of limitPrice that causes truncation of input a
+    function testBalanced2LimitTruncation() public {
+        // Two-asset balanced pool
+        int128[] memory q = new int128[](2);
+        q[0] = ABDKMath64x64.fromUInt(1_000_000);
+        q[1] = ABDKMath64x64.fromUInt(1_000_000);
+        s.init(q, stdTradeSize, stdSlippage);
+
+        // Compute b for constructing meaningful a and limits
+        int128 b = _computeB(q);
+
+        // Choose a large requested input so that the limitPrice will truncate it.
+        int128 aRequested = b.mul(ABDKMath64x64.fromUInt(10)); // a/b = 10 (within EXP_LIMIT)
+
+        // Small limit slightly above current price (r0 == 1). Use 0.1% above parity.
+        int128 limitPrice = ABDKMath64x64.fromInt(1).add(ABDKMath64x64.divu(1, 1000)); // 1.001
+
+        // Call optimized balanced2 and exact versions
+        (int128 inApprox, int128 outApprox) = LMSRStabilizedBalancedPair.swapAmountsForExactInput(s, 0, 1, aRequested, limitPrice);
+        (int128 inExact, int128 outExact) = s.swapAmountsForExactInput(0, 1, aRequested, limitPrice);
+
+        // Ensure exact returned something sensible
+        assertTrue(outExact > 0, "exact output should be positive");
+
+        // Relative error tolerance 0.001% (1e-5)
+        int128 relErr = (outApprox.sub(outExact)).abs().div(outExact);
+        int128 tol = ABDKMath64x64.divu(1, 100_000);
+        assertLt(relErr, tol, "balanced2 truncated output deviates from exact beyond tolerance");
+
+        // Input used should be close as well
+        int128 inRelErr = (inApprox.sub(inExact)).abs();
+        // If exact truncated, inExact likely equals aLimit computed by ln; allow small absolute difference tolerance of 1e-6 relative to b
+        int128 absTol = b.div(ABDKMath64x64.fromUInt(1_000_000)); // b * 1e-6
+        assertTrue(inRelErr <= absTol, "balanced2 truncated input differs from exact beyond small absolute tolerance");
+    }
+
+    /// @notice Test balanced2 with a limitPrice that does not truncate the provided input
+    function testBalanced2LimitNoTruncation() public {
+        // Two-asset balanced pool
+        int128[] memory q = new int128[](2);
+        q[0] = ABDKMath64x64.fromUInt(1_000_000);
+        q[1] = ABDKMath64x64.fromUInt(1_000_000);
+        s.init(q, stdTradeSize, stdSlippage);
+
+        // Small input a
+        int128 a = q[0].mul(ABDKMath64x64.divu(1, 1000)); // 0.1% of asset
+
+        // Very relaxed limit (2x current price) which should not truncate
+        int128 limitPrice = ABDKMath64x64.fromUInt(2);
+
+        (int128 inApprox, int128 outApprox) = LMSRStabilizedBalancedPair.swapAmountsForExactInput(s, 0, 1, a, limitPrice);
+        (int128 inExact, int128 outExact) = s.swapAmountsForExactInput(0, 1, a, limitPrice);
+
+        // Exact outputs must be positive
+        assertTrue(outExact > 0, "exact output should be positive");
+
+        // Expect almost exact match when no truncation occurs; use tight tolerance
+        int128 relErr = (outApprox.sub(outExact)).abs().div(outExact);
+        int128 tol = ABDKMath64x64.divu(1, 100_000); // 0.001%
+        assertLt(relErr, tol, "balanced2 no-truncate output deviates from exact beyond tolerance");
+
+        // AmountIn should equal provided a for both functions
+        assertEq(inApprox, a, "balanced2 should use full input when not truncated");
+        assertEq(inExact, a, "exact should use full input when not truncated");
+    }
+
+    /// @notice Test that balanced2 reverts when limitPrice <= current price (no partial fill allowed)
+    function testBalanced2LimitRevertWhenAtOrBelowCurrent() public {
+        // Two-asset balanced pool
+        int128[] memory q = new int128[](2);
+        q[0] = ABDKMath64x64.fromUInt(1_000_000);
+        q[1] = ABDKMath64x64.fromUInt(1_000_000);
+        s.init(q, stdTradeSize, stdSlippage);
+
+        int128 limitPrice = ABDKMath64x64.fromInt(1); // equal to current price
+
+        vm.expectRevert(bytes("LMSR: limitPrice <= current price"));
+        this._swapAmountsForExactInput_balanced2(0, 1, q[0].mul(ABDKMath64x64.divu(1, 1000)), limitPrice);
+    }
+
+    function _swapAmountsForExactInput_balanced2(
+        uint256 i,
+        uint256 j,
+        int128 a,
+        int128 limitPrice
+    ) external view returns (int128 amountIn, int128 amountOut) {
+        return LMSRStabilizedBalancedPair.swapAmountsForExactInput(s, i,j,a,limitPrice);
+    }
+
+
     function testRescalingAfterWithdrawal() public {
         // Initialize pool with almost balanced assets
         initAlmostBalanced();
@@ -858,4 +949,116 @@ contract LMSRStabilizedTest is Test {
         assertTrue(payout > 0, "payout must be positive even when one asset is zero");
     }
 
+    /// @notice Test that the balanced2 polynomial approximation is accurate for a two-asset balanced pool
+    function testBalanced2ApproxAccuracy() public {
+        // Create a minimal two-asset balanced pool
+        int128[] memory q = new int128[](2);
+        q[0] = ABDKMath64x64.fromUInt(1_000_000);
+        q[1] = ABDKMath64x64.fromUInt(1_000_000);
+        s.init(q, stdTradeSize, stdSlippage);
+
+        // Small trade (well within u <= 0.5 and delta <= 1%)
+        int128 a = q[0].mul(ABDKMath64x64.divu(1, 1000)); // 0.1% of asset
+
+        // Compute approx and exact
+        (int128 inApprox, int128 outApprox) = LMSRStabilizedBalancedPair.swapAmountsForExactInput(s, 0, 1, a, 0);
+        (int128 inExact, int128 outExact) = s.swapAmountsForExactInput(0, 1, a, 0);
+
+        // Sanity
+        assertTrue(outExact > 0, "Exact output should be positive");
+
+        // Relative error: |approx - exact| / exact
+        int128 relErr = (outApprox.sub(outExact)).abs().div(outExact);
+
+        // Require relative error < 0.001% (1e-5) => expressed as 1 / 100_000
+        int128 tolerance = ABDKMath64x64.divu(1, 100_000);
+        assertLt(relErr, tolerance, "balanced2 approximation relative error too large");
+
+        // AmountIn should equal requested a (no truncation)
+        assertEq(inApprox, a, "balanced2 approximation should use full input when no limitPrice");
+        assertEq(inExact, a, "exact computation should use full input when no limitPrice");
+    }
+
+    /// @notice Test that when the parity assumption is violated, the balanced2 helper falls back
+    /// to the exact implementation (we expect identical outputs).
+    function testBalanced2FallbackWhenParityViolated() public {
+        // Start with two-asset balanced pool (we'll mutate it)
+        int128[] memory q = new int128[](2);
+        q[0] = ABDKMath64x64.fromUInt(1_000_000);
+        q[1] = ABDKMath64x64.fromUInt(1_000_000);
+        s.init(q, stdTradeSize, stdSlippage);
+
+        // Prepare newQ starting from equal quantities; we'll grow q0 until delta > DELTA_MAX
+        int128[] memory newQ = new int128[](2);
+        newQ[0] = q[0];
+        newQ[1] = q[1];
+
+        // DELTA_MAX used by the library: 0.01
+        int128 DELTA_MAX = ABDKMath64x64.divu(1, 100);
+
+        // Iteratively increase q0 until the library's delta = (q0 - q1) / b > DELTA_MAX
+        // We cap iterations to avoid infinite loops in pathological cases.
+        bool reached = false;
+        for (uint iter = 0; iter < 64; iter++) {
+            // Update the state's cached qInternal with the candidate imbalance
+            _updateCachedQInternal(newQ);
+
+            // Compute the current b and delta using the state's parameters
+            int128 bNow = _computeB(); // uses s.qInternal and s.kappa
+            // avoid division by zero
+            if (bNow == int128(0)) { break; }
+
+            int128 deltaNow = newQ[0].sub(newQ[1]).div(bNow);
+            if (deltaNow < int128(0)) { deltaNow = deltaNow.neg(); }
+
+            if (deltaNow > DELTA_MAX) {
+                reached = true;
+                break;
+            }
+
+            // Increase q0 by 10% each iteration to quickly cross the threshold
+            newQ[0] = newQ[0].mul(ABDKMath64x64.fromUInt(11)).div(ABDKMath64x64.fromUInt(10));
+        }
+
+        // Ensure we actually achieved the desired imbalance for a meaningful test
+        _updateCachedQInternal(newQ);
+        int128 finalB = _computeB();
+        int128 finalDelta = newQ[0].sub(newQ[1]).div(finalB);
+        if (finalDelta < int128(0)) finalDelta = finalDelta.neg();
+        assertTrue(finalDelta > DELTA_MAX, "failed to create delta > DELTA_MAX in test");
+
+        // Small trade amount
+        int128 a = newQ[0].mul(ABDKMath64x64.divu(1, 1000));
+
+        // Call both functions; balanced2 should detect parity violation and fall back to exact
+        (int128 inApprox, int128 outApprox) = LMSRStabilizedBalancedPair.swapAmountsForExactInput(s, 0, 1, a, 0);
+        (int128 inExact, int128 outExact) = s.swapAmountsForExactInput(0, 1, a, 0);
+
+        // Because parity assumption is violated balanced2 should fall back to exact implementation
+        assertEq(inApprox, inExact, "fallback should return identical amountIn");
+        assertEq(outApprox, outExact, "fallback should return identical amountOut");
+    }
+
+    /// @notice Test that the balanced2 helper falls back when scaled input u = a/b is too large
+    function testBalanced2FallbackOnLargeInput() public {
+        // Two-asset balanced pool
+        int128[] memory q = new int128[](2);
+        q[0] = ABDKMath64x64.fromUInt(1_000_000);
+        q[1] = ABDKMath64x64.fromUInt(1_000_000);
+        s.init(q, stdTradeSize, stdSlippage);
+
+        // Compute b
+        int128 b = _computeB(q);
+
+        // Choose a so that u = a / b = 0.75 (> 0.5 U_MAX)
+        int128 a = b.mul(ABDKMath64x64.divu(3, 4)); // a/b = 0.75
+
+        // Call both functions and expect fallback (identical results)
+        (int128 inApprox, int128 outApprox) = LMSRStabilizedBalancedPair.swapAmountsForExactInput(s, 0, 1, a, 0);
+        (int128 inExact, int128 outExact) = s.swapAmountsForExactInput(0, 1, a, 0);
+
+        assertEq(inApprox, inExact, "fallback on large input should return identical amountIn");
+        assertEq(outApprox, outExact, "fallback on large input should return identical amountOut");
+    }
+
 }

test/PartyPool.t.sol

@@ -195,7 +195,7 @@ contract PartyPoolTest is Test {
         // Deploy pool with a small fee to test fee-handling paths (use 1000 ppm = 0.1%)
         uint256 feePpm = 1000;
 
-        pool = new PartyPool("LP", "LP", tokens, bases, tradeFrac, targetSlippage, feePpm, feePpm);
+        pool = new PartyPool("LP", "LP", tokens, bases, tradeFrac, targetSlippage, feePpm, feePpm, false);
 
         // Transfer initial deposit amounts into pool before initial mint (pool expects tokens already in contract)
         // We deposit equal amounts INIT_BAL for each token
@@ -224,7 +224,7 @@ contract PartyPoolTest is Test {
             bases10[i] = BASE;
         }
 
-        pool10 = new PartyPool("LP10", "LP10", tokens10, bases10, tradeFrac, targetSlippage, feePpm, feePpm);
+        pool10 = new PartyPool("LP10", "LP10", tokens10, bases10, tradeFrac, targetSlippage, feePpm, feePpm, false);
 
         // Mint additional tokens for pool10 initial deposit
         token0.mint(address(this), INIT_BAL);