diff --git a/src/LiquidityManager.sol b/src/LiquidityManager.sol index ad53d33..9e40340 100644 --- a/src/LiquidityManager.sol +++ b/src/LiquidityManager.sol @@ -49,11 +49,17 @@ contract LiquidityManager { // the fee growth of the aggregate position as of the last action on the individual position uint256 feeGrowthInside0LastX128; uint256 feeGrowthInside1LastX128; - // how many uncollected tokens are owed to the position, as of the last computation - uint128 tokensOwed0; - uint128 tokensOwed1; } + struct FeesOwed { + // how many uncollected tokens are owed + // updated by each position, as of the last computation + uint128 tokensOwed; + uint128 ethOwed; + } + + mapping(address => FeesOwed) private _feesOwed; + /// @dev The token ID position data mapping(bytes32 => TokenPosition) private _positions; @@ -67,6 +73,16 @@ contract LiquidityManager { return (token0isWeth) ? token1 : token0; } + function getFeesOwed(address token) public view returns (uint128 tokensOwed, uint128 ethOwed) { + (tokensOwed, ethOwed) = _feesOwed[token]; + } + + function updateFeesOwed(bool token0isEth, address token, uint128 tokensOwed0, uint128 tokensOwed1) internal { + FeesOwed storage feesOwed = _feesOwed[token]; + feesOwed = (feesOwed.tokensOwed += token0isEth ? tokensOwed1 : tokensOwed0, + feesOwed.ethOwed += token0isEth ? tokensOwed0 : tokensOwed1); + } + function posKey(address _token, uint24 _tickLower, uint24, tickUpper) internal pure returns (uint256) { return uint160(_token) << 48 + tickLower << 24 + tickUpper; } @@ -75,18 +91,14 @@ contract LiquidityManager { returns ( uint128 liquidity, uint256 feeGrowthInside0LastX128, - uint256 feeGrowthInside1LastX128, - uint128 tokensOwed0, - uint128 tokensOwed1 + uint256 feeGrowthInside1LastX128 ) { TokenPosition memory position = _positions[posKey(_token, tickLower, tickUpper)]; return ( position.liquidity, position.feeGrowthInside0LastX128, - position.feeGrowthInside1LastX128, - position.tokensOwed0, - position.tokensOwed1 + position.feeGrowthInside1LastX128 ); } @@ -148,32 +160,30 @@ contract LiquidityManager { (, uint256 feeGrowthInside0LastX128, uint256 feeGrowthInside1LastX128, , ) = pool.positions(positionKey); - TokenPosition memory position = _positions[posKey(getToken(params.token0, params.token1), tickLower, tickUpper)]; + (bool token0isWeth, address token) = getToken(params.token0, params.token1); + TokenPosition memory position = _positions[posKey(token, tickLower, tickUpper)]; if (liquidity == 0) { // create entry position = TokenPosition({ liquidity: liquidity, feeGrowthInside0LastX128: feeGrowthInside0LastX128, - feeGrowthInside1LastX128: feeGrowthInside1LastX128, - tokensOwed0: 0, - tokensOwed1: 0 + feeGrowthInside1LastX128: feeGrowthInside1LastX128 }); } else { // update entry - position.tokensOwed0 += uint128( + updateTokensOwed(token0isWeth, token, uint128( FullMath.mulDiv( feeGrowthInside0LastX128 - position.feeGrowthInside0LastX128, position.liquidity, FixedPoint128.Q128 ) - ); - position.tokensOwed1 += uint128( + ), uint128( FullMath.mulDiv( feeGrowthInside1LastX128 - position.feeGrowthInside1LastX128, position.liquidity, FixedPoint128.Q128 ) - ); + )); position.feeGrowthInside0LastX128 = feeGrowthInside0LastX128; position.feeGrowthInside1LastX128 = feeGrowthInside1LastX128; @@ -193,7 +203,8 @@ contract LiquidityManager { { require(params.liquidity > 0); - TokenPosition memory position = _positions[posKey(getToken(params.token0, params.token1), tickLower, tickUpper)]; + (bool token0isWeth, address token) = getToken(params.token0, params.token1); + TokenPosition memory position = _positions[posKey(token, tickLower, tickUpper)]; uint128 positionLiquidity = position.liquidity; require(positionLiquidity >= params.liquidity); @@ -209,24 +220,21 @@ contract LiquidityManager { // this is now updated to the current transaction (, uint256 feeGrowthInside0LastX128, uint256 feeGrowthInside1LastX128, , ) = pool.positions(positionKey); - position.tokensOwed0 += - uint128(amount0) + + updateTokensOwed(token0isWeth, token, uint128(amount0) + uint128( FullMath.mulDiv( feeGrowthInside0LastX128 - position.feeGrowthInside0LastX128, positionLiquidity, FixedPoint128.Q128 ) - ); - position.tokensOwed1 += - uint128(amount1) + + ), uint128(amount1) + uint128( FullMath.mulDiv( feeGrowthInside1LastX128 - position.feeGrowthInside1LastX128, positionLiquidity, FixedPoint128.Q128 ) - ); + )); position.feeGrowthInside0LastX128 = feeGrowthInside0LastX128; position.feeGrowthInside1LastX128 = feeGrowthInside1LastX128;