From 35ba973e854ba08c3167bba34bcd2a3d20e6da74 Mon Sep 17 00:00:00 2001 From: Junion <69495294+Jun1on@users.noreply.github.com> Date: Wed, 3 Jul 2024 15:48:07 -0400 Subject: [PATCH] test _afterSwap extension --- .forge-snapshots/FeeTakingFirstSwap.snap | 2 +- .forge-snapshots/FeeTakingSecondSwap.snap | 2 +- .../FeeTakingWithdrawTwoTokens.snap | 2 +- contracts/hooks/examples/FeeTaker.sol | 2 +- contracts/hooks/examples/FeeTaking.sol | 2 +- lib/forge-gas-snapshot | 2 +- lib/forge-std | 2 +- test/FeeTaking.t.sol | 107 ++++++++++++++++-- .../implementation/FeeTakingExtension.sol | 61 ++++++++++ 9 files changed, 168 insertions(+), 14 deletions(-) create mode 100644 test/shared/implementation/FeeTakingExtension.sol diff --git a/.forge-snapshots/FeeTakingFirstSwap.snap b/.forge-snapshots/FeeTakingFirstSwap.snap index 51d070ee..ae8cb8b8 100644 --- a/.forge-snapshots/FeeTakingFirstSwap.snap +++ b/.forge-snapshots/FeeTakingFirstSwap.snap @@ -1 +1 @@ -162044 \ No newline at end of file +112467 \ No newline at end of file diff --git a/.forge-snapshots/FeeTakingSecondSwap.snap b/.forge-snapshots/FeeTakingSecondSwap.snap index 74b56108..d18f8244 100644 --- a/.forge-snapshots/FeeTakingSecondSwap.snap +++ b/.forge-snapshots/FeeTakingSecondSwap.snap @@ -1 +1 @@ -86285 \ No newline at end of file +86214 \ No newline at end of file diff --git a/.forge-snapshots/FeeTakingWithdrawTwoTokens.snap b/.forge-snapshots/FeeTakingWithdrawTwoTokens.snap index 6565a566..98af7300 100644 --- a/.forge-snapshots/FeeTakingWithdrawTwoTokens.snap +++ b/.forge-snapshots/FeeTakingWithdrawTwoTokens.snap @@ -1 +1 @@ -71441 \ No newline at end of file +69428 \ No newline at end of file diff --git a/contracts/hooks/examples/FeeTaker.sol b/contracts/hooks/examples/FeeTaker.sol index 91d1fb7d..a58a9631 100644 --- a/contracts/hooks/examples/FeeTaker.sol +++ b/contracts/hooks/examples/FeeTaker.sol @@ -1,4 +1,4 @@ -// SPDX-License-Identifier: MIT +// SPDX-License-Identifier: UNLICENSED pragma solidity ^0.8.24; import {BaseHook} from "../../BaseHook.sol"; diff --git a/contracts/hooks/examples/FeeTaking.sol b/contracts/hooks/examples/FeeTaking.sol index c0099149..23ead10c 100644 --- a/contracts/hooks/examples/FeeTaking.sol +++ b/contracts/hooks/examples/FeeTaking.sol @@ -1,4 +1,4 @@ -// SPDX-License-Identifier: MIT +// SPDX-License-Identifier: UNLICENSED pragma solidity ^0.8.24; import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol"; diff --git a/lib/forge-gas-snapshot b/lib/forge-gas-snapshot index 2f884282..9161f7c0 160000 --- a/lib/forge-gas-snapshot +++ b/lib/forge-gas-snapshot @@ -1 +1 @@ -Subproject commit 2f884282b4cd067298e798974f5b534288b13bc2 +Subproject commit 9161f7c0b6c6788a89081e2b3b9c67592b71e689 diff --git a/lib/forge-std b/lib/forge-std index 2b58ecbc..75b3fcf0 160000 --- a/lib/forge-std +++ b/lib/forge-std @@ -1 +1 @@ -Subproject commit 2b58ecbcf3dfde7a75959dc7b4eb3d0670278de6 +Subproject commit 75b3fcf052cc7886327e4c2eac3d1a1f36942b41 diff --git a/test/FeeTaking.t.sol b/test/FeeTaking.t.sol index 9bbd0c34..6612380b 100644 --- a/test/FeeTaking.t.sol +++ b/test/FeeTaking.t.sol @@ -15,6 +15,7 @@ import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol"; import {HookEnabledSwapRouter} from "./utils/HookEnabledSwapRouter.sol"; import {StateLibrary} from "@uniswap/v4-core/src/libraries/StateLibrary.sol"; import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol"; +import {FeeTakingExtension} from "./shared/implementation/FeeTakingExtension.sol"; contract FeeTakingTest is Test, Deployers, GasSnapshot { using PoolIdLibrary for PoolKey; @@ -32,9 +33,11 @@ contract FeeTakingTest is Test, Deployers, GasSnapshot { TestERC20 token0; TestERC20 token1; FeeTaking feeTaking = FeeTaking(address(uint160(Hooks.AFTER_SWAP_FLAG | Hooks.AFTER_SWAP_RETURNS_DELTA_FLAG))); + FeeTakingExtension feeTakingExtension = + FeeTakingExtension(address(0x100000 | uint160(Hooks.AFTER_SWAP_FLAG | Hooks.AFTER_SWAP_RETURNS_DELTA_FLAG))); PoolId id; - function setUp() public { + function setUpNormal() public { deployFreshManagerAndRouters(); (currency0, currency1) = deployMintAndApprove2Currencies(); @@ -46,24 +49,46 @@ contract FeeTakingTest is Test, Deployers, GasSnapshot { FeeTakingImplementation impl = new FeeTakingImplementation(manager, 25, address(this), TREASURY, feeTaking); (, bytes32[] memory writes) = vm.accesses(address(impl)); vm.etch(address(feeTaking), address(impl).code); - // for each storage key that was written during the hook implementation, copy the value over unchecked { for (uint256 i = 0; i < writes.length; i++) { bytes32 slot = writes[i]; vm.store(address(feeTaking), slot, vm.load(address(impl), slot)); } } - - // key = PoolKey(currency0, currency1, 3000, 60, feeTaking); (key, id) = initPoolAndAddLiquidity(currency0, currency1, feeTaking, 3000, SQRT_PRICE_1_1, ZERO_BYTES); - token0.approve(address(feeTaking), type(uint256).max); - token1.approve(address(feeTaking), type(uint256).max); token0.approve(address(router), type(uint256).max); token1.approve(address(router), type(uint256).max); } + function setUpExtension() public { + deployFreshManagerAndRouters(); + (currency0, currency1) = deployMintAndApprove2Currencies(); + + router = new HookEnabledSwapRouter(manager); + token0 = TestERC20(Currency.unwrap(currency0)); + token1 = TestERC20(Currency.unwrap(currency1)); + + vm.record(); + FeeTakingExtension impl = new FeeTakingExtension(manager, 25, address(this), TREASURY); + (, bytes32[] memory writes) = vm.accesses(address(impl)); + vm.etch(address(feeTakingExtension), address(impl).code); + unchecked { + for (uint256 i = 0; i < writes.length; i++) { + bytes32 slot = writes[i]; + vm.store(address(feeTakingExtension), slot, vm.load(address(impl), slot)); + } + } + (key, id) = initPoolAndAddLiquidity(currency0, currency1, feeTakingExtension, 3000, SQRT_PRICE_1_1, ZERO_BYTES); + + token0.approve(address(router), type(uint256).max); + token1.approve(address(router), type(uint256).max); + token0.transfer(address(feeTakingExtension), 1e18); + token1.transfer(address(feeTakingExtension), 1e18); + } + function testSwapHooks() public { + setUpNormal(); assertEq(currency0.balanceOf(TREASURY), 0); assertEq(currency1.balanceOf(TREASURY), 0); @@ -90,7 +115,7 @@ contract FeeTakingTest is Test, Deployers, GasSnapshot { snapEnd(); uint128 input = uint128(-swapDelta2.amount0()); - assertTrue(output > 0); + assertTrue(input > 0); uint256 expectedFee2 = calculateFeeForExactOutput(input, feeTaking.swapFeeBips()); @@ -112,6 +137,7 @@ contract FeeTakingTest is Test, Deployers, GasSnapshot { // this would error had the hook not used ERC6909 function testEdgeCase() public { + setUpNormal(); // first, deplete the pool of token1 // Swap exact token0 for token1 // bool zeroForOne = true; @@ -200,6 +226,73 @@ contract FeeTakingTest is Test, Deployers, GasSnapshot { assertEq(currency3.balanceOf(TREASURY) / R, expectedFee / R); } + function testHookExtension() public { + setUpExtension(); + assertEq(currency0.balanceOf(TREASURY), 0); + assertEq(currency1.balanceOf(TREASURY), 0); + + // Swap exact token0 for token1 // + bool zeroForOne = true; + int256 amountSpecified = -1e12; + BalanceDelta swapDelta = swap(key, zeroForOne, amountSpecified, ZERO_BYTES); + + assertEq(feeTakingExtension.afterSwapCounter(), 1); + + uint128 output = uint128(swapDelta.amount1() - feeTakingExtension.DONATION_AMOUNT()); + assertTrue(output > 0); + + uint256 expectedFee = calculateFeeForExactInput(output, feeTakingExtension.swapFeeBips()); + + assertEq(manager.balanceOf(address(feeTakingExtension), CurrencyLibrary.toId(key.currency0)), 0); + assertEq( + manager.balanceOf(address(feeTakingExtension), CurrencyLibrary.toId(key.currency1)) / R, expectedFee / R + ); + + assertEq(currency0.balanceOf(address(feeTakingExtension)), 1 ether); + assertEq( + currency1.balanceOf(address(feeTakingExtension)), + 1 ether - uint256(int256(feeTakingExtension.DONATION_AMOUNT())) + ); + + // Swap token0 for exact token1 // + bool zeroForOne2 = true; + int256 amountSpecified2 = 1e12; // positive number indicates exact output swap + BalanceDelta swapDelta2 = swap(key, zeroForOne2, amountSpecified2, ZERO_BYTES); + return; + assertEq(feeTakingExtension.afterSwapCounter(), 2); + + uint128 input = uint128(-swapDelta2.amount0() + feeTakingExtension.DONATION_AMOUNT()); + assertTrue(input > 0); + + uint256 expectedFee2 = calculateFeeForExactOutput(input, feeTakingExtension.swapFeeBips()); + + assertEq( + manager.balanceOf(address(feeTakingExtension), CurrencyLibrary.toId(key.currency0)) / R, expectedFee2 / R + ); + assertEq( + manager.balanceOf(address(feeTakingExtension), CurrencyLibrary.toId(key.currency1)) / R, expectedFee / R + ); + + assertEq( + currency0.balanceOf(address(feeTakingExtension)), + 1 ether - uint256(int256(feeTakingExtension.DONATION_AMOUNT())) + ); + assertEq( + currency1.balanceOf(address(feeTakingExtension)), + 1 ether - uint256(int256(feeTakingExtension.DONATION_AMOUNT())) + ); + + // test withdrawing tokens // + Currency[] memory currencies = new Currency[](2); + currencies[0] = key.currency0; + currencies[1] = key.currency1; + feeTakingExtension.withdraw(currencies); + assertEq(manager.balanceOf(address(feeTakingExtension), CurrencyLibrary.toId(key.currency0)), 0); + assertEq(manager.balanceOf(address(feeTakingExtension), CurrencyLibrary.toId(key.currency1)), 0); + assertEq(currency0.balanceOf(TREASURY) / R, expectedFee2 / R); + assertEq(currency1.balanceOf(TREASURY) / R, expectedFee / R); + } + function calculateFeeForExactInput(uint256 outputAmount, uint128 feeBips) internal pure returns (uint256) { return outputAmount * TOTAL_BIPS / (TOTAL_BIPS - feeBips) - outputAmount; } diff --git a/test/shared/implementation/FeeTakingExtension.sol b/test/shared/implementation/FeeTakingExtension.sol new file mode 100644 index 00000000..70d0f3b7 --- /dev/null +++ b/test/shared/implementation/FeeTakingExtension.sol @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.24; + +import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol"; +import {SafeCast} from "@uniswap/v4-core/src/libraries/SafeCast.sol"; +import {Owned} from "solmate/auth/Owned.sol"; +import {FeeTaker} from "./../../../contracts/hooks/examples/FeeTaker.sol"; +import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol"; +import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol"; +import {BaseHook} from "./../../../contracts/BaseHook.sol"; +import {Currency, CurrencyLibrary} from "@uniswap/v4-core/src/types/Currency.sol"; + +contract FeeTakingExtension is FeeTaker, Owned { + using SafeCast for uint256; + using CurrencyLibrary for Currency; + + uint128 private constant TOTAL_BIPS = 10000; + uint128 public immutable swapFeeBips; + address public treasury; + + uint256 public afterSwapCounter; + int128 public DONATION_AMOUNT = 1 gwei; + + constructor(IPoolManager _poolManager, uint128 _swapFeeBips, address _owner, address _treasury) + FeeTaker(_poolManager) + Owned(_owner) + { + swapFeeBips = _swapFeeBips; + treasury = _treasury; + } + + function _afterSwap( + address, + PoolKey memory key, + IPoolManager.SwapParams memory params, + BalanceDelta, + bytes calldata + ) internal override returns (bytes4, int128) { + afterSwapCounter++; + bool currency0Specified = (params.amountSpecified < 0 == params.zeroForOne); + Currency currencyUnspecified = currency0Specified ? key.currency1 : key.currency0; + currencyUnspecified.transfer(address(manager), uint256(int256(DONATION_AMOUNT))); + manager.settle(currencyUnspecified); + return (BaseHook.afterSwap.selector, -DONATION_AMOUNT); + } + + function setTreasury(address _treasury) external onlyOwner { + treasury = _treasury; + } + + function _feeAmount(int128 amountUnspecified) internal view override returns (uint256) { + return uint128(amountUnspecified) * swapFeeBips / TOTAL_BIPS; + } + + function _recipient() internal view override returns (address) { + return treasury; + } + + // make this a no-op in testing + function validateHookAddress(BaseHook _this) internal pure override {} +}