exactOut implemented w awkward loops/int conversions

single hops passing on exactOut
Uniswap · Sep 5, 2023 · 679682d · 679682d
1 parent 5823d2f
commit 679682d
Show file tree

Hide file tree

Showing 10 changed files with 175 additions and 21 deletions.
diff --git a/.forge-snapshots/RouterBytecode.snap b/.forge-snapshots/RouterBytecode.snap
@@ -1 +1 @@
-4363
+4855
diff --git a/.forge-snapshots/RouterExactIn1Hop.snap b/.forge-snapshots/RouterExactIn1Hop.snap
@@ -1 +1 @@
-193825
+194162
diff --git a/.forge-snapshots/RouterExactIn2Hops.snap b/.forge-snapshots/RouterExactIn2Hops.snap
@@ -1 +1 @@
-271071
+271461
diff --git a/.forge-snapshots/RouterExactIn3Hops.snap b/.forge-snapshots/RouterExactIn3Hops.snap
@@ -1 +1 @@
-348334
+348765
diff --git a/.forge-snapshots/RouterExactInSingle.snap b/.forge-snapshots/RouterExactInSingle.snap
@@ -1 +1 @@
-192093
+192195
diff --git a/.forge-snapshots/RouterExactOut1Hop.snap b/.forge-snapshots/RouterExactOut1Hop.snap
@@ -0,0 +1 @@
+193514
diff --git a/.forge-snapshots/RouterExactOut2Hops.snap b/.forge-snapshots/RouterExactOut2Hops.snap
@@ -0,0 +1 @@
+271891
diff --git a/.forge-snapshots/RouterExactOut3Hops.snap b/.forge-snapshots/RouterExactOut3Hops.snap
@@ -0,0 +1 @@
+350296
diff --git a/contracts/Routing.sol b/contracts/Routing.sol
@@ -19,6 +19,7 @@ abstract contract Routing {
     error NotPoolManager();
     error InvalidSwapType();
     error TooLittleReceived();
+    error TooMuchRequested();
 
     struct SwapInfo {
         SwapType swapType;
@@ -27,7 +28,7 @@ abstract contract Routing {
     }
 
     struct PathKey {
-        Currency currencyOut;
+        Currency tradeCurrency;
         uint24 fee;
         int24 tickSpacing;
         IHooks hooks;
@@ -50,6 +51,14 @@ abstract contract Routing {
         uint128 amountOutMinimum;
     }
 
+    struct ExactOutputParams {
+        Currency currencyOut;
+        PathKey[] path;
+        address recipient;
+        uint128 amountOut;
+        uint128 amountInMaximum;
+    }
+
     enum SwapType {
         ExactInput,
         ExactInputSingle,
@@ -78,6 +87,8 @@ abstract contract Routing {
             _swapExactInput(abi.decode(swapInfo.params, (ExactInputParams)), swapInfo.msgSender);
         } else if (swapInfo.swapType == SwapType.ExactInputSingle) {
             _swapExactInputSingle(abi.decode(swapInfo.params, (ExactInputSingleParams)), swapInfo.msgSender);
+        } else if (swapInfo.swapType == SwapType.ExactOutput) {
+            _swapExactOutput(abi.decode(swapInfo.params, (ExactOutputParams)), swapInfo.msgSender);
         } else {
             revert InvalidSwapType();
         }
@@ -86,45 +97,66 @@ abstract contract Routing {
     }
 
     function _swapExactInputSingle(ExactInputSingleParams memory params, address msgSender) private {
-        _swapExactInputPrivate(
-            params.poolKey, params.zeroForOne, params.amountIn, params.sqrtPriceLimitX96, msgSender, true, true
+        _swapExactPrivate(
+            params.poolKey, params.zeroForOne, int256(int128(params.amountIn)), params.sqrtPriceLimitX96, msgSender, true, true
         );
     }
 
     function _swapExactInput(ExactInputParams memory params, address msgSender) private {
         for (uint256 i = 0; i < params.path.length; i++) {
             (PoolKey memory poolKey, bool zeroForOne) = _getPoolAndSwapDirection(params.path[i], params.currencyIn);
-            uint128 amountOut = _swapExactInputPrivate(
+            uint128 amountOut = uint128(-_swapExactPrivate(
                 poolKey,
                 zeroForOne,
-                params.amountIn,
+                int256(int128(params.amountIn)),
                 0,
                 msgSender,
                 i == 0,
                 i == params.path.length - 1
-            );
+            ));
 
             params.amountIn = amountOut;
-            params.currencyIn = params.path[i].currencyOut;
+            params.currencyIn = params.path[i].tradeCurrency;
         }
 
         if (params.amountIn < params.amountOutMinimum) revert TooLittleReceived();
     }
 
-    function _swapExactInputPrivate(
+    function _swapExactOutput(ExactOutputParams memory params, address msgSender) private {
+        for (uint256 i = params.path.length; i > 0; i--) {
+            (PoolKey memory poolKey, bool oneForZero) = _getPoolAndSwapDirection(params.path[i - 1], params.currencyOut);
+            uint128 amountIn = uint128(-_swapExactPrivate(
+                poolKey,
+                !oneForZero,
+                -int256(int128(params.amountOut)),
+                0,
+                msgSender,
+                i == 1,
+                i == params.path.length
+            ));
+
+            // console.log(amountIn);
+
+            params.amountOut = amountIn;
+            params.currencyOut = params.path[i - 1].tradeCurrency;
+        }
+        if (params.amountOut > params.amountInMaximum) revert TooMuchRequested();
+    }
+
+    function _swapExactPrivate(
         PoolKey memory poolKey,
         bool zeroForOne,
-        uint128 amountIn,
+        int256 amountSpecified,
         uint160 sqrtPriceLimitX96,
         address msgSender,
         bool settle,
         bool take
-    ) private returns (uint128 amountOut) {
+    ) private returns (int128 reciprocalAmount) {
         BalanceDelta delta = poolManager.swap(
             poolKey,
             IPoolManager.SwapParams(
                 zeroForOne,
-                int256(int128(amountIn)),
+                amountSpecified,
                 sqrtPriceLimitX96 == 0
                     ? (zeroForOne ? TickMath.MIN_SQRT_RATIO + 1 : TickMath.MAX_SQRT_RATIO - 1)
                     : sqrtPriceLimitX96
@@ -133,13 +165,13 @@ abstract contract Routing {
         );
 
         if (zeroForOne) {
-            amountOut = uint128(-delta.amount1());
+            reciprocalAmount = amountSpecified > 0 ? delta.amount1() : -delta.amount0();
             if (settle) _payAndSettle(poolKey.currency0, msgSender, delta.amount0());
-            if (take) poolManager.take(poolKey.currency1, msgSender, uint256(amountOut));
+            if (take) poolManager.take(poolKey.currency1, msgSender,  uint128(-delta.amount1()));
         } else {
-            amountOut = uint128(-delta.amount0());
+            reciprocalAmount = amountSpecified > 0 ? delta.amount0() : -delta.amount1();
             if (settle) _payAndSettle(poolKey.currency1, msgSender, delta.amount1());
-            if (take) poolManager.take(poolKey.currency0, msgSender, uint256(amountOut));
+            if (take) poolManager.take(poolKey.currency0, msgSender,  uint128(-delta.amount0()));
         }
     }
 
@@ -149,7 +181,7 @@ abstract contract Routing {
         returns (PoolKey memory poolKey, bool zeroForOne)
     {
         (Currency currency0, Currency currency1) =
-            currencyIn < params.currencyOut ? (currencyIn, params.currencyOut) : (params.currencyOut, currencyIn);
+            currencyIn < params.tradeCurrency ? (currencyIn, params.tradeCurrency) : (params.tradeCurrency, currencyIn);
 
         zeroForOne = currencyIn == currency0;
         poolKey = PoolKey(currency0, currency1, params.fee, params.tickSpacing, params.hooks);

diff --git a/test/Routing.t.sol b/test/Routing.t.sol
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.19;
 
+import 'forge-std/console.sol';
 import {Test} from "forge-std/Test.sol";
 import {Deployers} from "@uniswap/v4-core/test/foundry-tests/utils/Deployers.sol";
 import {GasSnapshot} from "forge-gas-snapshot/GasSnapshot.sol";
@@ -198,6 +199,107 @@ contract RoutingTest is Test, Deployers, GasSnapshot {
         assertEq(token3.balanceOf(address(router)), 0);
     }
 
+    function testRouter_swapExactOut_1Hop_zeroForOne() public {
+        uint256 amountOut = 1 ether;
+        uint256 expectedAmountIn = 1008049273448486163;
+
+        tokenPath.push(token0);
+        tokenPath.push(token1);
+        Routing.ExactOutputParams memory params = getExactOutputParams(tokenPath, amountOut);
+
+        uint256 prevBalance0 = token0.balanceOf(address(this));
+        uint256 prevBalance1 = token1.balanceOf(address(this));
+
+        snapStart("RouterExactOut1Hop");
+        router.swap(Routing.SwapType.ExactOutput, abi.encode(params));
+        snapEnd();
+
+        uint256 newBalance0 = token0.balanceOf(address(this));
+        uint256 newBalance1 = token1.balanceOf(address(this));
+
+        assertEq(prevBalance0 - newBalance0, expectedAmountIn);
+        assertEq(newBalance1 - prevBalance1, amountOut);
+    }
+
+    function testRouter_swapExactOut_1Hop_oneForZero() public {
+        uint256 amountOut = 1 ether;
+        uint256 expectedAmountIn = 1008049273448486163;
+
+        tokenPath.push(token1);
+        tokenPath.push(token0);
+        Routing.ExactOutputParams memory params = getExactOutputParams(tokenPath, amountOut);
+
+        uint256 prevBalance0 = token0.balanceOf(address(this));
+        uint256 prevBalance1 = token1.balanceOf(address(this));
+
+        snapStart("RouterExactOut1Hop");
+        router.swap(Routing.SwapType.ExactOutput, abi.encode(params));
+        snapEnd();
+
+        uint256 newBalance0 = token0.balanceOf(address(this));
+        uint256 newBalance1 = token1.balanceOf(address(this));
+
+        assertEq(prevBalance1 - newBalance1, expectedAmountIn);
+        assertEq(newBalance0 - prevBalance0, amountOut);
+    }
+
+    function testRouter_swapExactOut_2Hops() public {
+        uint256 amountOut = 1 ether;
+        uint256 expectedAmountIn = 1016204441757464409;
+
+        tokenPath.push(token0);
+        tokenPath.push(token1);
+        tokenPath.push(token2);
+        Routing.ExactOutputParams memory params = getExactOutputParams(tokenPath, amountOut);
+
+        uint256 prevBalance0 = token0.balanceOf(address(this));
+        uint256 prevBalance1 = token1.balanceOf(address(this));
+        uint256 prevBalance2 = token2.balanceOf(address(this));
+
+        snapStart("RouterExactOut2Hops");
+        router.swap(Routing.SwapType.ExactOutput, abi.encode(params));
+        snapEnd();
+
+        uint256 newBalance0 = token0.balanceOf(address(this));
+        uint256 newBalance1 = token1.balanceOf(address(this));
+        uint256 newBalance2 = token2.balanceOf(address(this));
+
+        assertEq(prevBalance0 - newBalance0, expectedAmountIn);
+        assertEq(prevBalance1 - newBalance1, 0);
+        assertEq(newBalance2 - prevBalance2, amountOut);
+        assertEq(token0.balanceOf(address(router)), 0);
+        assertEq(token1.balanceOf(address(router)), 0);
+        assertEq(token2.balanceOf(address(router)), 0);
+    }
+
+    function testRouter_swapExactOut_3Hops() public {
+        uint256 amountOut = 1 ether;
+        uint256 expectedAmountIn = 1024467570922834110;
+
+        tokenPath.push(token0);
+        tokenPath.push(token1);
+        tokenPath.push(token2);
+        tokenPath.push(token3);
+        Routing.ExactOutputParams memory params = getExactOutputParams(tokenPath, amountOut);
+
+        uint256 prevBalance0 = token0.balanceOf(address(this));
+        uint256 prevBalance3 = token3.balanceOf(address(this));
+
+        snapStart("RouterExactOut3Hops");
+        router.swap(Routing.SwapType.ExactOutput, abi.encode(params));
+        snapEnd();
+
+        uint256 newBalance0 = token0.balanceOf(address(this));
+        uint256 newBalance3 = token3.balanceOf(address(this));
+
+        assertEq(prevBalance0 - newBalance0, expectedAmountIn);
+        assertEq(newBalance3 - prevBalance3, amountOut);
+        assertEq(token0.balanceOf(address(router)), 0);
+        assertEq(token1.balanceOf(address(router)), 0);
+        assertEq(token2.balanceOf(address(router)), 0);
+        assertEq(token3.balanceOf(address(router)), 0);
+    }
+
     function createPoolKey(MockERC20 tokenA, MockERC20 tokenB) internal pure returns (PoolKey memory) {
         if (address(tokenA) > address(tokenB)) (tokenA, tokenB) = (tokenB, tokenA);
         return PoolKey(Currency.wrap(address(tokenA)), Currency.wrap(address(tokenB)), 3000, 60, IHooks(address(0)));
@@ -230,4 +332,21 @@ contract RoutingTest is Test, Deployers, GasSnapshot {
         params.amountIn = uint128(amountIn);
         params.amountOutMinimum = 0;
     }
+
+    function getExactOutputParams(MockERC20[] memory _tokenPath, uint256 amountOut)
+        internal
+        view
+        returns (Routing.ExactOutputParams memory params)
+    {
+        Routing.PathKey[] memory path = new Routing.PathKey[](_tokenPath.length - 1);
+        for (uint256 i = _tokenPath.length - 1; i > 0 ; i--) {
+            path[i - 1] = Routing.PathKey(Currency.wrap(address(_tokenPath[i - 1])), 3000, 60, IHooks(address(0)));
+        }
+
+        params.currencyOut = Currency.wrap(address(_tokenPath[_tokenPath.length - 1]));
+        params.path = path;
+        params.recipient = address(this);
+        params.amountOut = uint128(amountOut);
+        params.amountInMaximum = type(uint128).max;
+    }
 }