reworked router to prepay input

Author: saucepoint

test/EulerSwapHook.swaps.t.sol

@@ -12,6 +12,7 @@ import {MinimalRouter} from "./utils/MinimalRouter.sol";
 import {Currency, CurrencyLibrary} from "@uniswap/v4-core/src/types/Currency.sol";
 import {StateLibrary} from "@uniswap/v4-core/src/libraries/StateLibrary.sol";
 import {TickMath} from "@uniswap/v4-core/src/libraries/TickMath.sol";
+import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
 
 contract EulerSwapHookTest is EulerSwapTestBase {
     using StateLibrary for IPoolManager;
@@ -38,14 +39,10 @@ contract EulerSwapHookTest is EulerSwapTestBase {
         assertFalse(eulerSwap.poolKey().currency1 == CurrencyLibrary.ADDRESS_ZERO);
         (uint160 sqrtPriceX96,,,) = poolManager.getSlot0(eulerSwap.poolKey().toId());
         assertNotEq(sqrtPriceX96, 0);
-
-        // Seed the poolManager with balance so that transient withdrawing before depositing succeeds
-        assetTST.mint(address(poolManager), 1000e18);
-        assetTST2.mint(address(poolManager), 1000e18);
     }
 
     function test_SwapExactIn() public {
-        uint256 amountIn = 1001e18;
+        uint256 amountIn = 1e18;
         uint256 amountOut =
             periphery.quoteExactInput(address(eulerSwap), address(assetTST), address(assetTST2), amountIn);
 
@@ -55,17 +52,20 @@ contract EulerSwapHookTest is EulerSwapTestBase {
         assetTST.approve(address(minimalRouter), amountIn);
 
         bool zeroForOne = address(assetTST) < address(assetTST2);
-        minimalRouter.swap(eulerSwap.poolKey(), zeroForOne, true, amountIn, "");
+        BalanceDelta result = minimalRouter.swap(eulerSwap.poolKey(), zeroForOne, amountIn, 0, "");
         vm.stopPrank();
 
         assertEq(assetTST.balanceOf(anyone), 0);
         assertEq(assetTST2.balanceOf(anyone), amountOut);
+
+        assertEq(zeroForOne ? uint256(-int256(result.amount0())) : uint256(-int256(result.amount1())), amountIn);
+        assertEq(zeroForOne ? uint256(int256(result.amount1())) : uint256(int256(result.amount0())), amountOut);
     }
 
     /// @dev swapping with an amount that exceeds PoolManager's ERC20 token balance will revert
     /// if the router does not pre-pay the input
     function test_swapExactIn_revertWithoutTokenLiquidity() public {
-        uint256 amountIn = 1001e18; // input amount exceeds
+        uint256 amountIn = 1e18; // input amount exceeds PoolManager balance
 
         assetTST.mint(anyone, amountIn);
 
@@ -87,13 +87,35 @@ contract EulerSwapHookTest is EulerSwapTestBase {
         assetTST.mint(anyone, amountIn);
 
         vm.startPrank(anyone);
-        assetTST.approve(address(swapRouter), amountIn);
+        assetTST.approve(address(minimalRouter), amountIn);
+
         bool zeroForOne = address(assetTST) < address(assetTST2);
-        _swap(eulerSwap.poolKey(), zeroForOne, false, amountOut);
+        BalanceDelta result = minimalRouter.swap(eulerSwap.poolKey(), zeroForOne, amountIn, amountOut, "");
         vm.stopPrank();
 
         assertEq(assetTST.balanceOf(anyone), 0);
         assertEq(assetTST2.balanceOf(anyone), amountOut);
+
+        assertEq(zeroForOne ? uint256(-int256(result.amount0())) : uint256(-int256(result.amount1())), amountIn);
+        assertEq(zeroForOne ? uint256(int256(result.amount1())) : uint256(int256(result.amount0())), amountOut);
+    }
+
+    /// @dev swapping with an amount that exceeds PoolManager's ERC20 token balance will revert
+    /// if the router does not pre-pay the input
+    function test_SwapExactOut_revertWithoutTokenLiquidity() public {
+        uint256 amountOut = 1e18;
+        uint256 amountIn =
+            periphery.quoteExactOutput(address(eulerSwap), address(assetTST), address(assetTST2), amountOut);
+
+        assetTST.mint(anyone, amountIn);
+
+        vm.startPrank(anyone);
+        assetTST.approve(address(swapRouter), amountIn);
+        bool zeroForOne = address(assetTST) < address(assetTST2);
+        PoolKey memory poolKey = eulerSwap.poolKey();
+        vm.expectRevert();
+        _swap(poolKey, zeroForOne, false, amountOut);
+        vm.stopPrank();
     }
 
     function _swap(PoolKey memory key, bool zeroForOne, bool exactInput, uint256 amount) internal {

test/utils/MinimalRouter.sol

@@ -3,68 +3,91 @@ pragma solidity ^0.8.24;
 
 import {CurrencyLibrary, Currency} from "@uniswap/v4-core/src/types/Currency.sol";
 import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
-import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+import {BalanceDelta, toBalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
 import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
 import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
 import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
 import {CurrencySettler} from "@uniswap/v4-core/test/utils/CurrencySettler.sol";
 import {SafeCallback} from "v4-periphery/src/base/SafeCallback.sol";
 import {TickMath} from "@uniswap/v4-core/src/libraries/TickMath.sol";
+import {TransientStateLibrary} from "@uniswap/v4-core/src/libraries/TransientStateLibrary.sol";
 
 contract MinimalRouter is SafeCallback {
+    using TransientStateLibrary for IPoolManager;
     using CurrencySettler for Currency;
 
     uint160 public constant MIN_PRICE_LIMIT = TickMath.MIN_SQRT_PRICE + 1;
     uint160 public constant MAX_PRICE_LIMIT = TickMath.MAX_SQRT_PRICE - 1;
 
     constructor(IPoolManager _manager) SafeCallback(_manager) {}
 
-    function swap(PoolKey memory key, bool zeroForOne, bool exactInput, uint256 amount, bytes memory hookData)
+    /// @dev an unsafe swap function that does not check for slippage
+    /// @param key The pool key
+    /// @param zeroForOne The direction of the swap
+    /// @param amountIn The amount of input token, should be provided (as an estimate) for exact output swaps
+    /// @param amountOut The amount of output token can be provided as 0, for exact input swaps
+    /// @param hookData The data to pass to the hook
+    function swap(PoolKey memory key, bool zeroForOne, uint256 amountIn, uint256 amountOut, bytes memory hookData)
         external
         payable
         returns (BalanceDelta delta)
     {
         delta = abi.decode(
-            poolManager.unlock(abi.encode(msg.sender, key, zeroForOne, exactInput, amount, hookData)), (BalanceDelta)
+            poolManager.unlock(abi.encode(msg.sender, key, zeroForOne, amountIn, amountOut, hookData)), (BalanceDelta)
         );
 
         uint256 ethBalance = address(this).balance;
         if (ethBalance > 0) CurrencyLibrary.ADDRESS_ZERO.transfer(msg.sender, ethBalance);
     }
 
     function _unlockCallback(bytes calldata data) internal override returns (bytes memory) {
-        (address sender, PoolKey memory key, bool zeroForOne, bool exactInput, uint256 amount, bytes memory hookData) =
-            abi.decode(data, (address, PoolKey, bool, bool, uint256, bytes));
+        (
+            address sender,
+            PoolKey memory key,
+            bool zeroForOne,
+            uint256 amountIn,
+            uint256 amountOut,
+            bytes memory hookData
+        ) = abi.decode(data, (address, PoolKey, bool, uint256, uint256, bytes));
 
-        // for exact input swaps, send the input first to avoid PoolManager token balance issues
-        if (exactInput) {
-            zeroForOne
-                ? key.currency0.settle(poolManager, sender, amount, false)
-                : key.currency1.settle(poolManager, sender, amount, false);
-        }
+        // send the input first to avoid PoolManager token balance issues
+        zeroForOne
+            ? key.currency0.settle(poolManager, sender, amountIn, false)
+            : key.currency1.settle(poolManager, sender, amountIn, false);
 
-        BalanceDelta delta = poolManager.swap(
+        poolManager.swap(
             key,
             IPoolManager.SwapParams({
                 zeroForOne: zeroForOne,
-                amountSpecified: exactInput ? -int256(amount) : int256(amount),
+                amountSpecified: amountOut != 0 ? int256(amountOut) : -int256(amountIn),
                 sqrtPriceLimitX96: zeroForOne ? MIN_PRICE_LIMIT : MAX_PRICE_LIMIT
             }),
             hookData
         );
 
-        if (!exactInput && delta.amount0() < 0) {
-            key.currency0.settle(poolManager, sender, uint256(int256(-delta.amount0())), false);
-        } else if (delta.amount0() > 0) {
-            key.currency0.take(poolManager, sender, uint256(int256(delta.amount0())), false);
+        // observe deltas
+        int256 delta0 = poolManager.currencyDelta(address(this), key.currency0);
+        int256 delta1 = poolManager.currencyDelta(address(this), key.currency1);
+
+        // primarily take the output token, and excess amounts for exact output swaps
+        if (delta0 < 0) {
+            key.currency0.settle(poolManager, sender, uint256(-delta0), false);
+        } else if (delta0 > 0) {
+            key.currency0.take(poolManager, sender, uint256(delta0), false);
         }
 
-        if (!exactInput && delta.amount1() < 0) {
-            key.currency1.settle(poolManager, sender, uint256(int256(-delta.amount1())), false);
-        } else if (delta.amount1() > 0) {
-            key.currency1.take(poolManager, sender, uint256(int256(delta.amount1())), false);
+        if (delta1 < 0) {
+            key.currency1.settle(poolManager, sender, uint256(-delta1), false);
+        } else if (delta1 > 0) {
+            key.currency1.take(poolManager, sender, uint256(delta1), false);
         }
 
-        return abi.encode(delta);
+        // account for prepaid input against the observed deltas
+        BalanceDelta returnDelta = toBalanceDelta(int128(delta0), int128(delta1))
+            + toBalanceDelta(
+                zeroForOne ? -int128(int256(amountIn)) : int128(0), zeroForOne ? int128(0) : -int128(int256(amountIn))
+            );
+
+        return abi.encode(returnDelta);
     }
 }