feat: update stable swap impl

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