feat: update stable swap impl

.github/workflows/test.yml

@@ -25,5 +25,6 @@ jobs:
         run: forge test -vvv --isolate
         env:
           FOUNDRY_PROFILE: ci
+          ARB_FORK_URL: ${{ secrets.ARB_FORK_URL }}
           FORK_URL: ${{ secrets.FORK_URL }}
           TESTNET_FORK_URL: ${{ secrets.TESTNET_FORK_URL }}

README.md

@@ -9,6 +9,9 @@
 // testnet fork test for infinity, mainnet fork test for v2/v3 
 export FORK_URL=https://bsc-mainnet.nodereal.io/v1/xxx
 export TESTNET_FORK_URL=https://bsc-testnet.nodereal.io/v1/xxx
+
+(If not set, only Arb test will fail)
+export ARB_FORK_URL=https://xxx
 ```
 
 3. Run test with `forge test`

script/deployParameters/mainnet/DeployArbitrum.s.sol

@@ -15,7 +15,7 @@ import {RouterParameters} from "../../../src/base/RouterImmutables.sol";
 contract DeployArbitrum is DeployUniversalRouter {
     /// @notice contract address will be based on deployment salt
     function getDeploymentSalt() public pure override returns (bytes32) {
-        return keccak256("INFINITY-UNIVERSAL-ROUTER/UniversalRouter/0.0001");
+        return keccak256("INFINITY-UNIVERSAL-ROUTER/UniversalRouter/1.0.1");
     }
 
     function setUp() public override {
@@ -27,8 +27,8 @@ contract DeployArbitrum is DeployUniversalRouter {
             v3Deployer: 0x41ff9AA7e16B8B1a8a8dc4f0eFacd93D02d071c9,
             v2InitCodeHash: 0x57224589c67f3f30a6b0d7a1b54cf3153ab84563bc609ef41dfb34f8b2974d2d,
             v3InitCodeHash: 0x6ce8eb472fa82df5469c6ab6d485f17c3ad13c8cd7af59b3d4a8026c5ce0f7e2,
-            stableFactory: UNSUPPORTED_PROTOCOL,
-            stableInfo: UNSUPPORTED_PROTOCOL,
+            stableFactory: 0x5D5fBB19572c4A89846198c3DBEdB2B6eF58a77a,
+            stableInfo: 0x58B2F00f74a1877510Ec37b22f116Bf5D63Ab1b0,
             infiVault: UNSUPPORTED_PROTOCOL,
             infiClPoolManager: UNSUPPORTED_PROTOCOL,
             infiBinPoolManager: UNSUPPORTED_PROTOCOL,

script/deployParameters/mainnet/DeployBase.s.sol

@@ -15,7 +15,7 @@ import {RouterParameters} from "../../../src/base/RouterImmutables.sol";
 contract DeployBase is DeployUniversalRouter {
     /// @notice contract address will be based on deployment salt
     function getDeploymentSalt() public pure override returns (bytes32) {
-        return keccak256("INFINITY-UNIVERSAL-ROUTER/UniversalRouter/1.0.0");
+        return keccak256("INFINITY-UNIVERSAL-ROUTER/UniversalRouter/1.0.1");
     }
 
     function setUp() public override {

script/deployParameters/mainnet/DeployBsc.s.sol

@@ -15,7 +15,7 @@ import {RouterParameters} from "../../../src/base/RouterImmutables.sol";
 contract DeployBsc is DeployUniversalRouter {
     /// @notice contract address will be based on deployment salt
     function getDeploymentSalt() public pure override returns (bytes32) {
-        return keccak256("INFINITY-UNIVERSAL-ROUTER/UniversalRouter/1.0.0");
+        return keccak256("INFINITY-UNIVERSAL-ROUTER/UniversalRouter/1.0.1");
     }
 
     function setUp() public override {

snapshots/StableSwapArbTest.json

@@ -0,0 +1,4 @@
+{
+  "test_stableSwap_ExactInput0For1": "199963",
+  "test_stableSwap_ExactInput1For0": "198844"
+}

snapshots/StableSwapBusdUsdcTest.json

@@ -1,4 +1,4 @@
 {
-  "test_stableSwap_ExactInput0For1": "193988",
-  "test_stableSwap_ExactInput1For0": "194056"
+  "test_stableSwap_ExactInput0For1": "192796",
+  "test_stableSwap_ExactInput1For0": "193476"
 }

snapshots/StableSwapMultiHop.json

@@ -0,0 +1,12 @@
+{
+  "test_stableSwap_ExactInput0For1_DualAction_FromRouter": "284892",
+  "test_stableSwap_ExactInput0For1_DualAction_FromUser": "311128",
+  "test_stableSwap_ExactInput0For1_MultiCommand_FromRouter": "417203",
+  "test_stableSwap_ExactInput0For1_MultiHop_FromRouter": "281388",
+  "test_stableSwap_ExactInput0For1_MultiHop_FromUser": "307624",
+  "test_stableSwap_ExactInput0For1_SamePath_FromRouter": "246485",
+  "test_stableSwap_ExactInput0For1_SamePath_FromUser": "292156",
+  "test_stableSwap_ExactOut0For1_MultiHop_FromUser": "355264",
+  "test_stableSwap_ExactOutput0For1_MultiHop_FromRouter": "329029",
+  "test_stableSwap_ExactOutput0For1_SamePath_FromRouter": "294699"
+}

snapshots/UniversalRouterTest.json

@@ -1,4 +1,4 @@
 {
-  "UniversalRouterBytecodeSize": "24350",
+  "UniversalRouterBytecodeSize": "24510",
   "test_sweep_token": "55429"
 }

src/modules/pancakeswap/StableSwapRouter.sol

@@ -45,18 +45,32 @@ abstract contract StableSwapRouter is RouterImmutables, Permit2Payments, Ownable
         emit SetStableSwap(stableSwapFactory, stableSwapInfo);
     }
 
-    function _stableSwap(address[] calldata path, uint256[] calldata flag) private {
-        unchecked {
-            if (path.length - 1 != flag.length) revert StableInvalidPath();
-
-            for (uint256 i; i < flag.length; i++) {
-                (address input, address output) = (path[i], path[i + 1]);
-                (uint256 k, uint256 j, address swapContract) = stableSwapFactory.getStableInfo(input, output, flag[i]);
-                uint256 amountIn = ERC20(input).balanceOf(address(this));
-                ERC20(input).safeApprove(swapContract, amountIn);
-                IStableSwap(swapContract).exchange(k, j, amountIn, 0);
-            }
+    /// @dev if a single hop, path would be of size 2, and flag would be of size 1
+    /// if 2 hops, path would be of size 3, and flag would be of size 2
+    /// @return amtOut The amount of output tokens received after the swap
+    function _stableSwap(address[] calldata path, uint256[] calldata flag, uint256 amountIn)
+        private
+        returns (uint256 amtOut)
+    {
+        if (path.length - 1 != flag.length) revert StableInvalidPath();
+
+        uint256 outputTokenBal;
+        for (uint256 i; i < flag.length; i++) {
+            (address input, address output) = (path[i], path[i + 1]);
+
+            outputTokenBal = ERC20(output).balanceOf(address(this));
+
+            (uint256 k, uint256 j, address swapContract) = stableSwapFactory.getStableInfo(input, output, flag[i]);
+            ERC20(input).safeApprove(swapContract, amountIn);
+            IStableSwap(swapContract).exchange(k, j, amountIn, 0);
+
+            // Update amountIn for the next hop. this is done as swapContract do not return the output amount
+            // If this is the last hop, amountIn is the output amount
+            amountIn = ERC20(output).balanceOf(address(this)) - outputTokenBal;
         }
+
+        // after the swap iterations, amountIn is the output amount
+        amtOut = amountIn;
     }
 
     /// @notice Performs a PancakeSwap stable exact input swap
@@ -78,15 +92,15 @@ abstract contract StableSwapRouter is RouterImmutables, Permit2Payments, Ownable
             payOrPermit2Transfer(path[0], payer, address(this), amountIn);
         }
 
-        ERC20 tokenOut = ERC20(path[path.length - 1]);
-        uint256 balanceBefore = tokenOut.balanceOf(address(this));
-
-        _stableSwap(path, flag);
+        if (amountIn == ActionConstants.CONTRACT_BALANCE) {
+            amountIn = ERC20(path[0]).balanceOf(address(this));
+        }
 
-        uint256 amountOut = tokenOut.balanceOf(address(this)) - balanceBefore;
+        uint256 amountOut = _stableSwap(path, flag, amountIn);
         if (amountOut < amountOutMinimum) revert StableTooLittleReceived();
 
-        if (recipient != address(this)) pay(address(tokenOut), recipient, amountOut);
+        address tokenOut = path[path.length - 1];
+        if (recipient != address(this)) pay(tokenOut, recipient, amountOut);
     }
 
     /// @notice Performs a PancakeSwap stable exact output swap
@@ -106,7 +120,7 @@ abstract contract StableSwapRouter is RouterImmutables, Permit2Payments, Ownable
     ) internal {
         payOrPermit2Transfer(path[0], payer, address(this), amountIn);
 
-        _stableSwap(path, flag);
+        _stableSwap(path, flag, amountIn);
 
         if (recipient != address(this)) pay(path[path.length - 1], recipient, amountOut);
     }

test/stableSwap/StableSwap.t.sol

@@ -157,6 +157,26 @@ abstract contract StableSwapTest is Test {
         assertGt(ERC20(token1()).balanceOf(FROM), BALANCE); // token1 received
     }
 
+    function test_stableSwap_ExactInput0For1_Twice_FromRouter() public {
+        bytes memory commands =
+            abi.encodePacked(bytes1(uint8(Commands.STABLE_SWAP_EXACT_IN)), bytes1(uint8(Commands.STABLE_SWAP_EXACT_IN)));
+        deal(token0(), address(router), AMOUNT * 2);
+
+        address[] memory path = new address[](2);
+        path[0] = token0();
+        path[1] = token1();
+
+        // equivalent: abi.decode(inputs, (address, uint256, uint256, address[], uint256[], bool)
+        // recipient, amountIn, amountOutMin, path, flag, payerIsUser
+        bytes[] memory inputs = new bytes[](2);
+        inputs[0] = abi.encode(ActionConstants.MSG_SENDER, AMOUNT, 0, path, flag(), false);
+        inputs[1] = abi.encode(ActionConstants.MSG_SENDER, AMOUNT, 0, path, flag(), false);
+
+        router.execute(commands, inputs);
+        assertEq(ERC20(token0()).balanceOf(FROM), BALANCE); // no token0 taken from user, taken from router
+        assertGt(ERC20(token1()).balanceOf(FROM), BALANCE); // token1 received
+    }
+
     function test_stableSwap_exactInput1For0FromRouter() public {
         bytes memory commands = abi.encodePacked(bytes1(uint8(Commands.STABLE_SWAP_EXACT_IN)));
         deal(token1(), address(router), AMOUNT);
@@ -251,6 +271,26 @@ abstract contract StableSwapTest is Test {
         assertEq(ERC20(token1()).balanceOf(FROM), BALANCE); // no token1 taken from user, taken from router
     }
 
+    function test_stableSwap_exactOutput1For0FromRouter_Twice_FromRouter() public {
+        bytes memory commands = abi.encodePacked(
+            bytes1(uint8(Commands.STABLE_SWAP_EXACT_OUT)), bytes1(uint8(Commands.STABLE_SWAP_EXACT_OUT))
+        );
+        deal(token1(), address(router), BALANCE * 2);
+
+        // equivalent: abi.decode(inputs, (address, uint256, uint256, address[], uint256[], bool)
+        address[] memory path = new address[](2);
+        path[0] = token1();
+        path[1] = token0();
+
+        bytes[] memory inputs = new bytes[](2);
+        inputs[0] = abi.encode(ActionConstants.MSG_SENDER, AMOUNT, type(uint256).max, path, flag(), false);
+        inputs[1] = abi.encode(ActionConstants.MSG_SENDER, AMOUNT, type(uint256).max, path, flag(), false);
+
+        router.execute(commands, inputs);
+        assertGe(ERC20(token0()).balanceOf(FROM), BALANCE + AMOUNT * 2);
+        assertEq(ERC20(token1()).balanceOf(FROM), BALANCE); // no token1 taken from user, taken from router
+    }
+
     function token0() internal virtual returns (address);
     function token1() internal virtual returns (address);
     function flag() internal virtual returns (uint256[] memory);

test/stableSwap/StableSwapArb.t.sol

@@ -0,0 +1,125 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+import {Test, console} from "forge-std/Test.sol";
+import {ERC20} from "solmate/src/tokens/ERC20.sol";
+import {ActionConstants} from "infinity-periphery/src/libraries/ActionConstants.sol";
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+
+import {IPermit2} from "permit2/src/interfaces/IPermit2.sol";
+import {UniversalRouter} from "../../src/UniversalRouter.sol";
+import {Constants} from "../../src/libraries/Constants.sol";
+import {Commands} from "../../src/libraries/Commands.sol";
+import {RouterParameters} from "../../src/base/RouterImmutables.sol";
+import {IStableSwapFactory} from "../../src/interfaces/IStableSwapFactory.sol";
+import {IStableSwapInfo} from "../../src/interfaces/IStableSwapInfo.sol";
+
+/// @dev test stableSwap on arbitrum
+contract StableSwapArbTest is Test {
+    address constant RECIPIENT = address(10);
+    uint256 constant AMOUNT = 1 ether;
+    uint256 constant BALANCE = 100000 ether;
+    ERC20 constant WETH9 = ERC20(0x82aF49447D8a07e3bd95BD0d56f35241523fBab1);
+    IPermit2 constant PERMIT2 = IPermit2(0x31c2F6fcFf4F8759b3Bd5Bf0e1084A055615c768);
+    address constant FROM = address(1234);
+
+    /// @dev StableInfo refers to PancakeStableSwapTwoPoolInfo, threePoolInfo is not present as its not used in PCS
+    IStableSwapFactory STABLE_FACTORY = IStableSwapFactory(0x5D5fBB19572c4A89846198c3DBEdB2B6eF58a77a);
+    IStableSwapInfo STABLE_INFO = IStableSwapInfo(0x58B2F00f74a1877510Ec37b22f116Bf5D63Ab1b0);
+
+    UniversalRouter public router;
+
+    address public PENDLE = 0x0c880f6761F1af8d9Aa9C466984b80DAb9a8c9e8;
+    address public mPENDLE = 0xB688BA096b7Bb75d7841e47163Cd12D18B36A5bF;
+
+    function setUp() public {
+        // Arb: Jun-18-2025 08:50:07 AM +UTC
+        vm.createSelectFork(vm.envString("ARB_FORK_URL"), 348588274);
+
+        RouterParameters memory params = RouterParameters({
+            permit2: address(PERMIT2),
+            weth9: address(WETH9),
+            v2Factory: address(0),
+            v3Factory: address(0),
+            v3Deployer: address(0),
+            v2InitCodeHash: bytes32(0),
+            v3InitCodeHash: bytes32(0),
+            stableFactory: address(STABLE_FACTORY),
+            stableInfo: address(STABLE_INFO),
+            infiVault: address(0),
+            infiClPoolManager: address(0),
+            infiBinPoolManager: address(0),
+            v3NFTPositionManager: address(0),
+            infiClPositionManager: address(0),
+            infiBinPositionManager: address(0)
+        });
+        router = new UniversalRouter(params);
+
+        // pair doesn't exist, revert to keep this test simple without adding to lp etc
+        // Pendle-mPendle: https://arbiscan.io/address/0x73ed25e04Aa673ddf7411441098fC5ae19976CE0
+        if (STABLE_FACTORY.getPairInfo(PENDLE, mPENDLE).swapContract == address(0)) {
+            revert("Pair doesn't exist");
+        }
+
+        vm.startPrank(FROM);
+        deal(FROM, BALANCE);
+        deal(PENDLE, FROM, BALANCE);
+        deal(mPENDLE, FROM, BALANCE);
+        ERC20(PENDLE).approve(address(PERMIT2), type(uint256).max);
+        ERC20(mPENDLE).approve(address(PERMIT2), type(uint256).max);
+        PERMIT2.approve(PENDLE, address(router), type(uint160).max, type(uint48).max);
+        PERMIT2.approve(mPENDLE, address(router), type(uint160).max, type(uint48).max);
+    }
+
+    function test_stableSwap_ExactInput0For1() public {
+        bytes memory commands = abi.encodePacked(bytes1(uint8(Commands.STABLE_SWAP_EXACT_IN)));
+
+        // equivalent: abi.decode(inputs, (address, uint256, uint256, address[], uint256[], bool)
+        address[] memory path = new address[](2);
+        path[0] = PENDLE;
+        path[1] = mPENDLE;
+        bytes[] memory inputs = new bytes[](1);
+        inputs[0] = abi.encode(ActionConstants.MSG_SENDER, AMOUNT, 0, path, flag(), true);
+
+        router.execute(commands, inputs);
+        vm.snapshotGasLastCall("test_stableSwap_ExactInput0For1");
+        assertEq(ERC20(PENDLE).balanceOf(FROM), BALANCE - AMOUNT);
+        assertGt(ERC20(mPENDLE).balanceOf(FROM), BALANCE);
+    }
+
+    function test_stableSwap_ExactInput1For0() public {
+        bytes memory commands = abi.encodePacked(bytes1(uint8(Commands.STABLE_SWAP_EXACT_IN)));
+
+        // equivalent: abi.decode(inputs, (address, uint256, uint256, address[], uint256[], bool)
+        address[] memory path = new address[](2);
+        path[0] = mPENDLE;
+        path[1] = PENDLE;
+        bytes[] memory inputs = new bytes[](1);
+        inputs[0] = abi.encode(ActionConstants.MSG_SENDER, AMOUNT, 0, path, flag(), true);
+
+        router.execute(commands, inputs);
+        vm.snapshotGasLastCall("test_stableSwap_ExactInput1For0");
+        assertEq(ERC20(mPENDLE).balanceOf(FROM), BALANCE - AMOUNT);
+        assertGt(ERC20(PENDLE).balanceOf(FROM), BALANCE);
+    }
+
+    function test_stableSwap_exactInput0For1FromRouter_Arb() public {
+        bytes memory commands = abi.encodePacked(bytes1(uint8(Commands.STABLE_SWAP_EXACT_IN)));
+        deal(PENDLE, address(router), AMOUNT);
+        // equivalent: abi.decode(inputs, (address, uint256, uint256, address[], uint256[], bool)
+        address[] memory path = new address[](2);
+        path[0] = PENDLE;
+        path[1] = mPENDLE;
+        bytes[] memory inputs = new bytes[](1);
+        inputs[0] = abi.encode(ActionConstants.MSG_SENDER, AMOUNT, 0, path, flag(), false);
+
+        router.execute(commands, inputs);
+        assertEq(ERC20(PENDLE).balanceOf(FROM), BALANCE); // no token0 taken from user, taken from router
+        assertGt(ERC20(mPENDLE).balanceOf(FROM), BALANCE); // token1 received
+    }
+
+    function flag() internal pure returns (uint256[] memory pairFlag) {
+        pairFlag = new uint256[](1);
+        pairFlag[0] = 2; // 2 is the flag to indicate StableSwapTwoPool
+    }
+}