diff --git a/demo-contract/contracts/wtfswap/Pool.sol b/demo-contract/contracts/wtfswap/Pool.sol index 5692233..179d76d 100644 --- a/demo-contract/contracts/wtfswap/Pool.sol +++ b/demo-contract/contracts/wtfswap/Pool.sol @@ -8,6 +8,7 @@ import "./libraries/SqrtPriceMath.sol"; import "./libraries/TickMath.sol"; import "./libraries/LiquidityMath.sol"; import "./libraries/LowGasSafeMath.sol"; +import "./libraries/TransferHelper.sol"; import "./interfaces/IPool.sol"; import "./interfaces/IFactory.sol"; @@ -143,24 +144,49 @@ contract Pool is IPool { function collect( address recipient ) external override returns (uint128 amount0, uint128 amount1) { - // 获取当前用户的 position,TODO recipient 应该改为 msg.sender - Position storage position = positions[recipient]; - // TODO 把钱退给用户 recipient - // 修改 position 中的信息 - position.tokensOwed0 -= amount0; - position.tokensOwed1 -= amount1; + // 获取当前用户的 position + Position storage position = positions[msg.sender]; + // 把钱退给用户 recipient,只支持全部退还 + amount0 = position.tokensOwed0; + amount1 = position.tokensOwed1; + + if (amount0 > 0) { + position.tokensOwed0 -= amount0; + TransferHelper.safeTransfer(token0, recipient, amount0); + } + if (amount1 > 0) { + position.tokensOwed1 -= amount1; + TransferHelper.safeTransfer(token1, recipient, amount1); + } + + emit Collect(msg.sender, recipient, amount0, amount1); } function burn( uint128 amount ) external override returns (uint256 amount0, uint256 amount1) { // 修改 positions 中的信息 - positions[msg.sender].liquidity -= amount; + (int256 amount0Int, int256 amount1Int) = _modifyPosition( + ModifyPositionParams({ + owner: msg.sender, + liquidityDelta: -int256(int128(amount)).toInt128() + }) + ); // 获取燃烧后的 amount0 和 amount1 - // TODO 当前先写个假的 - (amount0, amount1) = (amount / 2, amount / 2); - positions[msg.sender].tokensOwed0 += amount0; - positions[msg.sender].tokensOwed1 += amount1; + amount0 = uint256(-amount0Int); + amount1 = uint256(-amount1Int); + + if (amount0 > 0 || amount1 > 0) { + ( + positions[msg.sender].tokensOwed0, + positions[msg.sender].tokensOwed1 + ) = ( + positions[msg.sender].tokensOwed0 + uint128(amount0), + positions[msg.sender].tokensOwed1 + uint128(amount1) + ); + } + + emit Burn(msg.sender, amount, amount0, amount1); } function swap( diff --git a/demo-contract/contracts/wtfswap/interfaces/IPool.sol b/demo-contract/contracts/wtfswap/interfaces/IPool.sol index 149d986..05ba839 100644 --- a/demo-contract/contracts/wtfswap/interfaces/IPool.sol +++ b/demo-contract/contracts/wtfswap/interfaces/IPool.sol @@ -20,8 +20,8 @@ interface ISwapCallback { interface IPool { struct Position { uint128 liquidity; - uint256 tokensOwed0; - uint256 tokensOwed1; + uint128 tokensOwed0; + uint128 tokensOwed1; } function factory() external view returns (address); diff --git a/demo-contract/contracts/wtfswap/libraries/TransferHelper.sol b/demo-contract/contracts/wtfswap/libraries/TransferHelper.sol new file mode 100644 index 0000000..41153bc --- /dev/null +++ b/demo-contract/contracts/wtfswap/libraries/TransferHelper.sol @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: GPL-2.0-or-later +pragma solidity >=0.6.0; + +import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; + +/// @title TransferHelper +/// @notice Contains helper methods for interacting with ERC20 tokens that do not consistently return true/false +library TransferHelper { + /// @notice Transfers tokens from msg.sender to a recipient + /// @dev Calls transfer on token contract, errors with TF if transfer fails + /// @param token The contract address of the token which will be transferred + /// @param to The recipient of the transfer + /// @param value The value of the transfer + function safeTransfer(address token, address to, uint256 value) internal { + (bool success, bytes memory data) = token.call( + abi.encodeWithSelector(IERC20.transfer.selector, to, value) + ); + require( + success && (data.length == 0 || abi.decode(data, (bool))), + "TF" + ); + } +}