Merge pull request #67 from propeller-heads/router/dc/ENG-4222-univ4-refactor-input

feat(univ4): Refactor execution input and handle all swap cases

foundry/src/TychoRouter.sol

@@ -449,8 +449,9 @@ contract TychoRouter is
         returns (bytes memory)
     {
         require(data.length >= 20, "Invalid data length");
-        bytes4 selector = bytes4(data[data.length - 24:data.length - 20]);
+        bytes4 selector = bytes4(data[data.length - 4:]);
-        address executor = address(uint160(bytes20(data[data.length - 20:])));
+        address executor =
+            address(uint160(bytes20(data[data.length - 24:data.length - 4])));
         bytes memory protocolData = data[:data.length - 24];
 
         if (!executors[executor]) {

foundry/src/executors/UniswapV4Executor.sol

@@ -18,121 +18,168 @@ import {IV4Router} from "@uniswap/v4-periphery/src/interfaces/IV4Router.sol";
 import {PathKey} from "@uniswap/v4-periphery/src/libraries/PathKey.sol";
 
 error UniswapV4Executor__InvalidDataLength();
-error UniswapV4Executor__SwapFailed();
 
 contract UniswapV4Executor is IExecutor, V4Router {
     using SafeERC20 for IERC20;
     using CurrencyLibrary for Currency;
 
+    struct UniswapV4Pool {
+        address intermediaryToken;
+        uint24 fee;
+        int24 tickSpacing;
+    }
+
     constructor(IPoolManager _poolManager) V4Router(_poolManager) {}
 
-    function swap(uint256, bytes calldata data)
+    function swap(uint256 amountIn, bytes calldata data)
         external
         payable
         returns (uint256 calculatedAmount)
     {
-        (address tokenIn, address tokenOut, bool isExactInput, uint256 amount) =
+        (
-            _decodeData(data);
+            address tokenIn,
+            address tokenOut,
+            uint256 amountOutMin,
+            bool zeroForOne,
+            address callbackExecutor,
+            bytes4 callbackSelector,
+            UniswapV4Executor.UniswapV4Pool[] memory pools
+        ) = _decodeData(data);
 
+        bytes memory swapData;
+        if (pools.length == 1) {
+            PoolKey memory key = PoolKey({
+                currency0: Currency.wrap(zeroForOne ? tokenIn : tokenOut),
+                currency1: Currency.wrap(zeroForOne ? tokenOut : tokenIn),
+                fee: pools[0].fee,
+                tickSpacing: pools[0].tickSpacing,
+                hooks: IHooks(address(0))
+            });
+            bytes memory actions = abi.encodePacked(
+                uint8(Actions.SWAP_EXACT_IN_SINGLE),
+                uint8(Actions.SETTLE_ALL),
+                uint8(Actions.TAKE_ALL)
+            );
+
+            bytes[] memory params = new bytes[](3);
+
+            params[0] = abi.encode(
+                IV4Router.ExactInputSingleParams({
+                    poolKey: key,
+                    zeroForOne: zeroForOne,
+                    amountIn: uint128(amountIn),
+                    amountOutMinimum: uint128(amountOutMin),
+                    hookData: bytes("")
+                })
+            );
+            params[1] = abi.encode(key.currency0, amountIn);
+            params[2] = abi.encode(key.currency1, amountOutMin);
+            swapData = abi.encode(actions, params);
+        } else {
+            PathKey[] memory path = new PathKey[](pools.length);
+            for (uint256 i = 0; i < pools.length; i++) {
+                path[i] = PathKey({
+                    intermediateCurrency: Currency.wrap(pools[i].intermediaryToken),
+                    fee: pools[i].fee,
+                    tickSpacing: pools[i].tickSpacing,
+                    hooks: IHooks(address(0)),
+                    hookData: bytes("")
+                });
+            }
+
+            bytes memory actions = abi.encodePacked(
+                uint8(Actions.SWAP_EXACT_IN),
+                uint8(Actions.SETTLE_ALL),
+                uint8(Actions.TAKE_ALL)
+            );
+
+            bytes[] memory params = new bytes[](3);
+
+            Currency currencyIn = Currency.wrap(tokenIn);
+            params[0] = abi.encode(
+                IV4Router.ExactInputParams({
+                    currencyIn: currencyIn,
+                    path: path,
+                    amountIn: uint128(amountIn),
+                    amountOutMinimum: uint128(amountOutMin)
+                })
+            );
+            params[1] = abi.encode(currencyIn, amountIn);
+            params[2] = abi.encode(Currency.wrap(tokenOut), amountOutMin);
+            swapData = abi.encode(actions, params);
+        }
+        bytes memory fullData =
+            abi.encodePacked(swapData, callbackExecutor, callbackSelector);
         uint256 tokenOutBalanceBefore;
-        uint256 tokenInBalanceBefore;
 
         tokenOutBalanceBefore = tokenOut == address(0)
             ? address(this).balance
             : IERC20(tokenOut).balanceOf(address(this));
 
-        tokenInBalanceBefore = tokenIn == address(0)
+        executeActions(fullData);
-            ? address(this).balance
-            : IERC20(tokenIn).balanceOf(address(this));
-
-        _executeActions(data);
 
         uint256 tokenOutBalanceAfter;
-        uint256 tokenInBalanceAfter;
 
         tokenOutBalanceAfter = tokenOut == address(0)
             ? address(this).balance
             : IERC20(tokenOut).balanceOf(address(this));
 
-        tokenInBalanceAfter = tokenIn == address(0)
+        calculatedAmount = tokenOutBalanceAfter - tokenOutBalanceBefore;
-            ? address(this).balance
-            : IERC20(tokenIn).balanceOf(address(this));
-
-        if (isExactInput) {
-            calculatedAmount = tokenOutBalanceAfter - tokenOutBalanceBefore;
-        } else {
-            calculatedAmount = tokenInBalanceBefore - tokenInBalanceAfter;
-        }
 
         return calculatedAmount;
     }
 
+    // necessary to convert bytes memory to bytes calldata
+    function executeActions(bytes memory unlockData) public {
+        // slither-disable-next-line unused-return
+        poolManager.unlock(unlockData);
+    }
+
     function _decodeData(bytes calldata data)
         internal
         pure
         returns (
             address tokenIn,
             address tokenOut,
-            bool isExactInput,
+            uint256 amountOutMin,
-            uint256 amount
+            bool zeroForOne,
+            address callbackExecutor,
+            bytes4 callbackSelector,
+            UniswapV4Pool[] memory pools
         )
     {
-        (bytes memory actions, bytes[] memory params) =
+        if (data.length < 123) {
-            abi.decode(data, (bytes, bytes[]));
+            revert UniswapV4Executor__InvalidDataLength();
+        }
 
-        // First byte of actions determines the swap type
+        tokenIn = address(bytes20(data[0:20]));
-        uint8 action = uint8(bytes1(actions[0]));
+        tokenOut = address(bytes20(data[20:40]));
+        amountOutMin = uint256(bytes32(data[40:72]));
+        zeroForOne = (data[72] != 0);
+        callbackExecutor = address(bytes20(data[73:93]));
+        callbackSelector = bytes4(data[93:97]);
 
-        if (action == uint8(Actions.SWAP_EXACT_IN_SINGLE)) {
+        uint256 poolsLength = (data.length - 97) / 26; // 26 bytes per pool object
-            IV4Router.ExactInputSingleParams memory swapParams =
+        pools = new UniswapV4Pool[](poolsLength);
-                abi.decode(params[0], (IV4Router.ExactInputSingleParams));
+        bytes memory poolsData = data[97:];
+        uint256 offset = 0;
+        for (uint256 i = 0; i < poolsLength; i++) {
+            address intermediaryToken;
+            uint24 fee;
+            int24 tickSpacing;
 
-            tokenIn = swapParams.zeroForOne
+            // slither-disable-next-line assembly
-                ? address(uint160(swapParams.poolKey.currency0.toId()))
+            assembly {
-                : address(uint160(swapParams.poolKey.currency1.toId()));
+                intermediaryToken := mload(add(poolsData, add(offset, 20)))
-            tokenOut = swapParams.zeroForOne
+                fee := shr(232, mload(add(poolsData, add(offset, 52))))
-                ? address(uint160(swapParams.poolKey.currency1.toId()))
+                tickSpacing := shr(232, mload(add(poolsData, add(offset, 55))))
-                : address(uint160(swapParams.poolKey.currency0.toId()));
+            }
-            isExactInput = true;
+            pools[i] = UniswapV4Pool(intermediaryToken, fee, tickSpacing);
-            amount = swapParams.amountIn;
+            offset += 26;
-        } else if (action == uint8(Actions.SWAP_EXACT_OUT_SINGLE)) {
-            IV4Router.ExactOutputSingleParams memory swapParams =
-                abi.decode(params[0], (IV4Router.ExactOutputSingleParams));
-
-            tokenIn = swapParams.zeroForOne
-                ? address(uint160(swapParams.poolKey.currency0.toId()))
-                : address(uint160(swapParams.poolKey.currency1.toId()));
-            tokenOut = swapParams.zeroForOne
-                ? address(uint160(swapParams.poolKey.currency1.toId()))
-                : address(uint160(swapParams.poolKey.currency0.toId()));
-            isExactInput = false;
-            amount = swapParams.amountOut;
-        } else if (action == uint8(Actions.SWAP_EXACT_IN)) {
-            IV4Router.ExactInputParams memory swapParams =
-                abi.decode(params[0], (IV4Router.ExactInputParams));
-
-            tokenIn = address(uint160(swapParams.currencyIn.toId()));
-            PathKey memory lastPath =
-                swapParams.path[swapParams.path.length - 1];
-            tokenOut = address(uint160(lastPath.intermediateCurrency.toId()));
-            isExactInput = true;
-            amount = swapParams.amountIn;
-        } else if (action == uint8(Actions.SWAP_EXACT_OUT)) {
-            IV4Router.ExactOutputParams memory swapParams =
-                abi.decode(params[0], (IV4Router.ExactOutputParams));
-
-            PathKey memory firstPath = swapParams.path[0];
-            tokenIn = address(uint160(firstPath.intermediateCurrency.toId()));
-            tokenOut = address(uint160(swapParams.currencyOut.toId()));
-            isExactInput = false;
-            amount = swapParams.amountOut;
         }
     }
 
-    function _pay(Currency token, address payer, uint256 amount)
+    function _pay(Currency token, address, uint256 amount) internal override {
-        internal
-        override
-    {
         IERC20(Currency.unwrap(token)).safeTransfer(
             address(poolManager), amount
         );

foundry/test/LibPrefixLengthEncodedByteArray.t.sol

@@ -51,14 +51,16 @@ contract LibPrefixLengthEncodedByteArrayTest is Test {
         assertEq(this.size(multiple), 3);
     }
 
-    function testFailInvalidLength() public view {
+    function test_RevertIf_InvalidLength() public {
         // Length prefix larger than remaining data
+        vm.expectRevert();
         bytes memory invalid = hex"0004414243";
         this.next(invalid);
     }
 
-    function testFailIncompletePrefix() public view {
+    function test_RevertIf_IncompletePrefix() public {
         // Only 1 byte instead of 2 bytes prefix
+        vm.expectRevert();
         bytes memory invalid = hex"01";
         this.next(invalid);
     }

foundry/test/TychoRouter.t.sol

@@ -857,15 +857,22 @@ contract TychoRouterTest is TychoRouterTestSetup {
         uint256 amountIn = 100 ether;
         deal(USDE_ADDR, tychoRouterAddr, amountIn);
 
-        bytes memory protocolData = UniswapV4Utils.encodeExactInputSingle(
+        UniswapV4Executor.UniswapV4Pool[] memory pools =
-            USDE_ADDR, USDT_ADDR, 100, true, 1, uint128(amountIn)
+            new UniswapV4Executor.UniswapV4Pool[](1);
-        );
+        pools[0] = UniswapV4Executor.UniswapV4Pool({
+            intermediaryToken: USDT_ADDR,
+            fee: uint24(100),
+            tickSpacing: int24(1)
+        });
 
-        // add executor and selector for callback
+        bytes memory protocolData = UniswapV4Utils.encodeExactInput(
-        bytes memory protocolDataWithCallBack = abi.encodePacked(
+            USDE_ADDR,
-            protocolData,
+            USDT_ADDR,
+            uint256(1),
+            true,
+            address(usv4Executor),
             SafeCallback.unlockCallback.selector,
-            address(usv4Executor)
+            pools
         );
 
         bytes memory swap = encodeSwap(
@@ -874,7 +881,7 @@ contract TychoRouterTest is TychoRouterTestSetup {
             uint24(0),
             address(usv4Executor),
             bytes4(0),
-            protocolDataWithCallBack
+            protocolData
         );
 
         bytes[] memory swaps = new bytes[](1);
@@ -882,6 +889,52 @@ contract TychoRouterTest is TychoRouterTestSetup {
 
         tychoRouter.exposedSwap(amountIn, 2, pleEncode(swaps));
 
-        assertTrue(IERC20(USDT_ADDR).balanceOf(tychoRouterAddr) == 99943852);
+        assertEq(IERC20(USDT_ADDR).balanceOf(tychoRouterAddr), 99943852);
+    }
+
+    function testSwapMultipleUSV4Callback() public {
+        // This test has two uniswap v4 hops that will be executed inside of the V4 pool manager
+        // USDE -> USDT -> WBTC
+        uint256 amountIn = 100 ether;
+        deal(USDE_ADDR, tychoRouterAddr, amountIn);
+
+        UniswapV4Executor.UniswapV4Pool[] memory pools =
+            new UniswapV4Executor.UniswapV4Pool[](2);
+        pools[0] = UniswapV4Executor.UniswapV4Pool({
+            intermediaryToken: USDT_ADDR,
+            fee: uint24(100),
+            tickSpacing: int24(1)
+        });
+        pools[1] = UniswapV4Executor.UniswapV4Pool({
+            intermediaryToken: WBTC_ADDR,
+            fee: uint24(3000),
+            tickSpacing: int24(60)
+        });
+
+        bytes memory protocolData = UniswapV4Utils.encodeExactInput(
+            USDE_ADDR,
+            WBTC_ADDR,
+            uint256(1),
+            true,
+            address(usv4Executor),
+            SafeCallback.unlockCallback.selector,
+            pools
+        );
+
+        bytes memory swap = encodeSwap(
+            uint8(0),
+            uint8(1),
+            uint24(0),
+            address(usv4Executor),
+            bytes4(0),
+            protocolData
+        );
+
+        bytes[] memory swaps = new bytes[](1);
+        swaps[0] = swap;
+
+        tychoRouter.exposedSwap(amountIn, 2, pleEncode(swaps));
+
+        assertEq(IERC20(WBTC_ADDR).balanceOf(tychoRouterAddr), 102718);
     }
 }

foundry/test/executors/UniswapV4Executor.t.sol

@@ -1,11 +1,12 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.26;
 
+import "../../src/executors/UniswapV4Executor.sol";
 import "./UniswapV4Utils.sol";
 import "@src/executors/UniswapV4Executor.sol";
 import {Constants} from "../Constants.sol";
 import {Test} from "../../lib/forge-std/src/Test.sol";
-import {console} from "forge-std/console.sol";
+import {SafeCallback} from "@uniswap/v4-periphery/src/base/SafeCallback.sol";
 
 contract UniswapV4ExecutorExposed is UniswapV4Executor {
     constructor(IPoolManager _poolManager) UniswapV4Executor(_poolManager) {}
@@ -16,8 +17,11 @@ contract UniswapV4ExecutorExposed is UniswapV4Executor {
         returns (
             address tokenIn,
             address tokenOut,
-            bool isExactInput,
+            uint256 amountOutMin,
-            uint256 amount
+            bool zeroForOne,
+            address callbackExecutor,
+            bytes4 callbackSelector,
+            UniswapV4Pool[] memory pools
         )
     {
         return _decodeData(data);
@@ -40,31 +44,84 @@ contract UniswapV4ExecutorTest is Test, Constants {
     }
 
     function testDecodeParams() public view {
-        uint24 expectedPoolFee = 500;
+        uint256 minAmountOut = 100;
-        uint128 expectedAmount = 100;
+        bool zeroForOne = true;
+        uint24 pool1Fee = 500;
+        int24 tickSpacing1 = 60;
+        uint24 pool2Fee = 1000;
+        int24 tickSpacing2 = -10;
 
-        bytes memory data = UniswapV4Utils.encodeExactInputSingle(
+        UniswapV4Executor.UniswapV4Pool[] memory pools =
-            USDE_ADDR, USDT_ADDR, expectedPoolFee, false, 1, expectedAmount
+            new UniswapV4Executor.UniswapV4Pool[](2);
+        pools[0] = UniswapV4Executor.UniswapV4Pool({
+            intermediaryToken: USDT_ADDR,
+            fee: pool1Fee,
+            tickSpacing: tickSpacing1
+        });
+        pools[1] = UniswapV4Executor.UniswapV4Pool({
+            intermediaryToken: USDE_ADDR,
+            fee: pool2Fee,
+            tickSpacing: tickSpacing2
+        });
+
+        bytes memory data = UniswapV4Utils.encodeExactInput(
+            USDE_ADDR,
+            USDT_ADDR,
+            minAmountOut,
+            zeroForOne,
+            address(uniswapV4Exposed),
+            SafeCallback.unlockCallback.selector,
+            pools
         );
 
-        (address tokenIn, address tokenOut, bool isExactInput, uint256 amount) =
+        (
-            uniswapV4Exposed.decodeData(data);
+            address tokenIn,
+            address tokenOut,
+            uint256 amountOutMin,
+            bool zeroForOneDecoded,
+            address callbackExecutor,
+            bytes4 callbackSelector,
+            UniswapV4Executor.UniswapV4Pool[] memory decodedPools
+        ) = uniswapV4Exposed.decodeData(data);
 
         assertEq(tokenIn, USDE_ADDR);
         assertEq(tokenOut, USDT_ADDR);
-        assertTrue(isExactInput);
+        assertEq(amountOutMin, minAmountOut);
-        assertEq(amount, expectedAmount);
+        assertEq(zeroForOneDecoded, zeroForOne);
+        assertEq(callbackExecutor, address(uniswapV4Exposed));
+        assertEq(callbackSelector, SafeCallback.unlockCallback.selector);
+        assertEq(decodedPools.length, 2);
+        assertEq(decodedPools[0].intermediaryToken, USDT_ADDR);
+        assertEq(decodedPools[0].fee, pool1Fee);
+        assertEq(decodedPools[0].tickSpacing, tickSpacing1);
+        assertEq(decodedPools[1].intermediaryToken, USDE_ADDR);
+        assertEq(decodedPools[1].fee, pool2Fee);
+        assertEq(decodedPools[1].tickSpacing, tickSpacing2);
     }
 
-    function testSwap() public {
+    function testSingleSwap() public {
         uint256 amountIn = 100 ether;
         deal(USDE_ADDR, address(uniswapV4Exposed), amountIn);
         uint256 usdeBalanceBeforePool = USDE.balanceOf(poolManager);
         uint256 usdeBalanceBeforeSwapExecutor =
             USDE.balanceOf(address(uniswapV4Exposed));
 
-        bytes memory data = UniswapV4Utils.encodeExactInputSingle(
+        UniswapV4Executor.UniswapV4Pool[] memory pools =
-            USDE_ADDR, USDT_ADDR, 100, true, 1, uint128(amountIn)
+            new UniswapV4Executor.UniswapV4Pool[](1);
+        pools[0] = UniswapV4Executor.UniswapV4Pool({
+            intermediaryToken: USDT_ADDR,
+            fee: uint24(100),
+            tickSpacing: int24(1)
+        });
+
+        bytes memory data = UniswapV4Utils.encodeExactInput(
+            USDE_ADDR,
+            USDT_ADDR,
+            uint256(1),
+            true,
+            address(uniswapV4Exposed),
+            SafeCallback.unlockCallback.selector,
+            pools
         );
 
         uint256 amountOut = uniswapV4Exposed.swap(amountIn, data);
@@ -75,4 +132,46 @@ contract UniswapV4ExecutorTest is Test, Constants {
         );
         assertTrue(USDT.balanceOf(address(uniswapV4Exposed)) == amountOut);
     }
+
+    function testMultipleSwap() public {
+        // USDE -> USDT -> WBTC
+        uint256 amountIn = 100 ether;
+        deal(USDE_ADDR, address(uniswapV4Exposed), amountIn);
+        uint256 usdeBalanceBeforePool = USDE.balanceOf(poolManager);
+        uint256 usdeBalanceBeforeSwapExecutor =
+            USDE.balanceOf(address(uniswapV4Exposed));
+
+        UniswapV4Executor.UniswapV4Pool[] memory pools =
+            new UniswapV4Executor.UniswapV4Pool[](2);
+        pools[0] = UniswapV4Executor.UniswapV4Pool({
+            intermediaryToken: USDT_ADDR,
+            fee: uint24(100),
+            tickSpacing: int24(1)
+        });
+        pools[1] = UniswapV4Executor.UniswapV4Pool({
+            intermediaryToken: WBTC_ADDR,
+            fee: uint24(3000),
+            tickSpacing: int24(60)
+        });
+
+        bytes memory data = UniswapV4Utils.encodeExactInput(
+            USDE_ADDR,
+            WBTC_ADDR,
+            uint256(1),
+            true,
+            address(uniswapV4Exposed),
+            SafeCallback.unlockCallback.selector,
+            pools
+        );
+
+        uint256 amountOut = uniswapV4Exposed.swap(amountIn, data);
+        assertEq(USDE.balanceOf(poolManager), usdeBalanceBeforePool + amountIn);
+        assertEq(
+            USDE.balanceOf(address(uniswapV4Exposed)),
+            usdeBalanceBeforeSwapExecutor - amountIn
+        );
+        assertTrue(
+            IERC20(WBTC_ADDR).balanceOf(address(uniswapV4Exposed)) == amountOut
+        );
+    }
 }

foundry/test/executors/UniswapV4Utils.sol

@@ -4,43 +4,34 @@ pragma solidity ^0.8.26;
 import "@src/executors/UniswapV4Executor.sol";
 
 library UniswapV4Utils {
-    function encodeExactInputSingle(
+    function encodeExactInput(
         address tokenIn,
         address tokenOut,
-        uint24 fee,
+        uint256 amountOutMin,
         bool zeroForOne,
-        uint24 tickSpacing,
+        address callbackExecutor,
-        uint128 amountIn
+        bytes4 callbackSelector,
+        UniswapV4Executor.UniswapV4Pool[] memory pools
     ) public pure returns (bytes memory) {
-        PoolKey memory key = PoolKey({
+        bytes memory encodedPools;
-            currency0: Currency.wrap(zeroForOne ? tokenIn : tokenOut),
-            currency1: Currency.wrap(zeroForOne ? tokenOut : tokenIn),
-            fee: fee,
-            tickSpacing: int24(tickSpacing),
-            hooks: IHooks(address(0))
-        });
 
-        bytes memory actions = abi.encodePacked(
+        for (uint256 i = 0; i < pools.length; i++) {
-            uint8(Actions.SWAP_EXACT_IN_SINGLE),
+            encodedPools = abi.encodePacked(
-            uint8(Actions.SETTLE_ALL),
+                encodedPools,
-            uint8(Actions.TAKE_ALL)
+                pools[i].intermediaryToken,
+                bytes3(pools[i].fee),
+                pools[i].tickSpacing
+            );
+        }
+
+        return abi.encodePacked(
+            tokenIn,
+            tokenOut,
+            amountOutMin,
+            zeroForOne,
+            callbackExecutor,
+            bytes4(callbackSelector),
+            encodedPools
         );
-
-        bytes[] memory params = new bytes[](3);
-
-        params[0] = abi.encode(
-            IV4Router.ExactInputSingleParams({
-                poolKey: key,
-                zeroForOne: zeroForOne,
-                amountIn: amountIn,
-                amountOutMinimum: 0,
-                hookData: bytes("")
-            })
-        );
-
-        params[1] = abi.encode(key.currency0, amountIn);
-        params[2] = abi.encode(key.currency1, 0);
-
-        return abi.encode(actions, params);
     }
 }