PartyPoolBalancedPair as subclass

src/Deploy.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.30;
 import {IERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
 import {PartyPlanner} from "./PartyPlanner.sol";
 import {PartyPool} from "./PartyPool.sol";
+import {PartyPoolBalancedPair} from "./PartyPoolBalancedPair.sol";
 import {PartyPoolMintImpl} from "./PartyPoolMintImpl.sol";
 import {PartyPoolSwapImpl} from "./PartyPoolSwapImpl.sol";
 
@@ -32,7 +33,21 @@ library Deploy {
         uint256 protocolFeePpm = 0;
         address protocolAddr = address(0);
 
-        return new PartyPool(
+        return _stable && tokens_.length == 2 ?
+        new PartyPoolBalancedPair(
+            name_,
+            symbol_,
+            tokens_,
+            bases_,
+            _kappa,
+            _swapFeePpm,
+            _flashFeePpm,
+            protocolFeePpm,
+            protocolAddr,
+            new PartyPoolSwapImpl(),
+            new PartyPoolMintImpl()
+        ) :
+        new PartyPool(
             name_,
             symbol_,
             tokens_,
@@ -42,7 +57,6 @@ library Deploy {
             _flashFeePpm,
             protocolFeePpm,
             protocolAddr,
-            _stable,
             new PartyPoolSwapImpl(),
             new PartyPoolMintImpl()
         );

src/PartyPlanner.sol

@@ -8,6 +8,7 @@ import {LMSRStabilized} from "./LMSRStabilized.sol";
 import {PartyPool} from "./PartyPool.sol";
 import {PartyPoolMintImpl} from "./PartyPoolMintImpl.sol";
 import {PartyPoolSwapImpl} from "./PartyPoolSwapImpl.sol";
+import {PartyPoolBalancedPair} from "./PartyPoolBalancedPair.sol";
 
 /// @title PartyPlanner
 /// @notice Factory contract for creating and tracking PartyPool instances
@@ -86,7 +87,21 @@ contract PartyPlanner is IPartyPlanner {
         require(_kappa > int128(0), "Planner: kappa must be > 0");
 
         // Create a new PartyPool instance (kappa-based constructor)
-        pool = new PartyPool(
+        pool = _stable && _tokens.length == 2 ?
+        new PartyPoolBalancedPair(
+            name_,
+            symbol_,
+            _tokens,
+            _bases,
+            _kappa,
+            _swapFeePpm,
+            _flashFeePpm,
+            PROTOCOL_FEE_PPM,
+            PROTOCOL_FEE_ADDRESS,
+            PartyPoolSwapImpl(SWAP_MINT_IMPL),
+            MINT_IMPL
+        ) :
+        new PartyPool(
             name_,
             symbol_,
             _tokens,
@@ -96,7 +111,6 @@ contract PartyPlanner is IPartyPlanner {
             _flashFeePpm,
             PROTOCOL_FEE_PPM,
             PROTOCOL_FEE_ADDRESS,
-            _stable,
             PartyPoolSwapImpl(SWAP_MINT_IMPL),
             MINT_IMPL
         );

src/PartyPool.sol

@@ -14,6 +14,7 @@ import {LMSRStabilizedBalancedPair} from "./LMSRStabilizedBalancedPair.sol";
 import {PartyPoolBase} from "./PartyPoolBase.sol";
 import {PartyPoolMintImpl} from "./PartyPoolMintImpl.sol";
 import {PartyPoolSwapImpl} from "./PartyPoolSwapImpl.sol";
+import {Proxy} from "../lib/openzeppelin-contracts/contracts/proxy/Proxy.sol";
 
 /// @title PartyPool - LMSR-backed multi-asset pool with LP ERC20 token
 /// @notice A multi-asset liquidity pool backed by the LMSRStabilized pricing model.
@@ -58,9 +59,6 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
     // @inheritdoc IPartyPool
     function allProtocolFeesOwed() external view returns (uint256[] memory) { return protocolFeesOwed; }
 
-    /// @notice If true and there are exactly two assets, an optimized 2-asset stable-pair path is used for some computations.
-    bool immutable private IS_STABLE_PAIR; // if true, the optimized LMSRStabilizedBalancedPair optimization path is enabled
-
     /// @notice Address of the Mint implementation contract for delegatecall
     PartyPoolMintImpl private immutable MINT_IMPL;
     function mintImpl() external view returns (PartyPoolMintImpl) { return MINT_IMPL; }
@@ -89,7 +87,6 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
     /// @param kappa_ liquidity parameter κ (Q64.64) used to derive b = κ * S(q)
     /// @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.
     /// @param swapMintImpl_ address of the SwapMint implementation contract
     /// @param mintImpl_ address of the Mint implementation contract
     constructor(
@@ -102,7 +99,6 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
         uint256 flashFeePpm_,
         uint256 protocolFeePpm_,             // NEW: protocol share of fees (ppm)
         address protocolFeeAddress_,         // NEW: recipient for collected protocol tokens
-        bool stable_,
         PartyPoolSwapImpl swapMintImpl_,
         PartyPoolMintImpl mintImpl_
     ) ERC20External(name_, symbol_) {
@@ -118,7 +114,6 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
         require(protocolFeePpm_ < 1_000_000, "Pool: protocol fee >= ppm");
         PROTOCOL_FEE_PPM = protocolFeePpm_;
         PROTOCOL_FEE_ADDRESS = protocolFeeAddress_;
-        IS_STABLE_PAIR = stable_ && tokens_.length == 2;
         SWAP_IMPL = swapMintImpl_;
         MINT_IMPL = mintImpl_;
 
@@ -250,6 +245,7 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
         }
     }
 
+
     /// @notice Swap input token i -> token j. Payer must approve token i.
     /// @dev This function transfers the exact gross input (including fee) from payer and sends the computed output to receiver.
     ///      Non-standard tokens (fee-on-transfer, rebasers) are rejected via balance checks.
@@ -405,9 +401,7 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
 
         // Compute internal amounts using LMSR (exact-input with price limit)
         // if _stablePair is true, use the optimized path
-        (amountInInternalUsed, amountOutInternal) =
+        (amountInInternalUsed, amountOutInternal) = _swapAmountsForExactInput(inputTokenIndex, outputTokenIndex, deltaInternalI, limitPrice);
-            IS_STABLE_PAIR ? LMSRStabilizedBalancedPair.swapAmountsForExactInput(lmsr, inputTokenIndex, outputTokenIndex, deltaInternalI, limitPrice)
-                        : lmsr.swapAmountsForExactInput(inputTokenIndex, outputTokenIndex, deltaInternalI, limitPrice);
 
         // Convert actual used input internal -> uint (ceil)
         amountInUintNoFee = _internalToUintCeil(amountInInternalUsed, bases[inputTokenIndex]);
@@ -667,8 +661,6 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
     }
 
 
-    /* Conversion helpers moved to PartyPoolBase (abstract) to centralize internal helpers and storage. */
-
     /// @notice Marginal price of `base` in terms of `quote` (p_quote / p_base) as Q64.64
     /// @dev Returns the LMSR marginal price directly (raw 64.64) for use by off-chain quoting logic.
     function price(uint256 baseTokenIndex, uint256 quoteTokenIndex) external view returns (int128) {
@@ -705,4 +697,10 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
         return pricePerQ.mul(factor);
     }
 
+
+    function _swapAmountsForExactInput(uint256 i, uint256 j, int128 a, int128 limitPrice) internal virtual view
+    returns (int128 amountIn, int128 amountOut) {
+        return lmsr.swapAmountsForExactInput(i, j, a, limitPrice);
+    }
+
 }

src/PartyPoolBalancedPair.sol

@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.30;
+
+import {IERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
+import {LMSRStabilizedBalancedPair} from "./LMSRStabilizedBalancedPair.sol";
+import {PartyPool} from "./PartyPool.sol";
+import {PartyPoolMintImpl} from "./PartyPoolMintImpl.sol";
+import {PartyPoolSwapImpl} from "./PartyPoolSwapImpl.sol";
+
+contract PartyPoolBalancedPair is PartyPool {
+    constructor(
+        string memory name_,
+        string memory symbol_,
+        IERC20[] memory tokens_,
+        uint256[] memory bases_,
+        int128 kappa_,
+        uint256 swapFeePpm_,
+        uint256 flashFeePpm_,
+        uint256 protocolFeePpm_,             // NEW: protocol share of fees (ppm)
+        address protocolFeeAddress_,         // NEW: recipient for collected protocol tokens
+        PartyPoolSwapImpl swapMintImpl_,
+        PartyPoolMintImpl mintImpl_
+    ) PartyPool(name_, symbol_, tokens_, bases_, kappa_, swapFeePpm_, flashFeePpm_, protocolFeePpm_, protocolFeeAddress_, swapMintImpl_, mintImpl_)
+    {}
+
+    function _swapAmountsForExactInput(uint256 i, uint256 j, int128 a, int128 limitPrice) internal virtual override view
+    returns (int128 amountIn, int128 amountOut) {
+        return LMSRStabilizedBalancedPair.swapAmountsForExactInput(lmsr, i, j, a, limitPrice);
+    }
+}