feat: update stable swap impl

foundry.toml

@@ -3,7 +3,7 @@ src = "src"
 out = 'foundry-out'
 libs = ["lib"]
 via_ir = true
-optimizer_runs = 20_000
+optimizer_runs = 10_000
 ffi = true
 fs_permissions = [
     { access = "read-write", path = ".forge-snapshots/" },

snapshots/BinNativePancakeSwapInfinityTest.json

@@ -1,6 +1,6 @@
 {
-  "test_infiBinSwap_ExactInSingle_NativeIn": "128929",
-  "test_infiBinSwap_ExactInSingle_NativeOut": "117627",
-  "test_infiBinSwap_ExactInSingle_NativeOut_RouterRecipient": "117916",
+  "test_infiBinSwap_ExactInSingle_NativeIn": "129706",
+  "test_infiBinSwap_ExactInSingle_NativeOut": "118434",
+  "test_infiBinSwap_ExactInSingle_NativeOut_RouterRecipient": "118729",
   "test_infiBinSwap_infiInitializeBinPool": "132846"
 }

snapshots/BinPancakeSwapInfinityTest.json

@@ -1,9 +1,9 @@
 {
-  "test_infiBinSwap_ExactInSingle": "143529",
-  "test_infiBinSwap_ExactIn_MultiHop": "174923",
-  "test_infiBinSwap_ExactIn_SingleHop": "145327",
-  "test_infiBinSwap_ExactOutSingle": "147891",
-  "test_infiBinSwap_ExactOut_MultiHop": "178771",
-  "test_infiBinSwap_ExactOut_SingleHop": "149699",
+  "test_infiBinSwap_ExactInSingle": "144336",
+  "test_infiBinSwap_ExactIn_MultiHop": "176216",
+  "test_infiBinSwap_ExactIn_SingleHop": "146134",
+  "test_infiBinSwap_ExactOutSingle": "148590",
+  "test_infiBinSwap_ExactOut_MultiHop": "179848",
+  "test_infiBinSwap_ExactOut_SingleHop": "150398",
   "test_infiBinSwap_InitializeBinPool": "152986"
 }

snapshots/CLNativePancakeSwapInfinityTest.json

@@ -1,5 +1,5 @@
 {
-  "test_infiClSwap_ExactInSingle_NativeIn": "161970",
-  "test_infiClSwap_ExactInSingle_NativeOut": "144706",
+  "test_infiClSwap_ExactInSingle_NativeIn": "162424",
+  "test_infiClSwap_ExactInSingle_NativeOut": "145166",
   "test_infiClSwap_infiInitializeClPool": "133628"
 }

snapshots/CLPancakeSwapInfinityTest.json

@@ -1,9 +1,9 @@
 {
-  "test_infiClSwap_ExactInSingle": "176570",
-  "test_infiClSwap_ExactIn_MultiHop": "241016",
-  "test_infiClSwap_ExactIn_SingleHop": "178361",
-  "test_infiClSwap_ExactOutSingle": "180894",
-  "test_infiClSwap_ExactOut_MultiHop": "244792",
-  "test_infiClSwap_ExactOut_SingleHop": "182697",
+  "test_infiClSwap_ExactInSingle": "177054",
+  "test_infiClSwap_ExactIn_MultiHop": "241651",
+  "test_infiClSwap_ExactIn_SingleHop": "178839",
+  "test_infiClSwap_ExactOutSingle": "181368",
+  "test_infiClSwap_ExactOut_MultiHop": "245419",
+  "test_infiClSwap_ExactOut_SingleHop": "183171",
   "test_infiClSwap_infiInitializeClPool": "153756"
 }

snapshots/StableSwapBusdUsdcTest.json

@@ -1,4 +1,4 @@
 {
-  "test_stableSwap_ExactInput0For1": "193988",
-  "test_stableSwap_ExactInput1For0": "194056"
+  "test_stableSwap_ExactInput0For1": "192810",
+  "test_stableSwap_ExactInput1For0": "193490"
 }

snapshots/StableSwapMultiHop.json

@@ -0,0 +1,11 @@
+{
+  "test_stableSwap_ExactInput0For1_DualAction_FromRouter": "284914",
+  "test_stableSwap_ExactInput0For1_DualAction_FromUser": "311144",
+  "test_stableSwap_ExactInput0For1_MultiHop_FromRouter": "281742",
+  "test_stableSwap_ExactInput0For1_MultiHop_FromUser": "307972",
+  "test_stableSwap_ExactInput0For1_SamePath_FromRouter": "246522",
+  "test_stableSwap_ExactInput0For1_SamePath_FromUser": "292183",
+  "test_stableSwap_ExactOut0For1_MultiHop_FromUser": "354076",
+  "test_stableSwap_ExactOutput0For1_MultiHop_FromRouter": "327846",
+  "test_stableSwap_ExactOutput0For1_SamePath_FromRouter": "291659"
+}

snapshots/UniversalRouterTest.json

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

snapshots/V2BnbCake.json

@@ -1,4 +1,4 @@
 {
-  "test_v2Swap_exactInput0For1": "116435",
-  "test_v2Swap_exactOutput0For1": "117038"
+  "test_v2Swap_exactInput0For1": "116453",
+  "test_v2Swap_exactOutput0For1": "117044"
 }

snapshots/V2MockBnb.json

@@ -1,4 +1,4 @@
 {
-  "test_v2Swap_exactInput0For1": "100167",
-  "test_v2Swap_exactOutput0For1": "100791"
+  "test_v2Swap_exactInput0For1": "100185",
+  "test_v2Swap_exactOutput0For1": "100797"
 }

snapshots/V3BnbCake.json

@@ -1,7 +1,7 @@
 {
-  "test_v3Swap_ExactInput0For1": "151904",
-  "test_v3Swap_ExactInput0For1_ContractBalance": "154811",
-  "test_v3Swap_exactInput_MultiHop": "243354",
-  "test_v3Swap_exactOutput0For1": "150606",
-  "test_v3Swap_exactOutput_MultiHop": "254284"
+  "test_v3Swap_ExactInput0For1": "151922",
+  "test_v3Swap_ExactInput0For1_ContractBalance": "154835",
+  "test_v3Swap_exactInput_MultiHop": "243390",
+  "test_v3Swap_exactOutput0For1": "150624",
+  "test_v3Swap_exactOutput_MultiHop": "254314"
 }

snapshots/V3ToInfinityMigrationNativeTest.json

@@ -1,4 +1,4 @@
 {
-  "test_infiBinPositionmanager_BinAddLiquidity_Native": "543857",
-  "test_infiCLPositionmanager_Mint_Native": "533474"
+  "test_infiBinPositionmanager_BinAddLiquidity_Native": "544766",
+  "test_infiCLPositionmanager_Mint_Native": "533949"
 }

snapshots/V3ToInfinityMigrationTest.json

@@ -1,5 +1,5 @@
 {
-  "test_infiBinPositionmanager_BinAddLiquidity": "591530",
-  "test_infiCLPositionmanager_Mint": "581171",
+  "test_infiBinPositionmanager_BinAddLiquidity": "592472",
+  "test_infiCLPositionmanager_Mint": "581676",
   "test_v3PositionManager_burn": "286990"
 }

src/modules/pancakeswap/StableSwapRouter.sol

@@ -45,16 +45,27 @@ 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
+    function _stableSwap(address[] calldata path, uint256[] calldata flag, uint256 amountIn) private {
+        if (path.length - 1 != flag.length) revert StableInvalidPath();
+
+        uint256 balanceBefore;
+        for (uint256 i; i < flag.length; i++) {
+            (address input, address output) = (path[i], path[i + 1]);
+
+            bool isLastHop = i == flag.length - 1;
+            if (!isLastHop) {
+                balanceBefore = ERC20(path[i + 1]).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);
+
+            if (!isLastHop) {
+                // if not last swap, update amountIn for the next hop. this is done as swapContract do not return the output amount
+                amountIn = ERC20(path[i + 1]).balanceOf(address(this)) - balanceBefore;
             }
         }
     }
@@ -78,10 +89,14 @@ abstract contract StableSwapRouter is RouterImmutables, Permit2Payments, Ownable
             payOrPermit2Transfer(path[0], payer, address(this), amountIn);
         }
 
+        if (amountIn == ActionConstants.CONTRACT_BALANCE) {
+            amountIn = ERC20(path[0]).balanceOf(address(this));
+        }
+
         ERC20 tokenOut = ERC20(path[path.length - 1]);
         uint256 balanceBefore = tokenOut.balanceOf(address(this));
 
-        _stableSwap(path, flag);
+        _stableSwap(path, flag, amountIn);
 
         uint256 amountOut = tokenOut.balanceOf(address(this)) - balanceBefore;
         if (amountOut < amountOutMinimum) revert StableTooLittleReceived();
@@ -106,7 +121,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);