proxied swapToLimit

script/DeployMock.sol

@@ -58,7 +58,7 @@ contract DeployMock is Script {
         }
 
         // call full newPool signature on factory which will take the deposits and mint initial LP
-        (PartyPool pool, ) = planner.newPool(
+        (IPartyPool pool, ) = planner.newPool(
             name,
             symbol,
             tokens,

src/Deploy.sol

@@ -1,17 +1,18 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.30;
 
-import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {IERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
-import {PartyPool} from "./PartyPool.sol";
-import {PartyPoolSwapMintImpl} from "./PartyPoolSwapMintImpl.sol";
-import {PartyPoolMintImpl} from "./PartyPoolMintImpl.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";
 
 library Deploy {
 
     function newPartyPlanner() internal returns (PartyPlanner) {
         return new PartyPlanner(
-            new PartyPoolSwapMintImpl(),
+            new PartyPoolSwapImpl(),
             new PartyPoolMintImpl(),
             0,               // protocolFeePpm = 0 for deploy helper
             address(0)       // protocolFeeAddress = address(0) for deploy helper
@@ -32,7 +33,8 @@ library Deploy {
         uint256 protocolFeePpm = 0;
         address protocolAddr = address(0);
 
-        return new PartyPool(
+        return _stable && tokens_.length == 2 ?
+        new PartyPoolBalancedPair(
             name_,
             symbol_,
             tokens_,
@@ -42,8 +44,20 @@ library Deploy {
             _flashFeePpm,
             protocolFeePpm,
             protocolAddr,
-            _stable,
+            new PartyPoolSwapImpl(),
-            new PartyPoolSwapMintImpl(),
+            new PartyPoolMintImpl()
+        ) :
+        new PartyPool(
+            name_,
+            symbol_,
+            tokens_,
+            bases_,
+            _kappa,
+            _swapFeePpm,
+            _flashFeePpm,
+            protocolFeePpm,
+            protocolAddr,
+            new PartyPoolSwapImpl(),
             new PartyPoolMintImpl()
         );
     }

src/IPartyPlanner.sol

@@ -1,14 +1,16 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.30;
 
-import "./PartyPool.sol";
+import "./IPartyPool.sol";
+import "./PartyPoolMintImpl.sol";
+import "./PartyPoolSwapImpl.sol";
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 
 /// @title IPartyPlanner
 /// @notice Interface for factory contract for creating and tracking PartyPool instances
 interface IPartyPlanner {
     // Event emitted when a new pool is created
-    event PartyStarted(PartyPool indexed pool, string name, string symbol, IERC20[] tokens);
+    event PartyStarted(IPartyPool indexed pool, string name, string symbol, IERC20[] tokens);
 
     /// @notice Creates a new PartyPool instance and initializes it with initial deposits (legacy signature).
     /// @dev Deprecated in favour of the kappa-based overload below; kept for backwards compatibility.
@@ -44,7 +46,7 @@ interface IPartyPlanner {
         uint256[] memory initialDeposits,
         uint256 initialLpAmount,
         uint256 deadline
-    ) external returns (PartyPool pool, uint256 lpAmount);
+    ) external returns (IPartyPool pool, uint256 lpAmount);
 
     /// @notice Creates a new PartyPool instance and initializes it with initial deposits (kappa-based).
     /// @param name_ LP token name
@@ -77,7 +79,7 @@ interface IPartyPlanner {
         uint256[] memory initialDeposits,
         uint256 initialLpAmount,
         uint256 deadline
-    ) external returns (PartyPool pool, uint256 lpAmount);
+    ) external returns (IPartyPool pool, uint256 lpAmount);
 
     /// @notice Checks if a pool is supported
     /// @param pool The pool address to check
@@ -92,7 +94,7 @@ interface IPartyPlanner {
     /// @param offset Starting index for pagination
     /// @param limit Maximum number of items to return
     /// @return pools Array of pool addresses for the requested page
-    function getAllPools(uint256 offset, uint256 limit) external view returns (PartyPool[] memory pools);
+    function getAllPools(uint256 offset, uint256 limit) external view returns (IPartyPool[] memory pools);
 
     /// @notice Returns the total number of unique tokens
     /// @return The total count of unique tokens
@@ -114,12 +116,12 @@ interface IPartyPlanner {
     /// @param offset Starting index for pagination
     /// @param limit Maximum number of items to return
     /// @return pools Array of pool addresses containing the specified token
-    function getPoolsByToken(IERC20 token, uint256 offset, uint256 limit) external view returns (PartyPool[] memory pools);
+    function getPoolsByToken(IERC20 token, uint256 offset, uint256 limit) external view returns (IPartyPool[] memory pools);
 
-    /// @notice Address of the SwapMint implementation contract used by all pools created by this factory
+    /// @notice Address of the mint implementation contract used by all pools created by this factory
     function mintImpl() external view returns (PartyPoolMintImpl);
 
-    /// @notice Address of the SwapMint implementation contract used by all pools created by this factory
+    /// @notice Address of the swap implementation contract used by all pools created by this factory
-    function swapMintImpl() external view returns (PartyPoolSwapMintImpl);
+    function swapMintImpl() external view returns (PartyPoolSwapImpl);
 
 }

src/IPartyPool.sol

@@ -1,7 +1,8 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.30;
 
-import "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";
+import {IERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
+import {IERC20Metadata} from "../lib/openzeppelin-contracts/contracts/token/ERC20/extensions/IERC20Metadata.sol";
 
 /// @title PartyPool - LMSR-backed multi-asset pool with LP ERC20 token
 /// @notice A multi-asset liquidity pool backed by the LMSRStabilized pricing model.

src/LMSRStabilized.sol

@@ -1,8 +1,7 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.30;
 
-import "forge-std/console2.sol";
+import {ABDKMath64x64} from "../lib/abdk-libraries-solidity/ABDKMath64x64.sol";
-import "@abdk/ABDKMath64x64.sol";
 
 /// @notice Stabilized LMSR library with incremental exp(z) caching for gas efficiency.
 /// - Stores b (64.64), M (shift), Z = sum exp(z_i), z[i] = (q_i / b) - M
@@ -41,21 +40,11 @@ library LMSRStabilized {
         }
 
         int128 total = _computeSizeMetric(s.qInternal);
-        console2.log("total (internal 64.64)");
-        console2.logInt(total);
         require(total > int128(0), "LMSR: total zero");
 
-        console2.log("LMSR.init: start");
-        console2.log("nAssets", s.nAssets);
-        console2.log("qInternal.length", s.qInternal.length);
-
         // Set kappa directly (caller provides kappa)
         s.kappa = kappa;
-        console2.log("kappa (64x64)");
-        console2.logInt(s.kappa);
         require(s.kappa > int128(0), "LMSR: kappa>0");
-
-        console2.log("LMSR.init: done");
     }
 
     /* --------------------
@@ -167,49 +156,31 @@ library LMSRStabilized {
         // push the marginal price p_i/p_j beyond the limit; if so, truncate `a`.
         // Marginal price ratio evolves as r(t) = r0 * exp(t/b) (since e_i multiplies by exp(t/b))
         if (limitPrice > int128(0)) {
-            console2.log("\n=== LimitPrice Logic Debug ===");
-            console2.log("Received limitPrice (64x64):");
-            console2.logInt(limitPrice);
-
-            console2.log("Current price ratio r0 (e_i/e_j, 64x64):");
-            console2.logInt(r0);
-
             // r0 must be positive; if r0 == 0 then no risk of exceeding limit by increasing r.
             require(r0 >= int128(0), "LMSR: r0<0");
             if (r0 == int128(0)) {
-                console2.log("r0 == 0 (input asset has zero weight), no limit truncation needed");
+//                console2.log("r0 == 0 (input asset has zero weight), no limit truncation needed");
             } else {
                 // If limitPrice <= current price, we revert (caller must choose a limit > current price to allow any fill)
                 if (limitPrice <= r0) {
-                    console2.log("Limit price is <= current price: reverting");
                     revert("LMSR: limitPrice <= current price");
                 }
 
                 // Compute a_limit directly from ln(limit / r0): a_limit = b * ln(limit / r0)
                 int128 ratioLimitOverR0 = limitPrice.div(r0);
-                console2.log("limitPrice/r0 (64x64):");
-                console2.logInt(ratioLimitOverR0);
                 require(ratioLimitOverR0 > int128(0), "LMSR: ratio<=0");
 
                 int128 aLimitOverB = _ln(ratioLimitOverR0); // > 0
-                console2.log("ln(limitPrice/r0) (64x64):");
-                console2.logInt(aLimitOverB);
 
                 // aLimit = b * aLimitOverB
                 int128 aLimit64 = b.mul(aLimitOverB);
-                console2.log("aLimit in 64x64 format:");
-                console2.logInt(aLimit64);
 
                 // If computed aLimit is less than the requested a, use the truncated value.
                 if (aLimit64 < a) {
-                    console2.log("TRUNCATING: a reduced from 64.64 value");
-                    console2.logInt(a);
-                    console2.log("to 64.64 value");
-                    console2.logInt(aLimit64);
                     amountIn = aLimit64; // Store the truncated input amount
                     a = aLimit64;        // Use truncated amount for calculations
                 } else {
-                    console2.log("Not truncating: aLimit64 >= a");
+//                    console2.log("Not truncating: aLimit64 >= a");
                 }
             }
         }
@@ -219,56 +190,24 @@ library LMSRStabilized {
         // Protect exp from enormous inputs (consistent with recenter thresholds)
         require(aOverB <= EXP_LIMIT, "LMSR: a/b too large (would overflow exp)");
 
-        console2.log("\n=== AmountOut Calculation Debug ===");
-        console2.log("Input amount (64.64):");
-        console2.logInt(a);
-        console2.log("a/b (64x64):");
-        console2.logInt(aOverB);
-
         // Use the closed-form fee-free formula:
         // y = b * ln(1 + r0 * (1 - exp(-a/b)))
-        console2.log("r0_for_calc (e_i/e_j):");
-        console2.logInt(r0);
-
         int128 expNeg = _exp(aOverB.neg()); // exp(-a/b)
-        console2.log("exp(-a/b):");
-        console2.logInt(expNeg);
-
         int128 oneMinusExpNeg = ONE.sub(expNeg);
-        console2.log("1 - exp(-a/b):");
-        console2.logInt(oneMinusExpNeg);
-
         int128 inner = ONE.add(r0.mul(oneMinusExpNeg));
-        console2.log("inner = 1 + r0 * (1 - exp(-a/b)):");
-        console2.logInt(inner);
 
         // If inner <= 0 then cap output to the current balance q_j (cannot withdraw more than q_j)
         if (inner <= int128(0)) {
-            console2.log("WARNING: inner <= 0, capping output to balance q_j");
             int128 qj64 = qInternal[j];
-            console2.log("Capped output (64.64):");
-            console2.logInt(qj64);
             return (amountIn, qj64);
         }
 
         int128 lnInner = _ln(inner);
-        console2.log("ln(inner):");
-        console2.logInt(lnInner);
-
         int128 b_lnInner = b.mul(lnInner);
-        console2.log("b*ln(inner):");
-        console2.logInt(b_lnInner);
-
         amountOut = b_lnInner;
-        console2.log("amountOut = b*ln(inner):");
-        console2.logInt(amountOut);
-
-        console2.log("amountOut (final 64.64 amount):");
-        console2.logInt(amountOut);
 
         // Safety check
         if (amountOut <= 0) {
-            console2.log("WARNING: x64 <= 0, returning 0");
             return (0, 0);
         }
     }
@@ -334,60 +273,34 @@ library LMSRStabilized {
         // Compute r0 = exp((q_i - q_j) / b) directly using invB
         int128 r0 = _exp(qInternal[i].sub(qInternal[j]).mul(invB));
 
-        console2.log("\n=== Max Input/Output Calculation ===");
-        console2.log("Limit price (64x64):");
-        console2.logInt(limitPrice);
-        console2.log("Current price ratio r0 (e_i/e_j, 64x64):");
-        console2.logInt(r0);
-
         // Mirror swapAmountsForExactInput behavior: treat invalid r0 as an error condition.
         // Revert if r0 is non-positive (no finite trade under a price limit).
         require(r0 > int128(0), "LMSR: r0<=0");
 
         // If current price already exceeds or equals limit, revert the same way swapAmountsForExactInput does.
         if (r0 >= limitPrice) {
-            console2.log("Limit price is <= current price: reverting");
             revert("LMSR: limitPrice <= current price");
         }
 
         // Calculate the price change factor: limitPrice/r0
         int128 priceChangeFactor = limitPrice.div(r0);
-        console2.log("Price change factor (limitPrice/r0):");
-        console2.logInt(priceChangeFactor);
 
         // ln(priceChangeFactor) gives us the maximum allowed delta in the exponent
         int128 maxDeltaExponent = _ln(priceChangeFactor);
-        console2.log("Max delta exponent ln(priceChangeFactor):");
-        console2.logInt(maxDeltaExponent);
 
         // Maximum input capable of reaching the price limit:
         // x_max = b * ln(limitPrice / r0)
         int128 amountInMax = b.mul(maxDeltaExponent);
-        console2.log("Max input to reach limit (64.64):");
-        console2.logInt(amountInMax);
 
         // The maximum output y corresponding to that input:
         // y = b * ln(1 + (e_i/e_j) * (1 - exp(-x_max/b)))
         int128 expTerm = ONE.sub(_exp(maxDeltaExponent.neg()));
-        console2.log("1 - exp(-maxDeltaExponent):");
-        console2.logInt(expTerm);
-
         int128 innerTerm = r0.mul(expTerm);
-        console2.log("e_i/e_j * expTerm:");
-        console2.logInt(innerTerm);
-
         int128 lnTerm = _ln(ONE.add(innerTerm));
-        console2.log("ln(1 + innerTerm):");
-        console2.logInt(lnTerm);
-
         int128 maxOutput = b.mul(lnTerm);
-        console2.log("Max output (b * lnTerm):");
-        console2.logInt(maxOutput);
 
         // Current balance of asset j (in 64.64)
         int128 qj64 = qInternal[j];
-        console2.log("Current j balance (64.64):");
-        console2.logInt(qj64);
 
         // Initialize outputs to the computed maxima
         amountIn = amountInMax;
@@ -395,7 +308,6 @@ library LMSRStabilized {
 
         // If the calculated maximum output exceeds the balance, cap output and solve for input.
         if (maxOutput > qj64) {
-            console2.log("Max output exceeds balance, capping to balance");
             amountOut = qj64;
 
             // Solve inverse relation for input given capped output:
@@ -405,19 +317,12 @@ library LMSRStabilized {
             //   a = -b * ln( (r0 + 1 - E) / r0 ) = b * ln( r0 / (r0 + 1 - E) )
             int128 E = _exp(amountOut.mul(invB)); // exp(y/b)
             int128 rhs = r0.add(ONE).sub(E); // r0 + 1 - E
-            console2.log("E = exp(y/b):");
-            console2.logInt(E);
-            console2.log("rhs = r0 + 1 - E:");
-            console2.logInt(rhs);
 
             // If rhs <= 0 due to numerical issues, fall back to amountInMax
             if (rhs <= int128(0)) {
-                console2.log("Numerical issue solving inverse; using amountInMax as fallback");
                 amountIn = amountInMax;
             } else {
                 amountIn = b.mul(_ln(r0.div(rhs)));
-                console2.log("Computed input required for capped output (64.64):");
-                console2.logInt(amountIn);
             }
         }
 
@@ -796,19 +701,9 @@ library LMSRStabilized {
         require(amountIn > int128(0), "LMSR: amountIn <= 0");
         require(amountOut > int128(0), "LMSR: amountOut <= 0");
 
-        console2.log("\n=== Applying Swap ===");
-        console2.log("Input asset:", i);
-        console2.log("Output asset:", j);
-        console2.log("Amount in (64.64):");
-        console2.logInt(amountIn);
-        console2.log("Amount out (64.64):");
-        console2.logInt(amountOut);
-
         // Update internal balances
         s.qInternal[i] = s.qInternal[i].add(amountIn);
         s.qInternal[j] = s.qInternal[j].sub(amountOut);
-
-        console2.log("=== Swap Applied (qInternal updated) ===\n");
     }
 
 
@@ -819,27 +714,16 @@ library LMSRStabilized {
     function updateForProportionalChange(State storage s, int128[] memory newQInternal) internal {
         require(newQInternal.length == s.nAssets, "LMSR: length mismatch");
 
-        console2.log("LMSR.updateForProportionalChange: start");
-
         // Compute new total for validation
         int128 newTotal = _computeSizeMetric(newQInternal);
-        console2.log("new total");
-        console2.logInt(newTotal);
 
         require(newTotal > int128(0), "LMSR: new total zero");
 
-        // With kappa formulation, b automatically scales with pool size
-        int128 newB = s.kappa.mul(newTotal);
-        console2.log("new effective b");
-        console2.logInt(newB);
-
         // Update the cached qInternal with new values
         for (uint i = 0; i < s.nAssets; ) {
             s.qInternal[i] = newQInternal[i];
             unchecked { i++; }
         }
-
-        console2.log("LMSR.updateForProportionalChange: end");
     }
 
     /// @notice Price-share of asset i: exp(z_i) / Z (64.64)
@@ -1064,8 +948,6 @@ library LMSRStabilized {
     /// @notice De-initialize the LMSR state when the entire pool is drained.
     /// This resets the state so the pool can be re-initialized by init(...) on next mint.
     function deinit(State storage s) internal {
-        console2.log("LMSR.deinit: resetting state");
-
         // Reset core state
         s.nAssets = 0;
         s.kappa = int128(0);

src/LMSRStabilizedBalancedPair.sol

@@ -1,9 +1,8 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.30;
 
-import "forge-std/console2.sol";
+import {ABDKMath64x64} from "../lib/abdk-libraries-solidity/ABDKMath64x64.sol";
-import "@abdk/ABDKMath64x64.sol";
+import {LMSRStabilized} from "./LMSRStabilized.sol";
-import "./LMSRStabilized.sol";
 
 /// @notice Specialized functions for the 2-asset stablecoin case
 library LMSRStabilizedBalancedPair {
@@ -40,7 +39,6 @@ library LMSRStabilizedBalancedPair {
 
         // 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);
         }
 
@@ -48,7 +46,6 @@ library LMSRStabilizedBalancedPair {
         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);
@@ -58,8 +55,6 @@ library LMSRStabilizedBalancedPair {
 
         // 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);
@@ -68,19 +63,13 @@ library LMSRStabilizedBalancedPair {
                 .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);
@@ -90,7 +79,6 @@ library LMSRStabilizedBalancedPair {
             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);
             }
 
@@ -101,63 +89,34 @@ library LMSRStabilizedBalancedPair {
                 .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");
+//                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);
         }
 
@@ -165,18 +124,13 @@ library LMSRStabilizedBalancedPair {
         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);
         }
 
@@ -200,39 +154,26 @@ library LMSRStabilizedBalancedPair {
         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);
         }
 
@@ -240,15 +181,8 @@ library LMSRStabilizedBalancedPair {
         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);
         }
 

src/PartyPlanner.sol

@@ -1,13 +1,15 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.30;
 
-import "./IPartyPlanner.sol";
+import {IERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
-import "./PartyPool.sol";
+import {SafeERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/utils/SafeERC20.sol";
-import "./LMSRStabilized.sol";
+import {IPartyPlanner} from "./IPartyPlanner.sol";
-import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {LMSRStabilized} from "./LMSRStabilized.sol";
-import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import {IPartyPool} from "./IPartyPool.sol";
-import {PartyPoolSwapMintImpl} from "./PartyPoolSwapMintImpl.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
@@ -20,8 +22,8 @@ contract PartyPlanner is IPartyPlanner {
     function mintImpl() external view returns (PartyPoolMintImpl) { return MINT_IMPL; }
 
     /// @notice Address of the SwapMint implementation contract used by all pools created by this factory
-    PartyPoolSwapMintImpl private immutable SWAP_MINT_IMPL;
+    PartyPoolSwapImpl private immutable SWAP_MINT_IMPL;
-    function swapMintImpl() external view returns (PartyPoolSwapMintImpl) { return SWAP_MINT_IMPL; }
+    function swapMintImpl() external view returns (PartyPoolSwapImpl) { return SWAP_MINT_IMPL; }
 
     /// @notice Protocol fee share (ppm) applied to fees collected by pools created by this planner
     uint256 private immutable PROTOCOL_FEE_PPM;
@@ -32,18 +34,18 @@ contract PartyPlanner is IPartyPlanner {
     function protocolFeeAddress() external view returns (address) { return PROTOCOL_FEE_ADDRESS; }
 
     // On-chain pool indexing
-    PartyPool[] private _allPools;
+    IPartyPool[] private _allPools;
     IERC20[] private _allTokens;
-    mapping(PartyPool => bool) private _poolSupported;
+    mapping(IPartyPool => bool) private _poolSupported;
     mapping(IERC20 => bool) private _tokenSupported;
-    mapping(IERC20 => PartyPool[]) private _poolsByToken;
+    mapping(IERC20 => IPartyPool[]) private _poolsByToken;
 
     /// @param _swapMintImpl address of the SwapMint implementation contract to be used by all pools
     /// @param _mintImpl address of the Mint implementation contract to be used by all pools
     /// @param _protocolFeePpm protocol fee share (ppm) to be used for pools created by this planner
     /// @param _protocolFeeAddress recipient address for protocol fees for pools created by this planner (may be address(0))
     constructor(
-        PartyPoolSwapMintImpl _swapMintImpl,
+        PartyPoolSwapImpl _swapMintImpl,
         PartyPoolMintImpl _mintImpl,
         uint256 _protocolFeePpm,
         address _protocolFeeAddress
@@ -75,7 +77,7 @@ contract PartyPlanner is IPartyPlanner {
         uint256[] memory initialDeposits,
         uint256 initialLpAmount,
         uint256 deadline
-    ) public returns (PartyPool pool, uint256 lpAmount) {
+    ) public returns (IPartyPool pool, uint256 lpAmount) {
         // Validate inputs
         require(deadline == 0 || block.timestamp <= deadline, "Planner: deadline exceeded");
         require(_tokens.length == initialDeposits.length, "Planner: tokens and deposits length mismatch");
@@ -86,7 +88,8 @@ 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,
@@ -96,8 +99,20 @@ contract PartyPlanner is IPartyPlanner {
             _flashFeePpm,
             PROTOCOL_FEE_PPM,
             PROTOCOL_FEE_ADDRESS,
-            _stable,
+            PartyPoolSwapImpl(SWAP_MINT_IMPL),
-            PartyPoolSwapMintImpl(SWAP_MINT_IMPL),
+            MINT_IMPL
+        ) :
+        new PartyPool(
+            name_,
+            symbol_,
+            _tokens,
+            _bases,
+            _kappa,
+            _swapFeePpm,
+            _flashFeePpm,
+            PROTOCOL_FEE_PPM,
+            PROTOCOL_FEE_ADDRESS,
+            PartyPoolSwapImpl(SWAP_MINT_IMPL),
             MINT_IMPL
         );
 
@@ -151,7 +166,7 @@ contract PartyPlanner is IPartyPlanner {
         uint256[] memory initialDeposits,
         uint256 initialLpAmount,
         uint256 deadline
-    ) external returns (PartyPool pool, uint256 lpAmount) {
+    ) external returns (IPartyPool pool, uint256 lpAmount) {
         // Validate fixed-point fractions: must be less than 1.0 in 64.64 fixed-point
         require(_tradeFrac < ONE, "Planner: tradeFrac must be < 1 (64.64)");
         require(_targetSlippage < ONE, "Planner: targetSlippage must be < 1 (64.64)");
@@ -179,7 +194,7 @@ contract PartyPlanner is IPartyPlanner {
     
     /// @inheritdoc IPartyPlanner
     function getPoolSupported(address pool) external view returns (bool) {
-        return _poolSupported[PartyPool(pool)];
+        return _poolSupported[IPartyPool(pool)];
     }
 
     /// @inheritdoc IPartyPlanner
@@ -188,19 +203,19 @@ contract PartyPlanner is IPartyPlanner {
     }
 
     /// @inheritdoc IPartyPlanner
-    function getAllPools(uint256 offset, uint256 limit) external view returns (PartyPool[] memory pools) {
+    function getAllPools(uint256 offset, uint256 limit) external view returns (IPartyPool[] memory pools) {
         uint256 totalPools = _allPools.length;
 
         // If offset is beyond array bounds, return empty array
         if (offset >= totalPools) {
-            return new PartyPool[](0);
+            return new IPartyPool[](0);
         }
 
         // Calculate actual number of pools to return (respecting bounds)
         uint256 itemsToReturn = (offset + limit > totalPools) ? (totalPools - offset) : limit;
 
         // Create result array of appropriate size
-        pools = new PartyPool[](itemsToReturn);
+        pools = new IPartyPool[](itemsToReturn);
 
         // Fill the result array
         for (uint256 i = 0; i < itemsToReturn; i++) {
@@ -244,20 +259,20 @@ contract PartyPlanner is IPartyPlanner {
     }
 
     /// @inheritdoc IPartyPlanner
-    function getPoolsByToken(IERC20 token, uint256 offset, uint256 limit) external view returns (PartyPool[] memory pools) {
+    function getPoolsByToken(IERC20 token, uint256 offset, uint256 limit) external view returns (IPartyPool[] memory pools) {
-        PartyPool[] storage tokenPools = _poolsByToken[token];
+        IPartyPool[] storage tokenPools = _poolsByToken[token];
         uint256 totalPools = tokenPools.length;
 
         // If offset is beyond array bounds, return empty array
         if (offset >= totalPools) {
-            return new PartyPool[](0);
+            return new IPartyPool[](0);
         }
 
         // Calculate actual number of pools to return (respecting bounds)
         uint256 itemsToReturn = (offset + limit > totalPools) ? (totalPools - offset) : limit;
 
         // Create result array of appropriate size
-        pools = new PartyPool[](itemsToReturn);
+        pools = new IPartyPool[](itemsToReturn);
 
         // Fill the result array
         for (uint256 i = 0; i < itemsToReturn; i++) {

src/PartyPool.sol

@@ -6,14 +6,15 @@ import {IERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/IERC20
 import {SafeERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/utils/SafeERC20.sol";
 import {Address} from "../lib/openzeppelin-contracts/contracts/utils/Address.sol";
 import {ReentrancyGuard} from "../lib/openzeppelin-contracts/contracts/utils/ReentrancyGuard.sol";
+import {ERC20External} from "./ERC20External.sol";
 import {IPartyFlashCallback} from "./IPartyFlashCallback.sol";
 import {IPartyPool} from "./IPartyPool.sol";
 import {LMSRStabilized} from "./LMSRStabilized.sol";
 import {LMSRStabilizedBalancedPair} from "./LMSRStabilizedBalancedPair.sol";
 import {PartyPoolBase} from "./PartyPoolBase.sol";
 import {PartyPoolMintImpl} from "./PartyPoolMintImpl.sol";
-import {PartyPoolSwapMintImpl} from "./PartyPoolSwapMintImpl.sol";
+import {PartyPoolSwapImpl} from "./PartyPoolSwapImpl.sol";
-import {ERC20External} from "./ERC20External.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,16 +59,13 @@ 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; }
 
     /// @notice Address of the SwapMint implementation contract for delegatecall
-    PartyPoolSwapMintImpl private immutable SWAP_MINT_IMPL;
+    PartyPoolSwapImpl private immutable SWAP_IMPL;
-    function swapMintImpl() external view returns (PartyPoolSwapMintImpl) { return SWAP_MINT_IMPL; }
+    function swapMintImpl() external view returns (PartyPoolSwapImpl) { return SWAP_IMPL; }
 
 
     /// @inheritdoc IPartyPool
@@ -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,8 +99,7 @@ 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_,
-        PartyPoolSwapMintImpl swapMintImpl_,
         PartyPoolMintImpl mintImpl_
     ) ERC20External(name_, symbol_) {
         require(tokens_.length > 1, "Pool: need >1 asset");
@@ -118,8 +114,7 @@ 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_;
-        SWAP_MINT_IMPL = swapMintImpl_;
         MINT_IMPL = mintImpl_;
 
         uint256 n = tokens_.length;
@@ -139,6 +134,47 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
     }
 
 
+    //
+    // Current marginal prices
+    //
+
+    /// @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) {
+        uint256 n = tokens.length;
+        require(baseTokenIndex < n && quoteTokenIndex < n, "price: idx");
+        require(lmsr.nAssets > 0, "price: uninit");
+        return lmsr.price(baseTokenIndex, quoteTokenIndex);
+    }
+
+    /// @notice Price of one LP token denominated in `quote` asset as Q64.64
+    /// @dev Computes LMSR poolPrice (quote per unit qTotal) and scales it by totalSupply() / qTotal
+    ///      to return price per LP token unit in quote asset (raw 64.64).
+    function poolPrice(uint256 quoteTokenIndex) external view returns (int128) {
+        uint256 n = tokens.length;
+        require(quoteTokenIndex < n, "poolPrice: idx");
+        require(lmsr.nAssets > 0, "poolPrice: uninit");
+
+        // price per unit of qTotal (Q64.64) from LMSR
+        int128 pricePerQ = lmsr.poolPrice(quoteTokenIndex);
+
+        // total internal q (qTotal) as Q64.64
+        int128 qTotal = _computeSizeMetric(lmsr.qInternal);
+        require(qTotal > int128(0), "poolPrice: qTotal zero");
+
+        // totalSupply as Q64.64
+        uint256 supply = totalSupply();
+        require(supply > 0, "poolPrice: zero supply");
+        int128 supplyQ64 = ABDKMath64x64.fromUInt(supply);
+
+        // factor = totalSupply / qTotal (Q64.64)
+        int128 factor = supplyQ64.div(qTotal);
+
+        // price per LP token = pricePerQ * factor (Q64.64)
+        return pricePerQ.mul(factor);
+    }
+
+
     /* ----------------------
        Initialization / Mint / Burn (LP token managed)
        ---------------------- */
@@ -220,54 +256,6 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
     }
 
     /// @inheritdoc IPartyPool
-    function swapToLimitAmounts(
-        uint256 inputTokenIndex,
-        uint256 outputTokenIndex,
-        int128 limitPrice
-    ) external view returns (uint256 amountIn, uint256 amountOut, uint256 fee) {
-        (uint256 grossIn, uint256 outUint,,,, uint256 feeUint) = _quoteSwapToLimit(inputTokenIndex, outputTokenIndex, limitPrice);
-        return (grossIn, outUint, feeUint);
-    }
-
-
-    /// @notice Transfer all protocol fees to the configured protocolFeeAddress and zero the ledger.
-    /// @dev Anyone can call; must have protocolFeeAddress != address(0) to be operational.
-    function collectProtocolFees() external nonReentrant {
-        address dest = PROTOCOL_FEE_ADDRESS;
-        require(dest != address(0), "collect: zero addr");
-
-        uint256 n = tokens.length;
-        for (uint256 i = 0; i < n; i++) {
-            uint256 owed = protocolFeesOwed[i];
-            if (owed == 0) continue;
-            uint256 bal = IERC20(tokens[i]).balanceOf(address(this));
-            require(bal >= owed, "collect: fee > bal");
-            protocolFeesOwed[i] = 0;
-            // transfer owed tokens to protocol destination
-            tokens[i].safeTransfer(dest, owed);
-            // update cached to effective onchain minus owed
-            cachedUintBalances[i] = bal - owed;
-        }
-    }
-
-    /// @dev Helper to record cached balances as effectiveBalance = onchain - owed. Reverts if owed > onchain.
-    function _recordCachedBalance(uint256 idx, uint256 onchainBal) internal {
-        uint256 owed = protocolFeesOwed[idx];
-        require(onchainBal >= owed, "balance < protocol owed");
-        cachedUintBalances[idx] = onchainBal - owed;
-    }
-
-    /// @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.
-    /// @param payer address of the account that pays for the swap
-    /// @param receiver address that will receive the output tokens
-    /// @param inputTokenIndex index of input asset
-    /// @param outputTokenIndex index of output asset
-    /// @param maxAmountIn maximum amount of token i (uint256) to transfer in (inclusive of fees)
-    /// @param limitPrice maximum acceptable marginal price (64.64 fixed point). Pass 0 to ignore.
-    /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore.
-    /// @return amountIn actual input used (uint256), amountOut actual output sent (uint256), fee fee taken from the input (uint256)
     function swap(
         address payer,
         address receiver,
@@ -320,62 +308,6 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
         return (totalTransferAmount, amountOutUint, feeUint);
     }
 
-    /// @notice Swap up to the price limit; computes max input to reach limit then performs swap.
-    /// @dev If balances prevent fully reaching the limit, the function caps and returns actuals.
-    ///      The payer must transfer the exact gross input computed by the view.
-    /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore.
-    function swapToLimit(
-        address payer,
-        address receiver,
-        uint256 inputTokenIndex,
-        uint256 outputTokenIndex,
-        int128 limitPrice,
-        uint256 deadline
-    ) external returns (uint256 amountInUsed, uint256 amountOut, uint256 fee) {
-        uint256 n = tokens.length;
-        require(inputTokenIndex < n && outputTokenIndex < n, "swapToLimit: idx");
-        require(limitPrice > int128(0), "swapToLimit: limit <= 0");
-        require(deadline == 0 || block.timestamp <= deadline, "swapToLimit: deadline exceeded");
-
-        // Read previous balances for affected assets
-        uint256 prevBalI = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
-        uint256 prevBalJ = IERC20(tokens[outputTokenIndex]).balanceOf(address(this));
-
-        // Compute amounts using the same path as views
-        (uint256 totalTransferAmount, uint256 amountOutUint, int128 amountInInternalMax, int128 amountOutInternal, uint256 amountInUsedUint, uint256 feeUint) =
-            _quoteSwapToLimit(inputTokenIndex, outputTokenIndex, limitPrice);
-
-        // Transfer the exact amount needed from payer and require exact receipt (revert on fee-on-transfer)
-        tokens[inputTokenIndex].safeTransferFrom(payer, address(this), totalTransferAmount);
-        uint256 balIAfter = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
-        require(balIAfter == prevBalI + totalTransferAmount, "swapToLimit: non-standard tokenIn");
-
-        // Transfer output to receiver and verify exact decrease
-        tokens[outputTokenIndex].safeTransfer(receiver, amountOutUint);
-        uint256 balJAfter = IERC20(tokens[outputTokenIndex]).balanceOf(address(this));
-        require(balJAfter == prevBalJ - amountOutUint, "swapToLimit: non-standard tokenOut");
-
-        // Accrue protocol share (floor) from the fee on input token
-        if (PROTOCOL_FEE_PPM > 0 && feeUint > 0 && PROTOCOL_FEE_ADDRESS != address(0)) {
-            uint256 protoShare = (feeUint * PROTOCOL_FEE_PPM) / 1_000_000; // floor
-            if (protoShare > 0) {
-                protocolFeesOwed[inputTokenIndex] += protoShare;
-            }
-        }
-
-        // Update caches to effective balances
-        _recordCachedBalance(inputTokenIndex, balIAfter);
-        _recordCachedBalance(outputTokenIndex, balJAfter);
-
-        // Apply swap to LMSR state with the internal amounts
-        lmsr.applySwap(inputTokenIndex, outputTokenIndex, amountInInternalMax, amountOutInternal);
-
-        // Maintain original event semantics (logs input without fee)
-        emit Swap(payer, receiver, tokens[inputTokenIndex], tokens[outputTokenIndex], amountInUsedUint, amountOutUint);
-
-        return (amountInUsedUint, amountOutUint, feeUint);
-    }
-
     /// @notice Internal quote for exact-input swap that mirrors swap() rounding and fee application
     /// @dev Returns amounts consistent with swap() semantics: grossIn includes fees (ceil), amountOut is floored.
     /// @return grossIn amount to transfer in (inclusive of fee), amountOutUint output amount (uint),
@@ -412,9 +344,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]);
@@ -436,70 +366,50 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
         require(amountOutUint > 0, "swap: output zero");
     }
 
-    /// @notice Internal quote for swap-to-limit that mirrors swapToLimit() rounding and fee application
+
-    /// @dev Computes the input required to reach limitPrice and the resulting output; all rounding matches swapToLimit.
+    /// @inheritdoc IPartyPool
-    /// @return grossIn amount to transfer in (inclusive of fee), amountOutUint output amount (uint),
+    function swapToLimitAmounts(
-    ///         amountInInternal and amountOutInternal (64.64), amountInUintNoFee input amount excluding fee (uint),
-    ///         feeUint fee taken from the gross input (uint)
-    function _quoteSwapToLimit(
         uint256 inputTokenIndex,
         uint256 outputTokenIndex,
         int128 limitPrice
-    )
+    ) external view returns (uint256 amountIn, uint256 amountOut, uint256 fee) {
-        internal
+        require(inputTokenIndex < tokens.length && outputTokenIndex < tokens.length, "swapToLimit: idx");
-        view
-        returns (
-            uint256 grossIn,
-            uint256 amountOutUint,
-            int128 amountInInternal,
-            int128 amountOutInternal,
-            uint256 amountInUintNoFee,
-            uint256 feeUint
-        )
-    {
-        uint256 n = tokens.length;
-        require(inputTokenIndex < n && outputTokenIndex < n, "swapToLimit: idx");
         require(limitPrice > int128(0), "swapToLimit: limit <= 0");
         require(lmsr.nAssets > 0, "swapToLimit: pool uninitialized");
 
-        // Compute internal maxima at the price limit
+        return SWAP_IMPL.swapToLimitAmounts(
-        (amountInInternal, amountOutInternal) = lmsr.swapAmountsForPriceLimit(inputTokenIndex, outputTokenIndex, limitPrice);
+            inputTokenIndex, outputTokenIndex, limitPrice,
-
+            bases, KAPPA, lmsr.qInternal, SWAP_FEE_PPM);
-        // Convert input to uint (ceil) and output to uint (floor)
-        amountInUintNoFee = _internalToUintCeil(amountInInternal, bases[inputTokenIndex]);
-        require(amountInUintNoFee > 0, "swapToLimit: input zero");
-
-        feeUint = 0;
-        grossIn = amountInUintNoFee;
-        if (SWAP_FEE_PPM > 0) {
-            feeUint = _ceilFee(amountInUintNoFee, SWAP_FEE_PPM);
-            grossIn += feeUint;
     }
 
-        amountOutUint = _internalToUintFloor(amountOutInternal, bases[outputTokenIndex]);
-        require(amountOutUint > 0, "swapToLimit: output zero");
-    }
 
-    /// @notice Compute fee and net amounts for a gross input (fee rounded up to favor the pool).
+    /// @inheritdoc IPartyPool
-    /// @return feeUint fee taken (uint) and netUint remaining for protocol use (uint)
+    function swapToLimit(
-    function _computeFee(uint256 gross) internal view returns (uint256 feeUint, uint256 netUint) {
+        address payer,
-        if (SWAP_FEE_PPM == 0) {
+        address receiver,
-            return (0, gross);
+        uint256 inputTokenIndex,
-        }
+        uint256 outputTokenIndex,
-        feeUint = _ceilFee(gross, SWAP_FEE_PPM);
+        int128 limitPrice,
-        netUint = gross - feeUint;
+        uint256 deadline
-    }
+    ) external nonReentrant returns (uint256 amountInUsed, uint256 amountOut, uint256 fee) {
-
+        bytes memory data = abi.encodeWithSignature(
-    /// @notice Convenience: return gross = net + fee(net) using ceiling for fee.
+            'swapToLimit(address,address,uint256,uint256,int128,uint256,uint256,uint256)',
-    function _addFee(uint256 netUint) internal view returns (uint256 gross) {
+            payer,
-        if (SWAP_FEE_PPM == 0) return netUint;
+            receiver,
-        uint256 fee = _ceilFee(netUint, SWAP_FEE_PPM);
+            inputTokenIndex,
-        return netUint + fee;
+            outputTokenIndex,
+            limitPrice,
+            deadline,
+            SWAP_FEE_PPM,
+            PROTOCOL_FEE_PPM
+        );
+        bytes memory result = Address.functionDelegateCall(address(SWAP_IMPL), data);
+        return abi.decode(result, (uint256,uint256,uint256));
     }
 
     function swapMintAmounts(uint256 inputTokenIndex, uint256 maxAmountIn) external view
     returns (uint256 amountInUsed, uint256 fee, uint256 lpMinted) {
-        return SWAP_MINT_IMPL.swapMintAmounts(
+        return MINT_IMPL.swapMintAmounts(
             inputTokenIndex,
             maxAmountIn,
             SWAP_FEE_PPM,
@@ -511,7 +421,7 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
 
     function burnSwapAmounts(uint256 lpAmount, uint256 inputTokenIndex) external view
     returns (uint256 amountOut) {
-        return SWAP_MINT_IMPL.burnSwapAmounts(
+        return MINT_IMPL.burnSwapAmounts(
             lpAmount,
             inputTokenIndex,
             SWAP_FEE_PPM,
@@ -537,32 +447,18 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
         uint256 deadline
     ) external returns (uint256 lpMinted) {
         bytes memory data = abi.encodeWithSignature(
-            "swapMint(address,address,uint256,uint256,uint256,uint256)",
+            "swapMint(address,address,uint256,uint256,uint256,uint256,uint256)",
             payer,
             receiver,
             inputTokenIndex,
             maxAmountIn,
             deadline,
-            SWAP_FEE_PPM
+            SWAP_FEE_PPM,
+            PROTOCOL_FEE_PPM
         );
 
-        bytes memory result = Address.functionDelegateCall(address(SWAP_MINT_IMPL), data);
+        bytes memory result = Address.functionDelegateCall(address(MINT_IMPL), data);
-        // New ABI: implementation returns (uint256 lpMinted, uint256 feeUintActual)
+        return abi.decode(result, (uint256));
-        (uint256 lpOut, uint256 feeUintActual) = abi.decode(result, (uint256, uint256));
-
-        // Accrue protocol share (floor) from the fee on the input token
-        if (PROTOCOL_FEE_PPM > 0 && feeUintActual > 0 && PROTOCOL_FEE_ADDRESS != address(0)) {
-            uint256 protoShare = (feeUintActual * PROTOCOL_FEE_PPM) / 1_000_000;
-            if (protoShare > 0) {
-                protocolFeesOwed[inputTokenIndex] += protoShare;
-            }
-        }
-
-        // Update cached balance for the input token to effective onchain - owed
-        uint256 bal = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
-        _recordCachedBalance(inputTokenIndex, bal);
-
-        return lpOut;
     }
 
     /// @notice Burn LP tokens then swap the redeemed proportional basket into a single asset `inputTokenIndex` and send to receiver.
@@ -581,32 +477,18 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
         uint256 deadline
     ) external returns (uint256 amountOutUint) {
         bytes memory data = abi.encodeWithSignature(
-            "burnSwap(address,address,uint256,uint256,uint256,uint256)",
+            "burnSwap(address,address,uint256,uint256,uint256,uint256,uint256)",
             payer,
             receiver,
             lpAmount,
             inputTokenIndex,
             deadline,
-            SWAP_FEE_PPM
+            SWAP_FEE_PPM,
+            PROTOCOL_FEE_PPM
         );
 
-        bytes memory result = Address.functionDelegateCall(address(SWAP_MINT_IMPL), data);
+        bytes memory result = Address.functionDelegateCall(address(MINT_IMPL), data);
-        // New ABI: implementation returns (uint256 amountOutUint, uint256 feeTokenUint)
+        return abi.decode(result, (uint256));
-        (uint256 outAmt, uint256 feeTokenUint) = abi.decode(result, (uint256, uint256));
-
-        // Accrue protocol share (floor) from the token-side fee computed by implementation
-        if (PROTOCOL_FEE_PPM > 0 && feeTokenUint > 0 && PROTOCOL_FEE_ADDRESS != address(0)) {
-            uint256 protoShare = (feeTokenUint * PROTOCOL_FEE_PPM) / 1_000_000;
-            if (protoShare > 0) {
-                protocolFeesOwed[inputTokenIndex] += protoShare;
-            }
-        }
-
-        // Update cached balance for the target token to effective onchain - owed
-        uint256 bal = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
-        _recordCachedBalance(inputTokenIndex, bal);
-
-        return outAmt;
     }
 
 
@@ -702,42 +584,41 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
     }
 
 
-    /* Conversion helpers moved to PartyPoolBase (abstract) to centralize internal helpers and storage. */
+    /// @notice Transfer all protocol fees to the configured protocolFeeAddress and zero the ledger.
+    /// @dev Anyone can call; must have protocolFeeAddress != address(0) to be operational.
+    function collectProtocolFees() external nonReentrant {
+        address dest = PROTOCOL_FEE_ADDRESS;
+        require(dest != address(0), "collect: zero addr");
 
-    /// @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) {
         uint256 n = tokens.length;
-        require(baseTokenIndex < n && quoteTokenIndex < n, "price: idx");
+        for (uint256 i = 0; i < n; i++) {
-        require(lmsr.nAssets > 0, "price: uninit");
+            uint256 owed = protocolFeesOwed[i];
-        return lmsr.price(baseTokenIndex, quoteTokenIndex);
+            if (owed == 0) continue;
+            uint256 bal = IERC20(tokens[i]).balanceOf(address(this));
+            require(bal >= owed, "collect: fee > bal");
+            protocolFeesOwed[i] = 0;
+            // transfer owed tokens to protocol destination
+            tokens[i].safeTransfer(dest, owed);
+            // update cached to effective onchain minus owed
+            cachedUintBalances[i] = bal - owed;
+        }
     }
 
-    /// @notice Price of one LP token denominated in `quote` asset as Q64.64
-    /// @dev Computes LMSR poolPrice (quote per unit qTotal) and scales it by totalSupply() / qTotal
-    ///      to return price per LP token unit in quote asset (raw 64.64).
-    function poolPrice(uint256 quoteTokenIndex) external view returns (int128) {
-        uint256 n = tokens.length;
-        require(quoteTokenIndex < n, "poolPrice: idx");
-        require(lmsr.nAssets > 0, "poolPrice: uninit");
 
-        // price per unit of qTotal (Q64.64) from LMSR
+    function _swapAmountsForExactInput(uint256 i, uint256 j, int128 a, int128 limitPrice) internal virtual view
-        int128 pricePerQ = lmsr.poolPrice(quoteTokenIndex);
+    returns (int128 amountIn, int128 amountOut) {
+        return lmsr.swapAmountsForExactInput(i, j, a, limitPrice);
+    }
 
-        // total internal q (qTotal) as Q64.64
+    /// @notice Compute fee and net amounts for a gross input (fee rounded up to favor the pool).
-        int128 qTotal = _computeSizeMetric(lmsr.qInternal);
+    /// @return feeUint fee taken (uint) and netUint remaining for protocol use (uint)
-        require(qTotal > int128(0), "poolPrice: qTotal zero");
+    function _computeFee(uint256 gross) internal view returns (uint256 feeUint, uint256 netUint) {
+        return _computeFee(gross, SWAP_FEE_PPM);
+    }
 
-        // totalSupply as Q64.64
+    /// @notice Convenience: return gross = net + fee(net) using ceiling for fee.
-        uint256 supply = totalSupply();
+    function _addFee(uint256 netUint) internal view returns (uint256 gross) {
-        require(supply > 0, "poolPrice: zero supply");
+        return _addFee(netUint, SWAP_FEE_PPM);
-        int128 supplyQ64 = ABDKMath64x64.fromUInt(supply);
-
-        // factor = totalSupply / qTotal (Q64.64)
-        int128 factor = supplyQ64.div(qTotal);
-
-        // price per LP token = pricePerQ * factor (Q64.64)
-        return pricePerQ.mul(factor);
     }
 
 }

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);
+    }
+}

src/PartyPoolBase.sol

@@ -116,4 +116,12 @@ abstract contract PartyPoolBase is ERC20Internal, ReentrancyGuard {
         }
         return total;
     }
+
+    /// @dev Helper to record cached balances as effectiveBalance = onchain - owed. Reverts if owed > onchain.
+    function _recordCachedBalance(uint256 idx, uint256 onchainBal) internal {
+        uint256 owed = protocolFeesOwed[idx];
+        require(onchainBal >= owed, "balance < protocol owed");
+        cachedUintBalances[idx] = onchainBal - owed;
+    }
+
 }

src/PartyPoolMintImpl.sol

@@ -1,12 +1,13 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.30;
 
-import "@abdk/ABDKMath64x64.sol";
+import {ABDKMath64x64} from "../lib/abdk-libraries-solidity/ABDKMath64x64.sol";
-import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {IERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
-import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import {SafeERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/utils/SafeERC20.sol";
-import "./PartyPoolBase.sol";
+import {ERC20Internal} from "./ERC20Internal.sol";
-import "./LMSRStabilized.sol";
+import {IPartyPool} from "./IPartyPool.sol";
-import {PartyPool} from "./PartyPool.sol";
+import {LMSRStabilized} from "./LMSRStabilized.sol";
+import {PartyPoolBase} from "./PartyPoolBase.sol";
 
 /// @title PartyPoolMintImpl - Implementation contract for mint and burn functions
 /// @notice This contract contains the mint and burn implementation that will be called via delegatecall
@@ -20,6 +21,10 @@ contract PartyPoolMintImpl is PartyPoolBase {
     event Mint(address indexed payer, address indexed receiver, uint256[] depositAmounts, uint256 lpMinted);
     event Burn(address indexed payer, address indexed receiver, uint256[] withdrawAmounts, uint256 lpBurned);
 
+    //
+    // Initialization Mint
+    //
+
     function initialMint(address receiver, uint256 lpTokens, int128 KAPPA) external
     returns (uint256 lpMinted) {
         uint256 n = tokens.length;
@@ -55,6 +60,11 @@ contract PartyPoolMintImpl is PartyPoolBase {
         emit Mint(address(0), receiver, depositAmounts, lpMinted);
     }
 
+
+    //
+    // Regular Mint and Burn
+    //
+
     function mint(address payer, address receiver, uint256 lpTokenAmount, uint256 deadline) external returns (uint256 lpMinted) {
         require(deadline == 0 || block.timestamp <= deadline, "mint: deadline exceeded");
         uint256 n = tokens.length;
@@ -192,7 +202,8 @@ contract PartyPoolMintImpl is PartyPoolBase {
         emit Burn(payer, receiver, withdrawAmounts, lpAmount);
     }
 
-    function mintAmounts(uint256 lpTokenAmount, uint256 numAssets, uint256 totalSupply, uint256[] memory cachedUintBalances) public pure
+    function mintAmounts(uint256 lpTokenAmount,
+        uint256 numAssets, uint256 totalSupply, uint256[] memory cachedUintBalances) public pure
     returns (uint256[] memory depositAmounts) {
         depositAmounts = new uint256[](numAssets);
 
@@ -232,4 +243,357 @@ contract PartyPoolMintImpl is PartyPoolBase {
 
         return withdrawAmounts;
     }
+
+
+    //
+    // Swap-Mint and Burn-Swap
+    //
+
+    /// @notice Single-token mint: deposit a single token, charge swap-LMSR cost, and mint LP.
+    /// @dev swapMint executes as an exact-in planned swap followed by proportional scaling of qInternal.
+    ///      The function emits SwapMint (gross, net, fee) and also emits Mint for LP issuance.
+    /// @param payer who transfers the input token
+    /// @param receiver who receives the minted LP tokens
+    /// @param inputTokenIndex index of the input token
+    /// @param maxAmountIn maximum uint token input (inclusive of fee)
+    /// @param deadline optional deadline
+    /// @param swapFeePpm fee in parts-per-million for this pool
+    /// @return lpMinted actual LP minted (uint)
+    function swapMint(
+        address payer,
+        address receiver,
+        uint256 inputTokenIndex,
+        uint256 maxAmountIn,
+        uint256 deadline,
+        uint256 swapFeePpm,
+        uint256 protocolFeePpm
+    ) external returns (uint256 lpMinted) {
+        uint256 n = tokens.length;
+        require(inputTokenIndex < n, "swapMint: idx");
+        require(maxAmountIn > 0, "swapMint: input zero");
+        require(deadline == 0 || block.timestamp <= deadline, "swapMint: deadline");
+        require(lmsr.nAssets > 0, "swapMint: uninit pool");
+
+        // compute fee on gross maxAmountIn to get an initial net estimate (we'll recompute based on actual used)
+        (, uint256 netUintGuess) = _computeFee(maxAmountIn, swapFeePpm);
+
+        // Convert the net guess to internal (floor)
+        int128 netInternalGuess = _uintToInternalFloor(netUintGuess, bases[inputTokenIndex]);
+        require(netInternalGuess > int128(0), "swapMint: input too small after fee");
+
+        // Use LMSR view to determine actual internal consumed and size-increase (ΔS) for mint
+        (int128 amountInInternalUsed, int128 sizeIncreaseInternal) = lmsr.swapAmountsForMint(inputTokenIndex, netInternalGuess);
+
+        // amountInInternalUsed may be <= netInternalGuess. Convert to uint (ceil) to determine actual transfer
+        uint256 amountInUint = _internalToUintCeil(amountInInternalUsed, bases[inputTokenIndex]);
+        require(amountInUint > 0, "swapMint: input zero after internal conversion");
+
+        // Compute fee on the actual used input and total transfer amount (ceiling)
+        uint256 feeUintActual = _ceilFee(amountInUint, swapFeePpm);
+        uint256 totalTransfer = amountInUint + feeUintActual;
+        require(totalTransfer > 0 && totalTransfer <= maxAmountIn, "swapMint: transfer exceeds max");
+
+        // Record pre-balance and transfer tokens from payer, require exact receipt (revert on fee-on-transfer)
+        uint256 prevBalI = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
+        tokens[inputTokenIndex].safeTransferFrom(payer, address(this), totalTransfer);
+        uint256 balIAfter = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
+        require(balIAfter == prevBalI + totalTransfer, "swapMint: non-standard tokenIn");
+
+        // Accrue protocol share (floor) from the fee on the input token
+        if (protocolFeePpm > 0 && feeUintActual > 0) {
+            uint256 protoShare = (feeUintActual * protocolFeePpm) / 1_000_000;
+            if (protoShare > 0) {
+                protocolFeesOwed[inputTokenIndex] += protoShare;
+            }
+        }
+        // Update cached balance for the input token to effective onchain - owed
+        _recordCachedBalance(inputTokenIndex, balIAfter);
+
+        // Compute old and new scaled size metrics to determine LP minted
+        int128 oldTotal = _computeSizeMetric(lmsr.qInternal);
+        require(oldTotal > int128(0), "swapMint: zero total");
+        uint256 oldScaled = ABDKMath64x64.mulu(oldTotal, LP_SCALE);
+
+        int128 newTotal = oldTotal.add(sizeIncreaseInternal);
+        uint256 newScaled = ABDKMath64x64.mulu(newTotal, LP_SCALE);
+
+        uint256 actualLpToMint;
+        // Use natural ERC20 function since base contract inherits from ERC20
+        uint256 currentSupply = _totalSupply;
+        if (currentSupply == 0) {
+            // If somehow supply zero (shouldn't happen as lmsr.nAssets>0), mint newScaled
+            actualLpToMint = newScaled;
+        } else {
+            require(oldScaled > 0, "swapMint: oldScaled zero");
+            uint256 delta = (newScaled > oldScaled) ? (newScaled - oldScaled) : 0;
+            if (delta > 0) {
+                // floor truncation rounds in favor of pool
+                actualLpToMint = (currentSupply * delta) / oldScaled;
+            } else {
+                actualLpToMint = 0;
+            }
+        }
+
+        require(actualLpToMint > 0, "swapMint: zero LP minted");
+
+        // Update LMSR internal state: scale qInternal proportionally by newTotal/oldTotal
+        int128[] memory newQInternal = new int128[](n);
+        for (uint256 idx = 0; idx < n; idx++) {
+            // newQInternal[idx] = qInternal[idx] * (newTotal / oldTotal)
+            newQInternal[idx] = lmsr.qInternal[idx].mul(newTotal).div(oldTotal);
+        }
+
+        // Update cached internal and kappa via updateForProportionalChange
+        lmsr.updateForProportionalChange(newQInternal);
+
+        // Use natural ERC20 function since base contract inherits from ERC20
+        _mint(receiver, actualLpToMint);
+
+        // Emit SwapMint event with gross transfer, net input and fee (planned exact-in)
+        emit IPartyPool.SwapMint(payer, receiver, inputTokenIndex, totalTransfer, amountInUint, feeUintActual);
+
+        // Emit standard Mint event which records deposit amounts and LP minted
+        emit IPartyPool.Mint(payer, receiver, new uint256[](n), actualLpToMint);
+        // Note: depositAmounts array omitted (empty) since swapMint uses single-token input
+
+        return actualLpToMint;
+    }
+
+    /// @notice Calculate the amounts for a swap mint operation
+    /// @dev This is a pure view function that computes swap mint amounts from provided state
+    /// @param inputTokenIndex index of the input token
+    /// @param maxAmountIn maximum amount of token to deposit (inclusive of fee)
+    /// @param swapFeePpm fee in parts-per-million
+    /// @param lmsrState current LMSR state
+    /// @param bases_ scaling bases for each token
+    /// @param totalSupply_ current total LP token supply
+    /// @return amountInUsed actual input amount used (excluding fee)
+    /// @return fee fee amount charged
+    /// @return lpMinted LP tokens that would be minted
+    function swapMintAmounts(
+        uint256 inputTokenIndex,
+        uint256 maxAmountIn,
+        uint256 swapFeePpm,
+        LMSRStabilized.State memory lmsrState,
+        uint256[] memory bases_,
+        uint256 totalSupply_
+    ) public pure returns (uint256 amountInUsed, uint256 fee, uint256 lpMinted) {
+        require(inputTokenIndex < bases_.length, "swapMintAmounts: idx");
+        require(maxAmountIn > 0, "swapMintAmounts: input zero");
+        require(lmsrState.nAssets > 0, "swapMintAmounts: uninit pool");
+
+        // Compute fee on gross maxAmountIn to get an initial net estimate
+        uint256 feeGuess = 0;
+        uint256 netUintGuess = maxAmountIn;
+        if (swapFeePpm > 0) {
+            feeGuess = (maxAmountIn * swapFeePpm + 999999) / 1000000; // ceil fee
+            netUintGuess = maxAmountIn - feeGuess;
+        }
+
+        // Convert the net guess to internal (floor)
+        int128 netInternalGuess = _uintToInternalFloorPure(netUintGuess, bases_[inputTokenIndex]);
+        require(netInternalGuess > int128(0), "swapMintAmounts: input too small after fee");
+
+        // Use LMSR view to determine actual internal consumed and size-increase (ΔS) for mint
+        (int128 amountInInternalUsed, int128 sizeIncreaseInternal) =
+            LMSRStabilized.swapAmountsForMint(lmsrState.nAssets, lmsrState.kappa, lmsrState.qInternal,
+                inputTokenIndex, netInternalGuess);
+
+        // amountInInternalUsed may be <= netInternalGuess. Convert to uint (ceil) to determine actual transfer
+        amountInUsed = _internalToUintCeilPure(amountInInternalUsed, bases_[inputTokenIndex]);
+        require(amountInUsed > 0, "swapMintAmounts: input zero after internal conversion");
+
+        // Compute fee on the actual used input (ceiling)
+        fee = 0;
+        if (swapFeePpm > 0) {
+            fee = (amountInUsed * swapFeePpm + 999999) / 1000000; // ceil fee
+        }
+        uint256 totalTransfer = amountInUsed + fee;
+        require(totalTransfer > 0 && totalTransfer <= maxAmountIn, "swapMintAmounts: transfer exceeds max");
+
+        // Compute old and new scaled size metrics to determine LP minted
+        int128 oldTotal = _computeSizeMetricPure(lmsrState.qInternal);
+        require(oldTotal > int128(0), "swapMintAmounts: zero total");
+        uint256 oldScaled = ABDKMath64x64.mulu(oldTotal, LP_SCALE);
+
+        int128 newTotal = oldTotal.add(sizeIncreaseInternal);
+        uint256 newScaled = ABDKMath64x64.mulu(newTotal, LP_SCALE);
+
+        if (totalSupply_ == 0) {
+            // If somehow supply zero (shouldn't happen as lmsr.nAssets>0), mint newScaled
+            lpMinted = newScaled;
+        } else {
+            require(oldScaled > 0, "swapMintAmounts: oldScaled zero");
+            uint256 delta = (newScaled > oldScaled) ? (newScaled - oldScaled) : 0;
+            if (delta > 0) {
+                // floor truncation rounds in favor of pool
+                lpMinted = (totalSupply_ * delta) / oldScaled;
+            } else {
+                lpMinted = 0;
+            }
+        }
+
+        require(lpMinted > 0, "swapMintAmounts: zero LP minted");
+    }
+
+    /// @notice Calculate the amounts for a burn swap operation
+    /// @dev This is a pure view function that computes burn swap amounts from provided state
+    /// @param lpAmount amount of LP tokens to burn
+    /// @param inputTokenIndex index of target asset to receive
+    /// @param swapFeePpm fee in parts-per-million
+    /// @param lmsrState current LMSR state
+    /// @param bases_ scaling bases for each token
+    /// @param totalSupply_ current total LP token supply
+    /// @return amountOut amount of target asset that would be received
+    function burnSwapAmounts(
+        uint256 lpAmount,
+        uint256 inputTokenIndex,
+        uint256 swapFeePpm,
+        LMSRStabilized.State memory lmsrState,
+        uint256[] memory bases_,
+        uint256 totalSupply_
+    ) public pure returns (uint256 amountOut) {
+        require(inputTokenIndex < bases_.length, "burnSwapAmounts: idx");
+        require(lpAmount > 0, "burnSwapAmounts: zero lp");
+        require(totalSupply_ > 0, "burnSwapAmounts: empty supply");
+
+        // alpha = lpAmount / supply as Q64.64
+        int128 alpha = ABDKMath64x64.divu(lpAmount, totalSupply_) // fraction of total supply to burn
+            .mul(ABDKMath64x64.divu(1000000-swapFeePpm, 1000000)); // adjusted for fee
+
+        // Use LMSR view to compute single-asset payout and burned size-metric
+        (int128 payoutInternal, ) = LMSRStabilized.swapAmountsForBurn(lmsrState.nAssets, lmsrState.kappa, lmsrState.qInternal,
+            inputTokenIndex, alpha);
+
+        // Convert payoutInternal -> uint (floor) to favor pool
+        amountOut = _internalToUintFloorPure(payoutInternal, bases_[inputTokenIndex]);
+        require(amountOut > 0, "burnSwapAmounts: output zero");
+    }
+
+    /// @notice Burn LP tokens then swap the redeemed proportional basket into a single asset `inputTokenIndex` and send to receiver.
+    /// @dev The function burns LP tokens (authorization via allowance if needed), sends the single-asset payout and updates LMSR state.
+    /// @param payer who burns LP tokens
+    /// @param receiver who receives the single asset
+    /// @param lpAmount amount of LP tokens to burn
+    /// @param inputTokenIndex index of target asset to receive
+    /// @param deadline optional deadline
+    /// @param swapFeePpm fee in parts-per-million for this pool (may be used for future fee logic)
+    /// @return amountOutUint uint amount of asset i sent to receiver
+    function burnSwap(
+        address payer,
+        address receiver,
+        uint256 lpAmount,
+        uint256 inputTokenIndex,
+        uint256 deadline,
+        uint256 swapFeePpm,
+        uint256 protocolFeePpm
+    ) external returns (uint256 amountOutUint) {
+        uint256 n = tokens.length;
+        require(inputTokenIndex < n, "burnSwap: idx");
+        require(lpAmount > 0, "burnSwap: zero lp");
+        require(deadline == 0 || block.timestamp <= deadline, "burnSwap: deadline");
+
+        uint256 supply = _totalSupply;
+        require(supply > 0, "burnSwap: empty supply");
+        require(_balances[payer] >= lpAmount, "burnSwap: insufficient LP");
+
+        // alpha = lpAmount / supply as Q64.64 (adjusted for fee)
+        int128 alpha = ABDKMath64x64.divu(lpAmount, supply) // fraction of total supply to burn
+            .mul(ABDKMath64x64.divu(1000000-swapFeePpm, 1000000)); // adjusted for fee
+
+        // Use LMSR view to compute single-asset payout and burned size-metric
+        (int128 payoutInternal, ) = lmsr.swapAmountsForBurn(inputTokenIndex, alpha);
+
+        // Convert payoutInternal -> uint (floor) to favor pool
+        amountOutUint = _internalToUintFloor(payoutInternal, bases[inputTokenIndex]);
+        require(amountOutUint > 0, "burnSwap: output zero");
+
+        // Compute gross payout (no swap fee) so we can determine token-side fee = gross - net
+        int128 alphaGross = ABDKMath64x64.divu(lpAmount, supply); // gross fraction (no swap fee)
+        (int128 payoutGrossInternal, ) = lmsr.swapAmountsForBurn(inputTokenIndex, alphaGross);
+        uint256 payoutGrossUint = _internalToUintFloor(payoutGrossInternal, bases[inputTokenIndex]);
+        uint256 feeTokenUint = (payoutGrossUint > amountOutUint) ? (payoutGrossUint - amountOutUint) : 0;
+
+        // Accrue protocol share (floor) from the token-side fee
+        if (protocolFeePpm > 0 && feeTokenUint > 0) {
+            uint256 protoShare = (feeTokenUint * protocolFeePpm) / 1_000_000;
+            if (protoShare > 0) {
+                protocolFeesOwed[inputTokenIndex] += protoShare;
+            }
+        }
+
+        // Transfer the payout to receiver
+        tokens[inputTokenIndex].safeTransfer(receiver, amountOutUint);
+
+        // Burn LP tokens from payer (authorization via allowance)
+        if (msg.sender != payer) {
+            uint256 allowed = _allowances[payer][msg.sender];
+            require(allowed >= lpAmount, "burnSwap: allowance insufficient");
+            _approve(payer, msg.sender, allowed - lpAmount);
+        }
+        _burn(payer, lpAmount);
+
+        // Update cached balances by reading on-chain balances for all tokens
+        int128[] memory newQInternal = new int128[](n);
+        for (uint256 idx = 0; idx < n; idx++) {
+            uint256 bal = IERC20(tokens[idx]).balanceOf(address(this));
+            cachedUintBalances[idx] = bal;
+            _recordCachedBalance(inputTokenIndex, bal);
+            newQInternal[idx] = _uintToInternalFloor(bal, bases[idx]);
+        }
+
+        // Emit BurnSwap with public-facing info only (do not expose ΔS or LP burned)
+        emit IPartyPool.BurnSwap(payer, receiver, inputTokenIndex, amountOutUint);
+
+        // If entire pool drained, deinit; else update proportionally
+        bool allZero = true;
+        for (uint256 idx = 0; idx < n; idx++) {
+            if (newQInternal[idx] != int128(0)) { allZero = false; break; }
+        }
+        if (allZero) {
+            lmsr.deinit();
+        } else {
+            lmsr.updateForProportionalChange(newQInternal);
+        }
+
+        emit IPartyPool.Burn(payer, receiver, new uint256[](n), lpAmount);
+        return amountOutUint;
+    }
+
+    /// @notice Pure version of _uintToInternalFloor for use in view functions
+    function _uintToInternalFloorPure(uint256 amount, uint256 base) internal pure returns (int128) {
+        // amount / base as Q64.64, floored
+        return ABDKMath64x64.divu(amount, base);
+    }
+
+    /// @notice Pure version of _internalToUintCeil for use in view functions
+    function _internalToUintCeilPure(int128 amount, uint256 base) internal pure returns (uint256) {
+        // Convert Q64.64 to uint with ceiling: ceil(amount * base)
+        // Use mulu which floors, then add remainder check for ceiling
+        uint256 floored = ABDKMath64x64.mulu(amount, base);
+        // Check if there's a fractional part by computing amount * base - floored
+        int128 baseQ64 = ABDKMath64x64.fromUInt(base);
+        int128 flooredQ64 = ABDKMath64x64.fromUInt(floored);
+        int128 product = amount.mul(baseQ64);
+        if (product > flooredQ64) {
+            return floored + 1; // Ceiling
+        }
+        return floored;
+    }
+
+    /// @notice Pure version of _internalToUintFloor for use in view functions
+    function _internalToUintFloorPure(int128 amount, uint256 base) internal pure returns (uint256) {
+        // Convert Q64.64 to uint with floor: floor(amount * base)
+        return ABDKMath64x64.mulu(amount, base);
+    }
+
+    /// @notice Pure version of _computeSizeMetric for use in view functions
+    function _computeSizeMetricPure(int128[] memory qInternal) internal pure returns (int128) {
+        int128 sum = int128(0);
+        for (uint256 i = 0; i < qInternal.length; i++) {
+            sum = sum.add(qInternal[i]);
+        }
+        return sum;
+    }
+
 }

src/PartyPoolSwapImpl.sol

@@ -0,0 +1,142 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.30;
+
+import {ABDKMath64x64} from "../lib/abdk-libraries-solidity/ABDKMath64x64.sol";
+import {IERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
+import {SafeERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/utils/SafeERC20.sol";
+import {LMSRStabilized} from "./LMSRStabilized.sol";
+import {PartyPoolBase} from "./PartyPoolBase.sol";
+import {IPartyPool} from "./IPartyPool.sol";
+
+/// @title PartyPoolSwapMintImpl - Implementation contract for swapMint and burnSwap functions
+/// @notice This contract contains the swapMint and burnSwap implementation that will be called via delegatecall
+/// @dev This contract inherits from PartyPoolBase to access storage and internal functions
+contract PartyPoolSwapImpl is PartyPoolBase {
+    using ABDKMath64x64 for int128;
+    using LMSRStabilized for LMSRStabilized.State;
+    using SafeERC20 for IERC20;
+
+    function swapToLimitAmounts(
+        uint256 inputTokenIndex,
+        uint256 outputTokenIndex,
+        int128 limitPrice,
+        uint256[] memory bases,
+        int128 kappa,
+        int128[] memory qInternal,
+        uint256 swapFeePpm
+    ) external pure returns (uint256 amountIn, uint256 amountOut, uint256 fee) {
+        // Compute internal maxima at the price limit
+        (int128 amountInInternal, int128 amountOutInternal) = LMSRStabilized.swapAmountsForPriceLimit(
+            bases.length, kappa, qInternal,
+            inputTokenIndex, outputTokenIndex, limitPrice);
+
+        // Convert input to uint (ceil) and output to uint (floor)
+        uint256 amountInUintNoFee = _internalToUintCeil(amountInInternal, bases[inputTokenIndex]);
+        require(amountInUintNoFee > 0, "swapToLimit: input zero");
+
+        fee = 0;
+        amountIn = amountInUintNoFee;
+        if (swapFeePpm > 0) {
+            fee = _ceilFee(amountInUintNoFee, swapFeePpm);
+            amountIn += fee;
+        }
+
+        amountOut = _internalToUintFloor(amountOutInternal, bases[outputTokenIndex]);
+        require(amountOut > 0, "swapToLimit: output zero");
+    }
+
+
+    function swapToLimit(
+        address payer,
+        address receiver,
+        uint256 inputTokenIndex,
+        uint256 outputTokenIndex,
+        int128 limitPrice,
+        uint256 deadline,
+        uint256 swapFeePpm,
+        uint256 protocolFeePpm
+    ) external returns (uint256 amountInUsed, uint256 amountOut, uint256 fee) {
+        uint256 n = tokens.length;
+        require(inputTokenIndex < n && outputTokenIndex < n, "swapToLimit: idx");
+        require(limitPrice > int128(0), "swapToLimit: limit <= 0");
+        require(deadline == 0 || block.timestamp <= deadline, "swapToLimit: deadline exceeded");
+
+        // Read previous balances for affected assets
+        uint256 prevBalI = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
+        uint256 prevBalJ = IERC20(tokens[outputTokenIndex]).balanceOf(address(this));
+
+        // Compute amounts using the same path as views
+        (uint256 totalTransferAmount, uint256 amountOutUint, int128 amountInInternalMax, int128 amountOutInternal, uint256 amountInUsedUint, uint256 feeUint) =
+            _quoteSwapToLimit(inputTokenIndex, outputTokenIndex, limitPrice, swapFeePpm);
+
+        // Transfer the exact amount needed from payer and require exact receipt (revert on fee-on-transfer)
+        tokens[inputTokenIndex].safeTransferFrom(payer, address(this), totalTransferAmount);
+        uint256 balIAfter = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
+        require(balIAfter == prevBalI + totalTransferAmount, "swapToLimit: non-standard tokenIn");
+
+        // Transfer output to receiver and verify exact decrease
+        tokens[outputTokenIndex].safeTransfer(receiver, amountOutUint);
+        uint256 balJAfter = IERC20(tokens[outputTokenIndex]).balanceOf(address(this));
+        require(balJAfter == prevBalJ - amountOutUint, "swapToLimit: non-standard tokenOut");
+
+        // Accrue protocol share (floor) from the fee on input token
+        if (protocolFeePpm > 0 && feeUint > 0 ) {
+            uint256 protoShare = (feeUint * protocolFeePpm) / 1_000_000; // floor
+            if (protoShare > 0) {
+                protocolFeesOwed[inputTokenIndex] += protoShare;
+            }
+        }
+
+        // Update caches to effective balances
+        _recordCachedBalance(inputTokenIndex, balIAfter);
+        _recordCachedBalance(outputTokenIndex, balJAfter);
+
+        // Apply swap to LMSR state with the internal amounts
+        lmsr.applySwap(inputTokenIndex, outputTokenIndex, amountInInternalMax, amountOutInternal);
+
+        // Maintain original event semantics (logs input without fee)
+        emit IPartyPool.Swap(payer, receiver, tokens[inputTokenIndex], tokens[outputTokenIndex], amountInUsedUint, amountOutUint);
+
+        return (amountInUsedUint, amountOutUint, feeUint);
+    }
+
+
+    /// @notice Internal quote for swap-to-limit that mirrors swapToLimit() rounding and fee application
+    /// @dev Computes the input required to reach limitPrice and the resulting output; all rounding matches swapToLimit.
+    /// @return grossIn amount to transfer in (inclusive of fee), amountOutUint output amount (uint),
+    ///         amountInInternal and amountOutInternal (64.64), amountInUintNoFee input amount excluding fee (uint),
+    ///         feeUint fee taken from the gross input (uint)
+    function _quoteSwapToLimit(
+        uint256 inputTokenIndex,
+        uint256 outputTokenIndex,
+        int128 limitPrice,
+        uint256 swapFeePpm
+    ) internal view
+        returns (
+            uint256 grossIn,
+            uint256 amountOutUint,
+            int128 amountInInternal,
+            int128 amountOutInternal,
+            uint256 amountInUintNoFee,
+            uint256 feeUint
+        )
+    {
+        // Compute internal maxima at the price limit
+        (amountInInternal, amountOutInternal) = lmsr.swapAmountsForPriceLimit(inputTokenIndex, outputTokenIndex, limitPrice);
+
+        // Convert input to uint (ceil) and output to uint (floor)
+        amountInUintNoFee = _internalToUintCeil(amountInInternal, bases[inputTokenIndex]);
+        require(amountInUintNoFee > 0, "swapToLimit: input zero");
+
+        feeUint = 0;
+        grossIn = amountInUintNoFee;
+        if (swapFeePpm > 0) {
+            feeUint = _ceilFee(amountInUintNoFee, swapFeePpm);
+            grossIn += feeUint;
+        }
+
+        amountOutUint = _internalToUintFloor(amountOutInternal, bases[outputTokenIndex]);
+        require(amountOutUint > 0, "swapToLimit: output zero");
+    }
+
+}

src/PartyPoolSwapMintImpl.sol

@@ -1,355 +0,0 @@
-// SPDX-License-Identifier: UNLICENSED
-pragma solidity ^0.8.30;
-
-import "@abdk/ABDKMath64x64.sol";
-import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
-import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
-import "./PartyPoolBase.sol";
-import "./LMSRStabilized.sol";
-
-/// @title PartyPoolSwapMintImpl - Implementation contract for swapMint and burnSwap functions
-/// @notice This contract contains the swapMint and burnSwap implementation that will be called via delegatecall
-/// @dev This contract inherits from PartyPoolBase to access storage and internal functions
-contract PartyPoolSwapMintImpl is PartyPoolBase {
-    using ABDKMath64x64 for int128;
-    using LMSRStabilized for LMSRStabilized.State;
-    using SafeERC20 for IERC20;
-
-    // Events that mirror the main contract events
-    event SwapMint(address indexed payer, address indexed receiver, uint256 indexed inputTokenIndex, uint256 totalTransfer, uint256 amountInUint, uint256 feeUintActual);
-    event BurnSwap(address indexed payer, address indexed receiver, uint256 indexed inputTokenIndex, uint256 amountOutUint);
-    event Mint(address indexed payer, address indexed receiver, uint256[] depositAmounts, uint256 lpMinted);
-    event Burn(address indexed payer, address indexed receiver, uint256[] withdrawAmounts, uint256 lpBurned);
-
-
-    /// @notice Single-token mint: deposit a single token, charge swap-LMSR cost, and mint LP.
-    /// @dev swapMint executes as an exact-in planned swap followed by proportional scaling of qInternal.
-    ///      The function emits SwapMint (gross, net, fee) and also emits Mint for LP issuance.
-    /// @param payer who transfers the input token
-    /// @param receiver who receives the minted LP tokens
-    /// @param inputTokenIndex index of the input token
-    /// @param maxAmountIn maximum uint token input (inclusive of fee)
-    /// @param deadline optional deadline
-    /// @param swapFeePpm fee in parts-per-million for this pool
-    /// @return lpMinted actual LP minted (uint)
-    function swapMint(
-        address payer,
-        address receiver,
-        uint256 inputTokenIndex,
-        uint256 maxAmountIn,
-        uint256 deadline,
-        uint256 swapFeePpm
-    ) external returns (uint256 lpMinted, uint256 feeUintActual) {
-        uint256 n = tokens.length;
-        require(inputTokenIndex < n, "swapMint: idx");
-        require(maxAmountIn > 0, "swapMint: input zero");
-        require(deadline == 0 || block.timestamp <= deadline, "swapMint: deadline");
-        require(lmsr.nAssets > 0, "swapMint: uninit pool");
-
-        // compute fee on gross maxAmountIn to get an initial net estimate (we'll recompute based on actual used)
-        (, uint256 netUintGuess) = _computeFee(maxAmountIn, swapFeePpm);
-
-        // Convert the net guess to internal (floor)
-        int128 netInternalGuess = _uintToInternalFloor(netUintGuess, bases[inputTokenIndex]);
-        require(netInternalGuess > int128(0), "swapMint: input too small after fee");
-
-        // Use LMSR view to determine actual internal consumed and size-increase (ΔS) for mint
-        (int128 amountInInternalUsed, int128 sizeIncreaseInternal) = lmsr.swapAmountsForMint(inputTokenIndex, netInternalGuess);
-
-        // amountInInternalUsed may be <= netInternalGuess. Convert to uint (ceil) to determine actual transfer
-        uint256 amountInUint = _internalToUintCeil(amountInInternalUsed, bases[inputTokenIndex]);
-        require(amountInUint > 0, "swapMint: input zero after internal conversion");
-
-        // Compute fee on the actual used input and total transfer amount (ceiling)
-        feeUintActual = _ceilFee(amountInUint, swapFeePpm);
-        uint256 totalTransfer = amountInUint + feeUintActual;
-        require(totalTransfer > 0 && totalTransfer <= maxAmountIn, "swapMint: transfer exceeds max");
-
-        // Record pre-balance and transfer tokens from payer, require exact receipt (revert on fee-on-transfer)
-        uint256 prevBalI = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
-        tokens[inputTokenIndex].safeTransferFrom(payer, address(this), totalTransfer);
-        uint256 balIAfter = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
-        require(balIAfter == prevBalI + totalTransfer, "swapMint: non-standard tokenIn");
-
-        // Update cached uint balances for token inputTokenIndex (implementation writes onchain value; wrapper will set effective)
-        cachedUintBalances[inputTokenIndex] = balIAfter;
-
-        // Compute old and new scaled size metrics to determine LP minted
-        int128 oldTotal = _computeSizeMetric(lmsr.qInternal);
-        require(oldTotal > int128(0), "swapMint: zero total");
-        uint256 oldScaled = ABDKMath64x64.mulu(oldTotal, LP_SCALE);
-
-        int128 newTotal = oldTotal.add(sizeIncreaseInternal);
-        uint256 newScaled = ABDKMath64x64.mulu(newTotal, LP_SCALE);
-
-        uint256 actualLpToMint;
-        // Use natural ERC20 function since base contract inherits from ERC20
-        uint256 currentSupply = _totalSupply;
-        if (currentSupply == 0) {
-            // If somehow supply zero (shouldn't happen as lmsr.nAssets>0), mint newScaled
-            actualLpToMint = newScaled;
-        } else {
-            require(oldScaled > 0, "swapMint: oldScaled zero");
-            uint256 delta = (newScaled > oldScaled) ? (newScaled - oldScaled) : 0;
-            if (delta > 0) {
-                // floor truncation rounds in favor of pool
-                actualLpToMint = (currentSupply * delta) / oldScaled;
-            } else {
-                actualLpToMint = 0;
-            }
-        }
-
-        require(actualLpToMint > 0, "swapMint: zero LP minted");
-
-        // Update LMSR internal state: scale qInternal proportionally by newTotal/oldTotal
-        int128[] memory newQInternal = new int128[](n);
-        for (uint256 idx = 0; idx < n; idx++) {
-            // newQInternal[idx] = qInternal[idx] * (newTotal / oldTotal)
-            newQInternal[idx] = lmsr.qInternal[idx].mul(newTotal).div(oldTotal);
-        }
-
-        // Update cached internal and kappa via updateForProportionalChange
-        lmsr.updateForProportionalChange(newQInternal);
-
-        // Use natural ERC20 function since base contract inherits from ERC20
-        _mint(receiver, actualLpToMint);
-
-        // Emit SwapMint event with gross transfer, net input and fee (planned exact-in)
-        emit SwapMint(payer, receiver, inputTokenIndex, totalTransfer, amountInUint, feeUintActual);
-
-        // Emit standard Mint event which records deposit amounts and LP minted
-        emit Mint(payer, receiver, new uint256[](n), actualLpToMint);
-        // Note: depositAmounts array omitted (empty) since swapMint uses single-token input
-
-        lpMinted = actualLpToMint;
-        return (lpMinted, feeUintActual);
-    }
-
-    /// @notice Calculate the amounts for a swap mint operation
-    /// @dev This is a pure view function that computes swap mint amounts from provided state
-    /// @param inputTokenIndex index of the input token
-    /// @param maxAmountIn maximum amount of token to deposit (inclusive of fee)
-    /// @param swapFeePpm fee in parts-per-million
-    /// @param lmsrState current LMSR state
-    /// @param bases_ scaling bases for each token
-    /// @param totalSupply_ current total LP token supply
-    /// @return amountInUsed actual input amount used (excluding fee)
-    /// @return fee fee amount charged
-    /// @return lpMinted LP tokens that would be minted
-    function swapMintAmounts(
-        uint256 inputTokenIndex,
-        uint256 maxAmountIn,
-        uint256 swapFeePpm,
-        LMSRStabilized.State memory lmsrState,
-        uint256[] memory bases_,
-        uint256 totalSupply_
-    ) public pure returns (uint256 amountInUsed, uint256 fee, uint256 lpMinted) {
-        require(inputTokenIndex < bases_.length, "swapMintAmounts: idx");
-        require(maxAmountIn > 0, "swapMintAmounts: input zero");
-        require(lmsrState.nAssets > 0, "swapMintAmounts: uninit pool");
-
-        // Compute fee on gross maxAmountIn to get an initial net estimate
-        uint256 feeGuess = 0;
-        uint256 netUintGuess = maxAmountIn;
-        if (swapFeePpm > 0) {
-            feeGuess = (maxAmountIn * swapFeePpm + 999999) / 1000000; // ceil fee
-            netUintGuess = maxAmountIn - feeGuess;
-        }
-
-        // Convert the net guess to internal (floor)
-        int128 netInternalGuess = _uintToInternalFloorPure(netUintGuess, bases_[inputTokenIndex]);
-        require(netInternalGuess > int128(0), "swapMintAmounts: input too small after fee");
-
-        // Use LMSR view to determine actual internal consumed and size-increase (ΔS) for mint
-        (int128 amountInInternalUsed, int128 sizeIncreaseInternal) =
-            LMSRStabilized.swapAmountsForMint(lmsrState.nAssets, lmsrState.kappa, lmsrState.qInternal,
-                inputTokenIndex, netInternalGuess);
-
-        // amountInInternalUsed may be <= netInternalGuess. Convert to uint (ceil) to determine actual transfer
-        amountInUsed = _internalToUintCeilPure(amountInInternalUsed, bases_[inputTokenIndex]);
-        require(amountInUsed > 0, "swapMintAmounts: input zero after internal conversion");
-
-        // Compute fee on the actual used input (ceiling)
-        fee = 0;
-        if (swapFeePpm > 0) {
-            fee = (amountInUsed * swapFeePpm + 999999) / 1000000; // ceil fee
-        }
-        uint256 totalTransfer = amountInUsed + fee;
-        require(totalTransfer > 0 && totalTransfer <= maxAmountIn, "swapMintAmounts: transfer exceeds max");
-
-        // Compute old and new scaled size metrics to determine LP minted
-        int128 oldTotal = _computeSizeMetricPure(lmsrState.qInternal);
-        require(oldTotal > int128(0), "swapMintAmounts: zero total");
-        uint256 oldScaled = ABDKMath64x64.mulu(oldTotal, LP_SCALE);
-
-        int128 newTotal = oldTotal.add(sizeIncreaseInternal);
-        uint256 newScaled = ABDKMath64x64.mulu(newTotal, LP_SCALE);
-
-        if (totalSupply_ == 0) {
-            // If somehow supply zero (shouldn't happen as lmsr.nAssets>0), mint newScaled
-            lpMinted = newScaled;
-        } else {
-            require(oldScaled > 0, "swapMintAmounts: oldScaled zero");
-            uint256 delta = (newScaled > oldScaled) ? (newScaled - oldScaled) : 0;
-            if (delta > 0) {
-                // floor truncation rounds in favor of pool
-                lpMinted = (totalSupply_ * delta) / oldScaled;
-            } else {
-                lpMinted = 0;
-            }
-        }
-
-        require(lpMinted > 0, "swapMintAmounts: zero LP minted");
-    }
-
-    /// @notice Calculate the amounts for a burn swap operation
-    /// @dev This is a pure view function that computes burn swap amounts from provided state
-    /// @param lpAmount amount of LP tokens to burn
-    /// @param inputTokenIndex index of target asset to receive
-    /// @param swapFeePpm fee in parts-per-million
-    /// @param lmsrState current LMSR state
-    /// @param bases_ scaling bases for each token
-    /// @param totalSupply_ current total LP token supply
-    /// @return amountOut amount of target asset that would be received
-    function burnSwapAmounts(
-        uint256 lpAmount,
-        uint256 inputTokenIndex,
-        uint256 swapFeePpm,
-        LMSRStabilized.State memory lmsrState,
-        uint256[] memory bases_,
-        uint256 totalSupply_
-    ) public pure returns (uint256 amountOut) {
-        require(inputTokenIndex < bases_.length, "burnSwapAmounts: idx");
-        require(lpAmount > 0, "burnSwapAmounts: zero lp");
-        require(totalSupply_ > 0, "burnSwapAmounts: empty supply");
-
-        // alpha = lpAmount / supply as Q64.64
-        int128 alpha = ABDKMath64x64.divu(lpAmount, totalSupply_) // fraction of total supply to burn
-            .mul(ABDKMath64x64.divu(1000000-swapFeePpm, 1000000)); // adjusted for fee
-
-        // Use LMSR view to compute single-asset payout and burned size-metric
-        (int128 payoutInternal, ) = LMSRStabilized.swapAmountsForBurn(lmsrState.nAssets, lmsrState.kappa, lmsrState.qInternal,
-            inputTokenIndex, alpha);
-
-        // Convert payoutInternal -> uint (floor) to favor pool
-        amountOut = _internalToUintFloorPure(payoutInternal, bases_[inputTokenIndex]);
-        require(amountOut > 0, "burnSwapAmounts: output zero");
-    }
-
-    /// @notice Burn LP tokens then swap the redeemed proportional basket into a single asset `inputTokenIndex` and send to receiver.
-    /// @dev The function burns LP tokens (authorization via allowance if needed), sends the single-asset payout and updates LMSR state.
-    /// @param payer who burns LP tokens
-    /// @param receiver who receives the single asset
-    /// @param lpAmount amount of LP tokens to burn
-    /// @param inputTokenIndex index of target asset to receive
-    /// @param deadline optional deadline
-    /// @param swapFeePpm fee in parts-per-million for this pool (may be used for future fee logic)
-    /// @return amountOutUint uint amount of asset i sent to receiver
-    function burnSwap(
-        address payer,
-        address receiver,
-        uint256 lpAmount,
-        uint256 inputTokenIndex,
-        uint256 deadline,
-        uint256 swapFeePpm
-    ) external returns (uint256 amountOutUint, uint256 feeTokenUint) {
-        uint256 n = tokens.length;
-        require(inputTokenIndex < n, "burnSwap: idx");
-        require(lpAmount > 0, "burnSwap: zero lp");
-        require(deadline == 0 || block.timestamp <= deadline, "burnSwap: deadline");
-
-        uint256 supply = _totalSupply;
-        require(supply > 0, "burnSwap: empty supply");
-        require(_balances[payer] >= lpAmount, "burnSwap: insufficient LP");
-
-        // alpha = lpAmount / supply as Q64.64 (adjusted for fee)
-        int128 alpha = ABDKMath64x64.divu(lpAmount, supply) // fraction of total supply to burn
-            .mul(ABDKMath64x64.divu(1000000-swapFeePpm, 1000000)); // adjusted for fee
-
-        // Use LMSR view to compute single-asset payout and burned size-metric
-        (int128 payoutInternal, ) = lmsr.swapAmountsForBurn(inputTokenIndex, alpha);
-
-        // Convert payoutInternal -> uint (floor) to favor pool
-        amountOutUint = _internalToUintFloor(payoutInternal, bases[inputTokenIndex]);
-        require(amountOutUint > 0, "burnSwap: output zero");
-
-        // Compute gross payout (no swap fee) so we can determine token-side fee = gross - net
-        int128 alphaGross = ABDKMath64x64.divu(lpAmount, supply); // gross fraction (no swap fee)
-        (int128 payoutGrossInternal, ) = lmsr.swapAmountsForBurn(inputTokenIndex, alphaGross);
-        uint256 payoutGrossUint = _internalToUintFloor(payoutGrossInternal, bases[inputTokenIndex]);
-        feeTokenUint = (payoutGrossUint > amountOutUint) ? (payoutGrossUint - amountOutUint) : 0;
-
-        // Transfer the payout to receiver
-        tokens[inputTokenIndex].safeTransfer(receiver, amountOutUint);
-
-        // Burn LP tokens from payer (authorization via allowance)
-        if (msg.sender != payer) {
-            uint256 allowed = _allowances[payer][msg.sender];
-            require(allowed >= lpAmount, "burnSwap: allowance insufficient");
-            _approve(payer, msg.sender, allowed - lpAmount);
-        }
-        _burn(payer, lpAmount);
-
-        // Update cached balances by reading on-chain balances for all tokens
-        int128[] memory newQInternal = new int128[](n);
-        for (uint256 idx = 0; idx < n; idx++) {
-            uint256 bal = IERC20(tokens[idx]).balanceOf(address(this));
-            // implementation writes raw onchain values; wrapper will set effective cached values
-            cachedUintBalances[idx] = bal;
-            newQInternal[idx] = _uintToInternalFloor(bal, bases[idx]);
-        }
-
-        // Emit BurnSwap with public-facing info only (do not expose ΔS or LP burned)
-        emit BurnSwap(payer, receiver, inputTokenIndex, amountOutUint);
-
-        // If entire pool drained, deinit; else update proportionally
-        bool allZero = true;
-        for (uint256 idx = 0; idx < n; idx++) {
-            if (newQInternal[idx] != int128(0)) { allZero = false; break; }
-        }
-        if (allZero) {
-            lmsr.deinit();
-        } else {
-            lmsr.updateForProportionalChange(newQInternal);
-        }
-
-        emit Burn(payer, receiver, new uint256[](n), lpAmount);
-        return (amountOutUint, feeTokenUint);
-    }
-
-    /// @notice Pure version of _uintToInternalFloor for use in view functions
-    function _uintToInternalFloorPure(uint256 amount, uint256 base) internal pure returns (int128) {
-        // amount / base as Q64.64, floored
-        return ABDKMath64x64.divu(amount, base);
-    }
-
-    /// @notice Pure version of _internalToUintCeil for use in view functions
-    function _internalToUintCeilPure(int128 amount, uint256 base) internal pure returns (uint256) {
-        // Convert Q64.64 to uint with ceiling: ceil(amount * base)
-        // Use mulu which floors, then add remainder check for ceiling
-        uint256 floored = ABDKMath64x64.mulu(amount, base);
-        // Check if there's a fractional part by computing amount * base - floored
-        int128 baseQ64 = ABDKMath64x64.fromUInt(base);
-        int128 flooredQ64 = ABDKMath64x64.fromUInt(floored);
-        int128 product = amount.mul(baseQ64);
-        if (product > flooredQ64) {
-            return floored + 1; // Ceiling
-        }
-        return floored;
-    }
-
-    /// @notice Pure version of _internalToUintFloor for use in view functions
-    function _internalToUintFloorPure(int128 amount, uint256 base) internal pure returns (uint256) {
-        // Convert Q64.64 to uint with floor: floor(amount * base)
-        return ABDKMath64x64.mulu(amount, base);
-    }
-
-    /// @notice Pure version of _computeSizeMetric for use in view functions
-    function _computeSizeMetricPure(int128[] memory qInternal) internal pure returns (int128) {
-        int128 sum = int128(0);
-        for (uint256 i = 0; i < qInternal.length; i++) {
-            sum = sum.add(qInternal[i]);
-        }
-        return sum;
-    }
-}

test/PartyPlanner.t.sol

@@ -1,12 +1,20 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.30;
 
-import "forge-std/Test.sol";
+import {CommonBase} from "../lib/forge-std/src/Base.sol";
-import "../src/LMSRStabilized.sol";
+import {StdAssertions} from "../lib/forge-std/src/StdAssertions.sol";
-import "../src/PartyPlanner.sol";
+import {StdChains} from "../lib/forge-std/src/StdChains.sol";
-import "../src/PartyPool.sol";
+import {StdCheats, StdCheatsSafe} from "../lib/forge-std/src/StdCheats.sol";
+import {StdUtils} from "../lib/forge-std/src/StdUtils.sol";
+import {Test} from "../lib/forge-std/src/Test.sol";
+import {ERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/ERC20.sol";
+import {IERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
 import {Deploy} from "../src/Deploy.sol";
-import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import {IPartyPool} from "../src/IPartyPool.sol";
+import {LMSRStabilized} from "../src/LMSRStabilized.sol";
+import {PartyPlanner} from "../src/PartyPlanner.sol";
+import {PartyPool} from "../src/PartyPool.sol";
+import {MockERC20} from "./PartyPlanner.t.sol";
 
 // Mock ERC20 token for testing
 contract MockERC20 is ERC20 {
@@ -88,7 +96,7 @@ contract PartyPlannerTest is Test {
         // Compute kappa then create pool via kappa overload
         int128 computedKappa = LMSRStabilized.computeKappaFromSlippage(tokens.length, tradeFrac, targetSlippage);
 
-        (PartyPool pool, uint256 lpAmount) = planner.newPool(
+        (IPartyPool pool, uint256 lpAmount) = planner.newPool(
             name,
             symbol,
             tokens,
@@ -117,7 +125,7 @@ contract PartyPlannerTest is Test {
         assertEq(planner.poolsByTokenCount(IERC20(address(tokenB))), initialTokenBCount + 1, "TokenB pool count should increase");
 
         // Verify pools can be retrieved
-        PartyPool[] memory allPools = planner.getAllPools(0, 10);
+        IPartyPool[] memory allPools = planner.getAllPools(0, 10);
         bool poolFound = false;
         for (uint256 i = 0; i < allPools.length; i++) {
             if (allPools[i] == pool) {
@@ -128,7 +136,7 @@ contract PartyPlannerTest is Test {
         assertTrue(poolFound, "Created pool should be in getAllPools result");
 
         // Verify pool appears in token-specific queries
-        PartyPool[] memory tokenAPools = planner.getPoolsByToken(IERC20(address(tokenA)), 0, 10);
+        IPartyPool[] memory tokenAPools = planner.getPoolsByToken(IERC20(address(tokenA)), 0, 10);
         bool poolInTokenA = false;
         for (uint256 i = 0; i < tokenAPools.length; i++) {
             if (tokenAPools[i] == pool) {
@@ -138,7 +146,7 @@ contract PartyPlannerTest is Test {
         }
         assertTrue(poolInTokenA, "Pool should be indexed under tokenA");
 
-        PartyPool[] memory tokenBPools = planner.getPoolsByToken(IERC20(address(tokenB)), 0, 10);
+        IPartyPool[] memory tokenBPools = planner.getPoolsByToken(IERC20(address(tokenB)), 0, 10);
         bool poolInTokenB = false;
         for (uint256 i = 0; i < tokenBPools.length; i++) {
             if (tokenBPools[i] == pool) {
@@ -167,7 +175,7 @@ contract PartyPlannerTest is Test {
         deposits1[1] = INITIAL_DEPOSIT_AMOUNT;
 
         int128 kappa1 = LMSRStabilized.computeKappaFromSlippage(tokens1.length, int128((1 << 64) - 1), int128(1 << 62));
-        (PartyPool pool1,) = planner.newPool(
+        (IPartyPool pool1,) = planner.newPool(
             "Pool 1", "LP1", tokens1, bases1,
             kappa1, 3000, 5000, false,
             payer, receiver, deposits1, 1000e18, 0
@@ -187,7 +195,7 @@ contract PartyPlannerTest is Test {
         deposits2[1] = INITIAL_DEPOSIT_AMOUNT / 1e12; // Adjust for 6 decimals
 
         int128 kappa2 = LMSRStabilized.computeKappaFromSlippage(tokens2.length, int128((1 << 64) - 1), int128(1 << 62));
-        (PartyPool pool2,) = planner.newPool(
+        (IPartyPool pool2,) = planner.newPool(
             "Pool 2", "LP2", tokens2, bases2,
             kappa2, 3000, 5000, false,
             payer, receiver, deposits2, 1000e18, 0
@@ -203,7 +211,7 @@ contract PartyPlannerTest is Test {
         assertEq(planner.poolsByTokenCount(IERC20(address(tokenC))), 1, "TokenC should be in 1 pool");
 
         // Verify tokenB appears in both pools
-        PartyPool[] memory tokenBPools = planner.getPoolsByToken(IERC20(address(tokenB)), 0, 10);
+        IPartyPool[] memory tokenBPools = planner.getPoolsByToken(IERC20(address(tokenB)), 0, 10);
         assertEq(tokenBPools.length, 2, "TokenB should have 2 pools");
 
         bool pool1Found = false;
@@ -274,7 +282,7 @@ contract PartyPlannerTest is Test {
     function test_poolIndexing_Pagination() public {
         // Create multiple pools for pagination testing
         uint256 numPools = 5;
-        PartyPool[] memory createdPools = new PartyPool[](numPools);
+        IPartyPool[] memory createdPools = new IPartyPool[](numPools);
 
         for (uint256 i = 0; i < numPools; i++) {
             IERC20[] memory tokens = new IERC20[](2);
@@ -290,7 +298,7 @@ contract PartyPlannerTest is Test {
             deposits[1] = INITIAL_DEPOSIT_AMOUNT;
 
             int128 kappaLoop = LMSRStabilized.computeKappaFromSlippage(tokens.length, int128((1 << 64) - 1), int128(1 << 62));
-            (PartyPool pool,) = planner.newPool(
+            (IPartyPool pool,) = planner.newPool(
                 string(abi.encodePacked("Pool ", vm.toString(i))),
                 string(abi.encodePacked("LP", vm.toString(i))),
                 tokens, bases,
@@ -304,19 +312,19 @@ contract PartyPlannerTest is Test {
         assertEq(planner.poolCount(), numPools, "Should have created all pools");
 
         // Test pagination - get first 3 pools
-        PartyPool[] memory page1 = planner.getAllPools(0, 3);
+        IPartyPool[] memory page1 = planner.getAllPools(0, 3);
         assertEq(page1.length, 3, "First page should have 3 pools");
 
         // Test pagination - get next 2 pools
-        PartyPool[] memory page2 = planner.getAllPools(3, 3);
+        IPartyPool[] memory page2 = planner.getAllPools(3, 3);
         assertEq(page2.length, 2, "Second page should have 2 pools");
 
         // Test pagination - offset beyond bounds
-        PartyPool[] memory emptyPage = planner.getAllPools(10, 3);
+        IPartyPool[] memory emptyPage = planner.getAllPools(10, 3);
         assertEq(emptyPage.length, 0, "Should return empty array for out of bounds offset");
 
         // Verify all pools are accessible through pagination
-        PartyPool[] memory allPools = planner.getAllPools(0, 10);
+        IPartyPool[] memory allPools = planner.getAllPools(0, 10);
         assertEq(allPools.length, numPools, "Should return all pools");
 
         for (uint256 i = 0; i < numPools; i++) {