add primitive tests

contracts/hooks/examples/BaseMiddleware.sol

@@ -0,0 +1,112 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.24;
+
+import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
+import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
+import {BeforeSwapDelta} from "@uniswap/v4-core/src/types/BeforeSwapDelta.sol";
+
+contract BaseMiddleware is IHooks {
+    error NotPoolManager();
+    error NotSelf();
+    error InvalidPool();
+    error LockFailure();
+    error HookNotImplemented();
+
+    /// @notice The address of the pool manager
+    IPoolManager public immutable poolManager;
+    IHooks public immutable implementation;
+
+    constructor(IPoolManager _poolManager, IHooks _implementation) {
+        poolManager = _poolManager;
+        implementation = _implementation;
+    }
+
+    //function getHookPermissions() public pure virtual returns (Hooks.Permissions memory);
+
+    function beforeInitialize(address, PoolKey calldata, uint160, bytes calldata) external virtual returns (bytes4) {
+        revert HookNotImplemented();
+    }
+
+    function afterInitialize(address, PoolKey calldata, uint160, int24, bytes calldata)
+        external
+        virtual
+        returns (bytes4)
+    {
+        revert HookNotImplemented();
+    }
+
+    function beforeAddLiquidity(address, PoolKey calldata, IPoolManager.ModifyLiquidityParams calldata, bytes calldata)
+        external
+        virtual
+        returns (bytes4)
+    {
+        revert HookNotImplemented();
+    }
+
+    function beforeRemoveLiquidity(
+        address,
+        PoolKey calldata,
+        IPoolManager.ModifyLiquidityParams calldata,
+        bytes calldata
+    ) external virtual returns (bytes4) {
+        revert HookNotImplemented();
+    }
+
+    function afterAddLiquidity(
+        address,
+        PoolKey calldata,
+        IPoolManager.ModifyLiquidityParams calldata,
+        BalanceDelta,
+        bytes calldata
+    ) external virtual returns (bytes4, BalanceDelta) {
+        revert HookNotImplemented();
+    }
+
+    function afterRemoveLiquidity(
+        address,
+        PoolKey calldata,
+        IPoolManager.ModifyLiquidityParams calldata,
+        BalanceDelta,
+        bytes calldata
+    ) external virtual returns (bytes4, BalanceDelta) {
+        revert HookNotImplemented();
+    }
+
+    function beforeSwap(
+        address sender,
+        PoolKey calldata key,
+        IPoolManager.SwapParams calldata params,
+        bytes calldata hookData
+    ) external returns (bytes4, BeforeSwapDelta, uint24) {
+        return implementation.beforeSwap(sender, key, params, hookData);
+    }
+
+    function afterSwap(
+        address sender,
+        PoolKey calldata key,
+        IPoolManager.SwapParams calldata params,
+        BalanceDelta delta,
+        bytes calldata hookData
+    ) external virtual returns (bytes4, int128) {
+        return implementation.afterSwap(sender, key, params, delta, hookData);
+    }
+
+    function beforeDonate(address, PoolKey calldata, uint256, uint256, bytes calldata)
+        external
+        virtual
+        returns (bytes4)
+    {
+        revert HookNotImplemented();
+    }
+
+    function afterDonate(address, PoolKey calldata, uint256, uint256, bytes calldata)
+        external
+        virtual
+        returns (bytes4)
+    {
+        revert HookNotImplemented();
+    }
+}

contracts/hooks/examples/FeeTakings.sol

@@ -0,0 +1,100 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import {BaseHook} from "../../BaseHook.sol";
+import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
+import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+import {BeforeSwapDelta, BeforeSwapDeltaLibrary} from "@uniswap/v4-core/src/types/BeforeSwapDelta.sol";
+import {Currency, CurrencyLibrary} from "@uniswap/v4-core/src/types/Currency.sol";
+import {SafeCast} from "@uniswap/v4-core/src/libraries/SafeCast.sol";
+import {Owned} from "solmate/auth/Owned.sol";
+import {IUnlockCallback} from "@uniswap/v4-core/src/interfaces/callback/IUnlockCallback.sol";
+
+contract FeeTakings is IUnlockCallback, Owned {
+    using SafeCast for uint256;
+
+    bytes internal constant ZERO_BYTES = bytes("");
+    uint128 private constant TOTAL_BIPS = 10000;
+    uint128 private constant MAX_BIPS = 100;
+    IPoolManager public poolManager;
+    uint128 public swapFeeBips;
+
+    struct CallbackData {
+        address to;
+        Currency[] currencies;
+    }
+
+    constructor(IPoolManager _poolManager, uint128 _swapFeeBips, address _owner) Owned(_owner) {
+        poolManager = _poolManager;
+        swapFeeBips = _swapFeeBips;
+    }
+
+    function getHookPermissions() public pure returns (Hooks.Permissions memory) {
+        return Hooks.Permissions({
+            beforeInitialize: false,
+            afterInitialize: false,
+            beforeAddLiquidity: false,
+            afterAddLiquidity: false,
+            beforeRemoveLiquidity: false,
+            afterRemoveLiquidity: false,
+            beforeSwap: false,
+            afterSwap: true,
+            beforeDonate: false,
+            afterDonate: false,
+            beforeSwapReturnDelta: false,
+            afterSwapReturnDelta: true,
+            afterAddLiquidityReturnDelta: false,
+            afterRemoveLiquidityReturnDelta: false
+        });
+    }
+
+    function afterSwap(
+        address,
+        PoolKey calldata key,
+        IPoolManager.SwapParams calldata params,
+        BalanceDelta delta,
+        bytes calldata
+    ) external returns (bytes4, int128) {
+        // fee will be in the unspecified token of the swap
+        bool currency0Specified = (params.amountSpecified < 0 == params.zeroForOne);
+        (Currency feeCurrency, int128 swapAmount) =
+            (currency0Specified) ? (key.currency1, delta.amount1()) : (key.currency0, delta.amount0());
+        // if fee is on output, get the absolute output amount
+        if (swapAmount < 0) swapAmount = -swapAmount;
+
+        uint256 feeAmount = 0;//(uint128(swapAmount) * swapFeeBips) / TOTAL_BIPS;
+        // mint ERC6909 instead of take to avoid edge case where PM doesn't have enough balance
+        //poolManager.mint(address(this), CurrencyLibrary.toId(feeCurrency), feeAmount);
+
+        return (BaseHook.afterSwap.selector, feeAmount.toInt128());
+    }
+
+    function setSwapFeeBips(uint128 _swapFeeBips) external onlyOwner {
+        require(_swapFeeBips <= MAX_BIPS);
+        swapFeeBips = _swapFeeBips;
+    }
+
+    function withdraw(address to, Currency[] calldata currencies) external onlyOwner {
+        poolManager.unlock(abi.encode(CallbackData(to, currencies)));
+    }
+
+    function unlockCallback(bytes calldata rawData)
+        external
+        override
+        returns (bytes memory)
+    {
+        CallbackData memory data = abi.decode(rawData, (CallbackData));
+        uint256 length = data.currencies.length;
+        for (uint256 i = 0; i < length;) {
+            uint256 amount = poolManager.balanceOf(address(this), CurrencyLibrary.toId(data.currencies[i]));
+            poolManager.burn(address(this), CurrencyLibrary.toId(data.currencies[i]), amount);
+            poolManager.take(data.currencies[i], data.to, amount);
+            unchecked {
+                i++;
+            }
+        }
+        return ZERO_BYTES;
+    }
+}

test/BaseMiddleware.t.sol

@@ -0,0 +1,114 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {Test} from "forge-std/Test.sol";
+import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
+import {FeeTakings} from "../contracts/hooks/examples/FeeTakings.sol";
+import {BaseMiddleware} from "../contracts/hooks/examples/BaseMiddleware.sol";
+import {BaseMiddlewareImplementation} from "./shared/implementation/BaseMiddlewareImplementation.sol";
+import {PoolManager} from "@uniswap/v4-core/src/PoolManager.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {Deployers} from "@uniswap/v4-core/test/utils/Deployers.sol";
+import {TestERC20} from "@uniswap/v4-core/src/test/TestERC20.sol";
+import {CurrencyLibrary, Currency} from "@uniswap/v4-core/src/types/Currency.sol";
+import {PoolId, PoolIdLibrary} from "@uniswap/v4-core/src/types/PoolId.sol";
+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 {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
+
+contract BaseMiddlewareTest is Test, Deployers {
+    using PoolIdLibrary for PoolKey;
+    using StateLibrary for IPoolManager;
+
+    uint160 constant SQRT_RATIO_10_1 = 250541448375047931186413801569;
+
+    address constant TREASURY = address(0x1234567890123456789012345678901234567890);
+    uint128 private constant TOTAL_BIPS = 10000;
+
+    HookEnabledSwapRouter router;
+    TestERC20 token0;
+    TestERC20 token1;
+    FeeTakings feeTaking = new FeeTakings(manager, 25, address(this));
+    BaseMiddleware baseMiddleware =
+        BaseMiddleware(address(uint160(Hooks.AFTER_SWAP_FLAG | Hooks.AFTER_SWAP_RETURNS_DELTA_FLAG)));
+    PoolId id;
+
+    function setUp() public {
+        deployFreshManagerAndRouters();
+        (currency0, currency1) = deployMintAndApprove2Currencies();
+
+        router = new HookEnabledSwapRouter(manager);
+        token0 = TestERC20(Currency.unwrap(currency0));
+        token1 = TestERC20(Currency.unwrap(currency1));
+
+        vm.record();
+        BaseMiddlewareImplementation impl = new BaseMiddlewareImplementation(manager, IHooks(address(feeTaking)), baseMiddleware);
+        (, bytes32[] memory writes) = vm.accesses(address(impl));
+        vm.etch(address(baseMiddleware), 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(baseMiddleware), slot, vm.load(address(impl), slot));
+            }
+        }
+
+        // key = PoolKey(currency0, currency1, 3000, 60, baseMiddleware);
+        (key, id) = initPoolAndAddLiquidity(currency0, currency1, baseMiddleware, 3000, SQRT_PRICE_1_1, ZERO_BYTES);
+
+        token0.approve(address(baseMiddleware), type(uint256).max);
+        token1.approve(address(baseMiddleware), type(uint256).max);
+        token0.approve(address(router), type(uint256).max);
+        token1.approve(address(router), type(uint256).max);
+    }
+
+    function testNormal() public {
+        // Swap exact token0 for token1 //
+        bool zeroForOne = true;
+        int256 amountSpecified = -1e12;
+        BalanceDelta swapDelta = swap(key, zeroForOne, amountSpecified, ZERO_BYTES);
+        // ---------------------------- //
+
+        // uint128 output = uint128(swapDelta.amount1());
+        // assertTrue(output > 0);
+
+        // uint256 expectedFee = calculateFeeForExactInput(output, 25);
+
+        // assertEq(manager.balanceOf(address(baseMiddleware), CurrencyLibrary.toId(key.currency0)), 0);
+        // assertEq(manager.balanceOf(address(baseMiddleware), CurrencyLibrary.toId(key.currency1)), expectedFee);
+
+        // // 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);
+        // // ---------------------------- //
+
+        // uint128 input = uint128(-swapDelta2.amount0());
+        // assertTrue(output > 0);
+
+        // uint256 expectedFee2 = calculateFeeForExactOutput(input, 25);
+
+        // assertEq(manager.balanceOf(address(baseMiddleware), CurrencyLibrary.toId(key.currency0)), expectedFee2);
+        // assertEq(manager.balanceOf(address(baseMiddleware), CurrencyLibrary.toId(key.currency1)), expectedFee);
+
+        // // test withdrawing tokens //
+        // Currency[] memory currencies = new Currency[](2);
+        // currencies[0] = key.currency0;
+        // currencies[1] = key.currency1;
+        // feeTaking.withdraw(TREASURY, currencies);
+        // assertEq(manager.balanceOf(address(baseMiddleware), CurrencyLibrary.toId(key.currency0)), 0);
+        // assertEq(manager.balanceOf(address(baseMiddleware), CurrencyLibrary.toId(key.currency1)), 0);
+        // assertEq(currency0.balanceOf(TREASURY), expectedFee2);
+        // assertEq(currency1.balanceOf(TREASURY), expectedFee);
+    }
+
+    function calculateFeeForExactInput(uint256 outputAmount, uint128 feeBips) internal pure returns (uint256) {
+        return outputAmount * TOTAL_BIPS / (TOTAL_BIPS - feeBips) - outputAmount;
+    }
+
+    function calculateFeeForExactOutput(uint256 inputAmount, uint128 feeBips) internal pure returns (uint256) {
+        return (inputAmount * feeBips) / (TOTAL_BIPS + feeBips);
+    }
+}

test/shared/implementation/BaseMiddlewareImplementation.sol

@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {BaseHook} from "../../../contracts/BaseHook.sol";
+import {BaseMiddleware} from "../../../contracts/hooks/examples/BaseMiddleware.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {Hooks} from "@uniswap/v4-core/src/libraries/Hooks.sol";
+import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
+
+contract BaseMiddlewareImplementation is BaseMiddleware {
+    constructor(IPoolManager _poolManager, IHooks _implementation, BaseMiddleware addressToEtch)
+        BaseMiddleware(_poolManager, _implementation)
+    {
+        //Hooks.validateHookPermissions(addressToEtch, getHookPermissions());
+    }
+
+    // make this a no-op in testing
+    //function validateHookAddress(BaseHook _this) internal pure override {}
+}