Use safe transfer for take

contracts/PoolManager.sol

@@ -5,6 +5,7 @@ import {Hooks} from './libraries/Hooks.sol';
 import {Pool} from './libraries/Pool.sol';
 import {Tick} from './libraries/Tick.sol';
 import {SafeCast} from './libraries/SafeCast.sol';
+import {TransferHelper} from './libraries/TransferHelper.sol';
 
 import {IERC20Minimal} from './interfaces/external/IERC20Minimal.sol';
 import {NoDelegateCall} from './NoDelegateCall.sol';
@@ -20,6 +21,7 @@ contract PoolManager is IPoolManager, NoDelegateCall, ERC1155, IERC1155Receiver
     using SafeCast for *;
     using Pool for *;
     using Hooks for IHooks;
+    using TransferHelper for IERC20Minimal;
 
     /// @inheritdoc IPoolManager
     int24 public constant override MAX_TICK_SPACING = type(int16).max;
@@ -260,7 +262,7 @@ contract PoolManager is IPoolManager, NoDelegateCall, ERC1155, IERC1155Receiver
     ) external override noDelegateCall onlyByLocker {
         _accountDelta(token, amount.toInt256());
         reservesOf[token] -= amount;
-        token.transfer(to, amount);
+        token.safeTransfer(to, amount);
     }
 
     /// @inheritdoc IPoolManager

contracts/libraries/TransferHelper.sol

@@ -5,6 +5,7 @@ import {IERC20Minimal} from '../interfaces/external/IERC20Minimal.sol';
 
 /// @title TransferHelper
 /// @notice Contains helper methods for interacting with ERC20 tokens that do not consistently return true/false
+/// @dev implementation from https://github.com/Rari-Capital/solmate/blob/main/src/utils/SafeTransferLib.sol#L63
 library TransferHelper {
     /// @notice Transfers tokens from msg.sender to a recipient
     /// @dev Calls transfer on token contract, errors with TF if transfer fails
@@ -16,9 +17,29 @@ library TransferHelper {
         address to,
         uint256 value
     ) internal {
-        (bool success, bytes memory data) = address(token).call(
-            abi.encodeWithSelector(IERC20Minimal.transfer.selector, to, value)
-        );
-        require(success && (data.length == 0 || abi.decode(data, (bool))), 'TF');
+        bool success;
+
+        assembly {
+            // Get a pointer to some free memory.
+            let freeMemoryPointer := mload(0x40)
+
+            // Write the abi-encoded calldata into memory, beginning with the function selector.
+            mstore(freeMemoryPointer, 0xa9059cbb00000000000000000000000000000000000000000000000000000000)
+            mstore(add(freeMemoryPointer, 4), to) // Append the "to" argument.
+            mstore(add(freeMemoryPointer, 36), value) // Append the "value" argument.
+
+            success := and(
+                // Set success to whether the call reverted, if not we check it either
+                // returned exactly 1 (can't just be non-zero data), or had no return data.
+                or(and(eq(mload(0), 1), gt(returndatasize(), 31)), iszero(returndatasize())),
+                // We use 68 because the length of our calldata totals up like so: 4 + 32 * 2.
+                // We use 0 and 32 to copy up to 32 bytes of return data into the scratch space.
+                // Counterintuitively, this call must be positioned second to the or() call in the
+                // surrounding and() call or else returndatasize() will be zero during the computation.
+                call(gas(), token, 0, freeMemoryPointer, 68, 0, 32)
+            )
+        }
+
+        require(success, 'TRANSFER_FAILED');
     }
 }

contracts/test/PoolTakeTest.sol

@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity =0.8.13;
+
+import {IERC20Minimal} from '../interfaces/external/IERC20Minimal.sol';
+
+import {ILockCallback} from '../interfaces/callback/ILockCallback.sol';
+import {IPoolManager} from '../interfaces/IPoolManager.sol';
+
+contract PoolTakeTest is ILockCallback {
+    IPoolManager public immutable manager;
+
+    constructor(IPoolManager _manager) {
+        manager = _manager;
+    }
+
+    struct CallbackData {
+        address sender;
+        IPoolManager.PoolKey key;
+        uint256 amount0;
+        uint256 amount1;
+    }
+
+    function take(
+        IPoolManager.PoolKey memory key,
+        uint256 amount0,
+        uint256 amount1
+    ) external {
+        manager.lock(abi.encode(CallbackData(msg.sender, key, amount0, amount1)));
+    }
+
+    function lockAcquired(bytes calldata rawData) external returns (bytes memory) {
+        require(msg.sender == address(manager));
+
+        CallbackData memory data = abi.decode(rawData, (CallbackData));
+
+        if (data.amount0 > 0) {
+            uint256 balBefore = data.key.token0.balanceOf(data.sender);
+            manager.take(data.key.token0, data.sender, data.amount0);
+            uint256 balAfter = data.key.token0.balanceOf(data.sender);
+            require(balAfter - balBefore == data.amount0);
+
+            data.key.token0.transferFrom(data.sender, address(manager), uint256(data.amount0));
+            manager.settle(data.key.token0);
+        }
+
+        if (data.amount1 > 0) {
+            uint256 balBefore = data.key.token1.balanceOf(data.sender);
+            manager.take(data.key.token1, data.sender, data.amount1);
+            uint256 balAfter = data.key.token1.balanceOf(data.sender);
+            require(balAfter - balBefore == data.amount1);
+
+            data.key.token1.transferFrom(data.sender, address(manager), uint256(data.amount1));
+            manager.settle(data.key.token1);
+        }
+
+        return abi.encode(0);
+    }
+}

contracts/test/TestInvalidERC20.sol

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity =0.8.13;
+
+import {IERC20Minimal} from '../interfaces/external/IERC20Minimal.sol';
+
+// Regular ERC20 but it doesn't return true on transfer.
+contract TestInvalidERC20 is IERC20Minimal {
+    mapping(address => uint256) public override balanceOf;
+    mapping(address => mapping(address => uint256)) public override allowance;
+
+    constructor(uint256 amountToMint) {
+        mint(msg.sender, amountToMint);
+    }
+
+    function mint(address to, uint256 amount) public {
+        uint256 balanceNext = balanceOf[to] + amount;
+        require(balanceNext >= amount, 'overflow balance');
+        balanceOf[to] = balanceNext;
+    }
+
+    function transfer(address recipient, uint256 amount) external override returns (bool) {
+        uint256 balanceBefore = balanceOf[msg.sender];
+        require(balanceBefore >= amount, 'insufficient balance');
+        balanceOf[msg.sender] = balanceBefore - amount;
+
+        uint256 balanceRecipient = balanceOf[recipient];
+        require(balanceRecipient + amount >= balanceRecipient, 'recipient balance overflow');
+        balanceOf[recipient] = balanceRecipient + amount;
+
+        emit Transfer(msg.sender, recipient, amount);
+        return false;
+    }
+
+    function approve(address spender, uint256 amount) external override returns (bool) {
+        allowance[msg.sender][spender] = amount;
+        emit Approval(msg.sender, spender, amount);
+        return true;
+    }
+
+    function transferFrom(
+        address sender,
+        address recipient,
+        uint256 amount
+    ) external override returns (bool) {
+        uint256 allowanceBefore = allowance[sender][msg.sender];
+        require(allowanceBefore >= amount, 'allowance insufficient');
+
+        allowance[sender][msg.sender] = allowanceBefore - amount;
+
+        uint256 balanceRecipient = balanceOf[recipient];
+        require(balanceRecipient + amount >= balanceRecipient, 'overflow balance recipient');
+        balanceOf[recipient] = balanceRecipient + amount;
+        uint256 balanceSender = balanceOf[sender];
+        require(balanceSender >= amount, 'underflow balance sender');
+        balanceOf[sender] = balanceSender - amount;
+
+        emit Transfer(sender, recipient, amount);
+        return false;
+    }
+}

package.json

@@ -19,8 +19,7 @@
   "files": [
     "contracts/interfaces",
     "contracts/libraries",
-    "artifacts/contracts/PoolFactory.sol/PoolFactory.json",
-    "artifacts/contracts/Pool.sol/Pool.json",
+    "artifacts/contracts/PoolManager.sol/PoolManager.json",
     "artifacts/contracts/interfaces/**/*.json",
     "!artifacts/contracts/interfaces/**/*.dbg.json"
   ],

test/PoolManager.spec.ts

@@ -10,12 +10,13 @@ import {
   PoolModifyPositionTest,
   PoolSwapTest,
   TestERC20,
+  PoolTakeTest,
 } from '../typechain'
 import { MAX_TICK_SPACING } from './shared/constants'
 import { expect } from './shared/expect'
 import { tokensFixture } from './shared/fixtures'
+import { encodeSqrtPriceX96, expandTo18Decimals, FeeAmount, getPoolId, MaxUint128 } from './shared/utilities'
 import { deployMockContract, MockedContract } from './shared/mockContract'
-import { encodeSqrtPriceX96, expandTo18Decimals, FeeAmount, getPoolId } from './shared/utilities'
 
 const createFixtureLoader = waffle.createFixtureLoader
 
@@ -31,6 +32,7 @@ describe('PoolManager', () => {
   let swapTest: PoolSwapTest
   let modifyPositionTest: PoolModifyPositionTest
   let donateTest: PoolDonateTest
+  let takeTest: PoolTakeTest
   let hooksMock: MockedContract
   let testHooksEmpty: EmptyTestHooks
   let tokens: { token0: TestERC20; token1: TestERC20; token2: TestERC20 }
@@ -41,6 +43,7 @@ describe('PoolManager', () => {
     const swapTestFactory = await ethers.getContractFactory('PoolSwapTest')
     const modifyPositionTestFactory = await ethers.getContractFactory('PoolModifyPositionTest')
     const donateTestFactory = await ethers.getContractFactory('PoolDonateTest')
+    const takeTestFactory = await ethers.getContractFactory('PoolTakeTest')
     const hooksTestEmptyFactory = await ethers.getContractFactory('EmptyTestHooks')
     const tokens = await tokensFixture()
     const manager = (await poolManagerFactory.deploy()) as PoolManager
@@ -64,13 +67,14 @@ describe('PoolManager', () => {
       swapTest: (await swapTestFactory.deploy(manager.address)) as PoolSwapTest,
       modifyPositionTest: (await modifyPositionTestFactory.deploy(manager.address)) as PoolModifyPositionTest,
       donateTest: (await donateTestFactory.deploy(manager.address)) as PoolDonateTest,
+      takeTest: (await takeTestFactory.deploy(manager.address)) as PoolTakeTest,
       tokens,
       hooksMock,
       testHooksEmpty,
     }
 
     for (const token of [tokens.token0, tokens.token1, tokens.token2]) {
-      for (const spender of [result.swapTest, result.modifyPositionTest, result.donateTest]) {
+      for (const spender of [result.swapTest, result.modifyPositionTest, result.donateTest, result.takeTest]) {
         await token.connect(wallet).approve(spender.address, constants.MaxUint256)
       }
     }
@@ -86,7 +90,7 @@ describe('PoolManager', () => {
   })
 
   beforeEach('deploy fixture', async () => {
-    ;({ manager, tokens, lockTest, modifyPositionTest, swapTest, donateTest, hooksMock, testHooksEmpty } =
+    ;({ manager, tokens, lockTest, modifyPositionTest, swapTest, donateTest, takeTest, hooksMock, testHooksEmpty } =
       await loadFixture(fixture))
   })
 
@@ -845,6 +849,72 @@ describe('PoolManager', () => {
     })
   })
 
+  describe('#take', () => {
+    it('fails if no liquidity', async () => {
+      await tokens.token0.connect(wallet).transfer(ADDRESS_ZERO, constants.MaxUint256.div(2))
+      await expect(
+        takeTest.connect(wallet).take(
+          {
+            token0: tokens.token0.address,
+            token1: tokens.token1.address,
+            fee: 100,
+            hooks: ADDRESS_ZERO,
+            tickSpacing: 10,
+          },
+          100,
+          0
+        )
+      ).to.be.reverted
+    })
+
+    it('fails for invalid tokens that dont return true on transfer', async () => {
+      const tokenFactory = await ethers.getContractFactory('TestInvalidERC20')
+      const invalidToken = (await tokenFactory.deploy(BigNumber.from(2).pow(255))) as TestERC20
+      const token0Invalid = invalidToken.address.toLowerCase() < tokens.token0.address.toLowerCase()
+      const key = {
+        token0: token0Invalid ? invalidToken.address : tokens.token0.address,
+        token1: token0Invalid ? tokens.token0.address : invalidToken.address,
+        fee: 100,
+        hooks: ADDRESS_ZERO,
+        tickSpacing: 10,
+      }
+      await invalidToken.approve(modifyPositionTest.address, constants.MaxUint256)
+      await manager.initialize(key, encodeSqrtPriceX96(1, 1))
+      await modifyPositionTest.modifyPosition(key, {
+        tickLower: -60,
+        tickUpper: 60,
+        liquidityDelta: 100,
+      })
+
+      await tokens.token0.connect(wallet).approve(takeTest.address, MaxUint128)
+      await invalidToken.connect(wallet).approve(takeTest.address, MaxUint128)
+
+      await expect(takeTest.connect(wallet).take(key, token0Invalid ? 1 : 0, token0Invalid ? 0 : 1)).to.be.reverted
+      await expect(takeTest.connect(wallet).take(key, token0Invalid ? 0 : 1, token0Invalid ? 1 : 0)).to.not.be.reverted
+    })
+
+    it('succeeds if has liquidity', async () => {
+      const key = {
+        token0: tokens.token0.address,
+        token1: tokens.token1.address,
+        fee: 100,
+        hooks: ADDRESS_ZERO,
+        tickSpacing: 10,
+      }
+      await manager.initialize(key, encodeSqrtPriceX96(1, 1))
+      await modifyPositionTest.modifyPosition(key, {
+        tickLower: -60,
+        tickUpper: 60,
+        liquidityDelta: 100,
+      })
+
+      await tokens.token0.connect(wallet).approve(takeTest.address, MaxUint128)
+
+      await expect(takeTest.connect(wallet).take(key, 1, 0)).to.be.not.be.reverted
+      await expect(takeTest.connect(wallet).take(key, 0, 1)).to.be.not.be.reverted
+    })
+  })
+
   describe('#donate', () => {
     it('fails if not initialized', async () => {
       await expect(

test/__snapshots__/PoolManager.gas.spec.ts.snap

@@ -66,62 +66,62 @@ Object {
 exports[`PoolManager gas tests #swapExact0For1 first swap in block moves tick, no initialized crossings 1`] = `
 Object {
   "calldataByteLength": 324,
-  "gasUsed": 229688,
+  "gasUsed": 229568,
 }
 `;
 
 exports[`PoolManager gas tests #swapExact0For1 first swap in block with no tick movement 1`] = `
 Object {
   "calldataByteLength": 324,
-  "gasUsed": 223794,
+  "gasUsed": 223673,
 }
 `;
 
 exports[`PoolManager gas tests #swapExact0For1 first swap in block, large swap crossing a single initialized tick 1`] = `
 Object {
   "calldataByteLength": 324,
-  "gasUsed": 248636,
+  "gasUsed": 248516,
 }
 `;
 
 exports[`PoolManager gas tests #swapExact0For1 first swap in block, large swap crossing several initialized ticks 1`] = `
 Object {
   "calldataByteLength": 324,
-  "gasUsed": 290680,
+  "gasUsed": 290560,
 }
 `;
 
 exports[`PoolManager gas tests #swapExact0For1 first swap in block, large swap, no initialized crossings 1`] = `
 Object {
   "calldataByteLength": 324,
-  "gasUsed": 241674,
+  "gasUsed": 241553,
 }
 `;
 
 exports[`PoolManager gas tests #swapExact0For1 second swap in block moves tick, no initialized crossings 1`] = `
 Object {
   "calldataByteLength": 324,
-  "gasUsed": 229688,
+  "gasUsed": 229568,
 }
 `;
 
 exports[`PoolManager gas tests #swapExact0For1 second swap in block with no tick movement 1`] = `
 Object {
   "calldataByteLength": 324,
-  "gasUsed": 223883,
+  "gasUsed": 223762,
 }
 `;
 
 exports[`PoolManager gas tests #swapExact0For1 second swap in block, large swap crossing a single initialized tick 1`] = `
 Object {
   "calldataByteLength": 324,
-  "gasUsed": 242860,
+  "gasUsed": 242739,
 }
 `;
 
 exports[`PoolManager gas tests #swapExact0For1 second swap in block, large swap crossing several initialized ticks 1`] = `
 Object {
   "calldataByteLength": 324,
-  "gasUsed": 284878,
+  "gasUsed": 284757,
 }
 `;