feat: Do not use V4Router for uniswap v4

Do our own implementation. In the end this is much cleaner than I expected and I was able to do a few improvements:
- we don't need to use actions. using function selectors and delegate call instead
- we don't need to convert into V4 Router types
- we don't need to check balances to get the amount out, we can just use the returned value

--- don't change below this line ---
ENG-4437 Took 3 hours 56 minutes


Took 3 minutes

Took 11 minutes

foundry/src/TychoRouter.sol

@@ -786,8 +786,8 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
         returns (bytes memory)
     {
         if (data.length < 24) revert TychoRouter__InvalidDataLength();
-        _handleCallback(data);
+        bytes memory result = _handleCallback(data);
-        return "";
+        return result;
     }
 
     function _balanceOf(address token, address owner)

foundry/src/executors/UniswapV4Executor.sol

@@ -2,28 +2,48 @@
 pragma solidity ^0.8.26;
 
 import "@interfaces/IExecutor.sol";
+import {ICallback} from "@interfaces/ICallback.sol";
+import {TokenTransfer} from "./TokenTransfer.sol";
 import {
     IERC20,
     SafeERC20
 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+
 import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
 import {
     Currency, CurrencyLibrary
 } from "@uniswap/v4-core/src/types/Currency.sol";
 import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
+import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+import {TickMath} from "@uniswap/v4-core/src/libraries/TickMath.sol";
 import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
-import {V4Router} from "@uniswap/v4-periphery/src/V4Router.sol";
-import {Actions} from "@uniswap/v4-periphery/src/libraries/Actions.sol";
-import {IV4Router} from "@uniswap/v4-periphery/src/interfaces/IV4Router.sol";
 import {PathKey} from "@uniswap/v4-periphery/src/libraries/PathKey.sol";
-import {ICallback} from "@interfaces/ICallback.sol";
+import {IUnlockCallback} from
-import {TokenTransfer} from "./TokenTransfer.sol";
+    "@uniswap/v4-core/src/interfaces/callback/IUnlockCallback.sol";
+import {SafeCast} from "@uniswap/v4-core/src/libraries/SafeCast.sol";
+import {TransientStateLibrary} from
+    "@uniswap/v4-core/src/libraries/TransientStateLibrary.sol";
 
 error UniswapV4Executor__InvalidDataLength();
+error UniswapV4Executor__NotPoolManager();
+error UniswapV4Executor__DeltaNotPositive(Currency currency);
+error UniswapV4Executor__DeltaNotNegative(Currency currency);
+error UniswapV4Executor__V4TooMuchRequested(
+    uint256 maxAmountInRequested, uint256 amountRequested
+);
 
-contract UniswapV4Executor is IExecutor, V4Router, ICallback, TokenTransfer {
+contract UniswapV4Executor is
+    IExecutor,
+    IUnlockCallback,
+    ICallback,
+    TokenTransfer
+{
     using SafeERC20 for IERC20;
     using CurrencyLibrary for Currency;
+    using SafeCast for *;
+    using TransientStateLibrary for IPoolManager;
+
+    IPoolManager public immutable poolManager;
 
     struct UniswapV4Pool {
         address intermediaryToken;
@@ -32,9 +52,20 @@ contract UniswapV4Executor is IExecutor, V4Router, ICallback, TokenTransfer {
     }
 
     constructor(IPoolManager _poolManager, address _permit2)
-        V4Router(_poolManager)
         TokenTransfer(_permit2)
-    {}
+    {
+        poolManager = _poolManager;
+    }
+
+    /**
+     * @dev Modifier to restrict access to only the pool manager.
+     */
+    modifier poolManagerOnly() virtual {
+        if (msg.sender != address(poolManager)) {
+            revert UniswapV4Executor__NotPoolManager();
+        }
+        _;
+    }
 
     function swap(uint256 amountIn, bytes calldata data)
         external
@@ -70,26 +101,14 @@ contract UniswapV4Executor is IExecutor, V4Router, ICallback, TokenTransfer {
                 tickSpacing: pools[0].tickSpacing,
                 hooks: IHooks(address(0))
             });
-            bytes memory actions = abi.encodePacked(
+            swapData = abi.encodeWithSelector(
-                uint8(Actions.SWAP_EXACT_IN_SINGLE),
+                this.swapExactInputSingle.selector,
-                uint8(Actions.SETTLE_ALL),
+                key,
-                uint8(Actions.TAKE)
+                zeroForOne,
+                amountIn,
+                receiver,
+                bytes("")
             );
-
-            bytes[] memory params = new bytes[](3);
-
-            params[0] = abi.encode(
-                IV4Router.ExactInputSingleParams({
-                    poolKey: key,
-                    zeroForOne: zeroForOne,
-                    amountIn: uint128(amountIn),
-                    amountOutMinimum: uint128(0),
-                    hookData: bytes("")
-                })
-            );
-            params[1] = abi.encode(tokenIn, amountIn); // currency to settle
-            params[2] = abi.encode(tokenOut, receiver, uint256(0)); // currency to take. 0 means to take the full amount
-            swapData = abi.encode(actions, params);
         } else {
             PathKey[] memory path = new PathKey[](pools.length);
             for (uint256 i = 0; i < pools.length; i++) {
@@ -102,51 +121,20 @@ contract UniswapV4Executor is IExecutor, V4Router, ICallback, TokenTransfer {
                 });
             }
 
-            bytes memory actions = abi.encodePacked(
-                uint8(Actions.SWAP_EXACT_IN),
-                uint8(Actions.SETTLE_ALL),
-                uint8(Actions.TAKE)
-            );
-
-            bytes[] memory params = new bytes[](3);
-
             Currency currencyIn = Currency.wrap(tokenIn);
-            params[0] = abi.encode(
+            swapData = abi.encodeWithSelector(
-                IV4Router.ExactInputParams({
+                this.swapExactInput.selector,
-                    currencyIn: currencyIn,
+                currencyIn,
-                    path: path,
+                path,
-                    amountIn: uint128(amountIn),
+                amountIn,
-                    amountOutMinimum: uint128(0)
+                receiver
-                })
             );
-            params[1] = abi.encode(currencyIn, amountIn);
-            params[2] =
-                abi.encode(Currency.wrap(tokenOut), receiver, uint256(0));
-            swapData = abi.encode(actions, params);
-        }
-        uint256 tokenOutBalanceBefore;
-
-        tokenOutBalanceBefore = tokenOut == address(0)
-            ? receiver.balance
-            : IERC20(tokenOut).balanceOf(receiver);
-
-        executeActions(swapData);
-
-        uint256 tokenOutBalanceAfter;
-
-        tokenOutBalanceAfter = tokenOut == address(0)
-            ? receiver.balance
-            : IERC20(tokenOut).balanceOf(receiver);
-
-        calculatedAmount = tokenOutBalanceAfter - tokenOutBalanceBefore;
-
-        return calculatedAmount;
         }
 
-    // necessary to convert bytes memory to bytes calldata
+        bytes memory result = poolManager.unlock(swapData);
-    function executeActions(bytes memory unlockData) public {
+        uint128 amountOut = abi.decode(result, (uint128));
-        // slither-disable-next-line unused-return
+
-        poolManager.unlock(unlockData);
+        return amountOut;
     }
 
     function _decodeData(bytes calldata data)
@@ -191,6 +179,9 @@ contract UniswapV4Executor is IExecutor, V4Router, ICallback, TokenTransfer {
         }
     }
 
+    /**
+     * @notice Handles the callback from the pool manager. This is used for callbacks from the router.
+     */
     function handleCallback(bytes calldata data)
         external
         returns (bytes memory)
@@ -199,15 +190,244 @@ contract UniswapV4Executor is IExecutor, V4Router, ICallback, TokenTransfer {
         return _unlockCallback(data);
     }
 
-    function verifyCallback(bytes calldata) public view onlyPoolManager {}
+    function verifyCallback(bytes calldata) public view poolManagerOnly {}
 
-    function _pay(Currency token, address, uint256 amount) internal override {
+    /**
+     * @notice Handles the unlock callback from the pool manager. This is used for swaps against the executor directly (bypassing the router).
+     */
+    function unlockCallback(bytes calldata data)
+        external
+        poolManagerOnly
+        returns (bytes memory)
+    {
+        return _unlockCallback(data);
+    }
+
+    /**
+     * @dev Internal function to handle the unlock callback.
+     * The executor address is needed to perform the call. If the router is being used, the executor address is in
+     * transient storage. If it is not, then address(this) should be used.
+     */
+    function _unlockCallback(bytes calldata data)
+        internal
+        returns (bytes memory)
+    {
+        address executor;
+        // slither-disable-next-line assembly
+        assembly {
+            executor := tload(0)
+        }
+
+        if (executor == address(0)) {
+            executor = address(this);
+        }
+        // slither-disable-next-line low-level-calls
+        (bool success, bytes memory returnData) = executor.delegatecall(data);
+        if (!success) {
+            revert(
+                string(
+                    returnData.length > 0
+                        ? returnData
+                        : abi.encodePacked("Uniswap v4 Callback failed")
+                )
+            );
+        }
+        return returnData;
+    }
+
+    /**
+     * @notice Performs an exact input single swap. It settles and takes the tokens after the swap.
+     * @param poolKey The key of the pool to swap in.
+     * @param zeroForOne Whether the swap is from token0 to token1 (true) or vice versa (false).
+     * @param amountIn The amount of tokens to swap in.
+     * @param receiver The address of the receiver.
+     * @param hookData Additional data for hook contracts.
+     */
+    function swapExactInputSingle(
+        PoolKey memory poolKey,
+        bool zeroForOne,
+        uint128 amountIn,
+        address receiver,
+        bytes calldata hookData
+    ) external returns (uint128) {
+        uint128 amountOut = _swap(
+            poolKey, zeroForOne, -int256(uint256(amountIn)), hookData
+        ).toUint128();
+
+        Currency currencyIn = zeroForOne ? poolKey.currency0 : poolKey.currency1;
+        uint256 amount = _getFullDebt(currencyIn);
+        if (amount > amountIn) {
+            revert UniswapV4Executor__V4TooMuchRequested(amountIn, amount);
+        }
+        _settle(currencyIn, address(this), amount);
+
+        Currency currencyOut =
+            zeroForOne ? poolKey.currency1 : poolKey.currency0;
+        _take(currencyOut, receiver, _mapTakeAmount(amountOut, currencyOut));
+        return amountOut;
+    }
+
+    /**
+     * @notice Performs an exact input swap along a path. It settles and takes the tokens after the swap.
+     * @param currencyIn The currency of the input token.
+     * @param path The path to swap along.
+     * @param amountIn The amount of tokens to swap in.
+     * @param receiver The address of the receiver.
+     */
+    function swapExactInput(
+        Currency currencyIn,
+        PathKey[] calldata path,
+        uint128 amountIn,
+        address receiver
+    ) external returns (uint128) {
+        uint128 amountOut = 0;
+        Currency swapCurrencyIn = currencyIn;
+        uint256 swapAmountIn = amountIn;
+        unchecked {
+            uint256 pathLength = path.length;
+            PathKey calldata pathKey;
+
+            for (uint256 i = 0; i < pathLength; i++) {
+                pathKey = path[i];
+                (PoolKey memory poolKey, bool zeroForOne) =
+                    pathKey.getPoolAndSwapDirection(swapCurrencyIn);
+                // The output delta will always be positive, except for when interacting with certain hook pools
+                amountOut = _swap(
+                    poolKey,
+                    zeroForOne,
+                    -int256(uint256(swapAmountIn)),
+                    pathKey.hookData
+                ).toUint128();
+
+                swapAmountIn = amountOut;
+                swapCurrencyIn = pathKey.intermediateCurrency;
+            }
+        }
+
+        uint256 amount = _getFullDebt(currencyIn);
+        if (amount > amountIn) {
+            revert UniswapV4Executor__V4TooMuchRequested(amountIn, amount);
+        }
+        _settle(currencyIn, address(this), amount);
+
+        _take(
+            swapCurrencyIn, // at the end of the loop this is actually currency out
+            receiver,
+            _mapTakeAmount(amountOut, swapCurrencyIn)
+        );
+        return amountOut;
+    }
+
+    function _swap(
+        PoolKey memory poolKey,
+        bool zeroForOne,
+        int256 amountSpecified,
+        bytes calldata hookData
+    ) private returns (int128 reciprocalAmount) {
+        unchecked {
+            // slither-disable-next-line calls-loop
+            BalanceDelta delta = poolManager.swap(
+                poolKey,
+                IPoolManager.SwapParams(
+                    zeroForOne,
+                    amountSpecified,
+                    zeroForOne
+                        ? TickMath.MIN_SQRT_PRICE + 1
+                        : TickMath.MAX_SQRT_PRICE - 1
+                ),
+                hookData
+            );
+
+            reciprocalAmount = (zeroForOne == amountSpecified < 0)
+                ? delta.amount1()
+                : delta.amount0();
+        }
+    }
+
+    /**
+     * @notice Obtains the full amount owed by this contract (negative delta).
+     * @param currency The currency to get the delta for.
+     * @return amount The amount owed by this contract.
+     */
+    function _getFullCredit(Currency currency)
+        internal
+        view
+        returns (uint256 amount)
+    {
+        int256 _amount = poolManager.currencyDelta(address(this), currency);
+        // If the amount is negative, it should be settled not taken.
+        if (_amount < 0) revert UniswapV4Executor__DeltaNotPositive(currency);
+        amount = uint256(_amount);
+    }
+
+    /// @notice Obtain the full amount owed by this contract (negative delta)
+    /// @param currency Currency to get the delta for
+    /// @return amount The amount owed by this contract as a uint256
+    function _getFullDebt(Currency currency)
+        internal
+        view
+        returns (uint256 amount)
+    {
+        int256 _amount = poolManager.currencyDelta(address(this), currency);
+        // If the amount is positive, it should be taken not settled.
+        if (_amount > 0) revert UniswapV4Executor__DeltaNotNegative(currency);
+        // Casting is safe due to limits on the total supply of a pool
+        amount = uint256(-_amount);
+    }
+
+    /**
+     * @notice Pays and settles a currency to the pool manager.
+     * @dev The implementing contract must ensure that the `payer` is a secure address.
+     * @param currency The currency to settle.
+     * @param payer The address of the payer.
+     * @param amount The amount to send.
+     * @dev Returns early if the amount is 0.
+     */
+    function _settle(Currency currency, address payer, uint256 amount)
+        internal
+    {
+        if (amount == 0) return;
+
+        poolManager.sync(currency);
+        if (currency.isAddressZero()) {
+            // slither-disable-next-line unused-return
+            poolManager.settle{value: amount}();
+        } else {
+            _pay(currency, payer, amount);
+            // slither-disable-next-line unused-return
+            poolManager.settle();
+        }
+    }
+
+    function _pay(Currency token, address, uint256 amount) internal {
         IERC20(Currency.unwrap(token)).safeTransfer(
             address(poolManager), amount
         );
     }
 
-    function msgSender() public view override returns (address) {
+    /**
-        return address(this);
+     * @notice Takes an amount of currency out of the pool manager.
+     * @param currency The currency to take.
+     * @param recipient The address to receive the currency.
+     * @param amount The amount to take.
+     * @dev Returns early if the amount is 0.
+     */
+    function _take(Currency currency, address recipient, uint256 amount)
+        internal
+    {
+        if (amount == 0) return;
+        poolManager.take(currency, recipient, amount);
+    }
+
+    function _mapTakeAmount(uint256 amount, Currency currency)
+        internal
+        view
+        returns (uint256)
+    {
+        if (amount == 0) {
+            return _getFullCredit(currency);
+        } else {
+            return amount;
+        }
     }
 }