From 7d256ec36ffa4d06fc48bf2c4741a37a0993c9cd Mon Sep 17 00:00:00 2001 From: Michael De Luca Date: Mon, 6 Jan 2025 02:42:46 -0500 Subject: [PATCH] feat: `setClaimRecipient` --- src/WrappedMToken.sol | 35 +++++--- src/interfaces/IWrappedMToken.sol | 15 +++- test/integration/UniswapV3.t.sol | 2 +- test/unit/WrappedMToken.t.sol | 125 ++++++++++++++++++++-------- test/utils/WrappedMTokenHarness.sol | 20 +++-- 5 files changed, 142 insertions(+), 55 deletions(-) diff --git a/src/WrappedMToken.sol b/src/WrappedMToken.sol index 096231c..25718a6 100644 --- a/src/WrappedMToken.sol +++ b/src/WrappedMToken.sol @@ -34,9 +34,10 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { /** * @dev Struct to represent an account's balance and yield earning details - * @param isEarning Whether the account is actively earning yield. - * @param balance The present amount of tokens held by the account. - * @param lastIndex The index of the last interaction for the account (0 for non-earning accounts). + * @param isEarning Whether the account is actively earning yield. + * @param balance The present amount of tokens held by the account. + * @param lastIndex The index of the last interaction for the account (0 for non-earning accounts). + * @param hasClaimRecipient Whether the account has an explicitly set claim recipient. */ struct Account { // First Slot @@ -44,6 +45,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { uint240 balance; // Second slot uint128 lastIndex; + bool hasClaimRecipient; } /* ============ Variables ============ */ @@ -87,6 +89,8 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { /// @dev Array of indices at which earning was enabled or disabled. uint128[] internal _enableDisableEarningIndices; + mapping(address account => address claimRecipient) internal _claimRecipients; + /* ============ Constructor ============ */ /** @@ -190,6 +194,13 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { _stopEarningFor(account_, currentIndex()); } + /// @inheritdoc IWrappedMToken + function setClaimRecipient(address claimRecipient_) external { + _accounts[msg.sender].hasClaimRecipient = (_claimRecipients[msg.sender] = claimRecipient_) != address(0); + + emit ClaimRecipientSet(msg.sender, claimRecipient_); + } + /* ============ Temporary Admin Migration ============ */ /// @inheritdoc IWrappedMToken @@ -227,13 +238,14 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { } /// @inheritdoc IWrappedMToken - function claimOverrideRecipientFor(address account_) public view returns (address recipient_) { - return - address( - uint160( - uint256(_getFromRegistrar(keccak256(abi.encode(CLAIM_OVERRIDE_RECIPIENT_KEY_PREFIX, account_)))) - ) - ); + function claimRecipientFor(address account_) public view returns (address recipient_) { + if (_accounts[account_].hasClaimRecipient) return _claimRecipients[account_]; + + address claimOverrideRecipient_ = address( + uint160(uint256(_getFromRegistrar(keccak256(abi.encode(CLAIM_OVERRIDE_RECIPIENT_KEY_PREFIX, account_))))) + ); + + return claimOverrideRecipient_ == address(0) ? account_ : claimOverrideRecipient_; } /// @inheritdoc IWrappedMToken @@ -428,8 +440,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { totalEarningSupply += yield_; } - address claimOverrideRecipient_ = claimOverrideRecipientFor(account_); - address claimRecipient_ = claimOverrideRecipient_ == address(0) ? account_ : claimOverrideRecipient_; + address claimRecipient_ = claimRecipientFor(account_); // Emit the appropriate `Claimed` and `Transfer` events, depending on the claim override recipient emit Claimed(account_, claimRecipient_, yield_); diff --git a/src/interfaces/IWrappedMToken.sol b/src/interfaces/IWrappedMToken.sol index a74ca47..66d905c 100644 --- a/src/interfaces/IWrappedMToken.sol +++ b/src/interfaces/IWrappedMToken.sol @@ -20,6 +20,13 @@ interface IWrappedMToken is IMigratable, IERC20Extended { */ event Claimed(address indexed account, address indexed recipient, uint240 yield); + /** + * @notice Emitted when `account` set their yield claim recipient. + * @param account The account that set their yield claim recipient. + * @param claimRecipient The account that will receive the yield. + */ + event ClaimRecipientSet(address indexed account, address indexed claimRecipient); + /** * @notice Emitted when Wrapped M earning is enabled. * @param index The index at the moment earning is enabled. @@ -159,6 +166,12 @@ interface IWrappedMToken is IMigratable, IERC20Extended { */ function stopEarningFor(address account) external; + /** + * @notice Explicitly sets the recipient of any yield claimed for the caller. + * @param claimRecipient The account that will receive the caller's yield. + */ + function setClaimRecipient(address claimRecipient) external; + /* ============ Temporary Admin Migration ============ */ /** @@ -207,7 +220,7 @@ interface IWrappedMToken is IMigratable, IERC20Extended { * @param account The account being queried. * @return recipient The address of the recipient, if any, to override as the destination of claimed yield. */ - function claimOverrideRecipientFor(address account) external view returns (address recipient); + function claimRecipientFor(address account) external view returns (address recipient); /// @notice The current index of Wrapped M's earning mechanism. function currentIndex() external view returns (uint128 index); diff --git a/test/integration/UniswapV3.t.sol b/test/integration/UniswapV3.t.sol index e06721d..eb79457 100644 --- a/test/integration/UniswapV3.t.sol +++ b/test/integration/UniswapV3.t.sol @@ -56,7 +56,7 @@ contract UniswapV3IntegrationTests is TestBase { _deployV2Components(); _migrate(); - _poolClaimRecipient = _wrappedMToken.claimOverrideRecipientFor(_pool); + _poolClaimRecipient = _wrappedMToken.claimRecipientFor(_pool); _wrapperBalanceOfM = _mToken.balanceOf(address(_wrappedMToken)); _poolBalanceOfUSDC = IERC20(_USDC).balanceOf(_pool); diff --git a/test/unit/WrappedMToken.t.sol b/test/unit/WrappedMToken.t.sol index 53823ab..3171581 100644 --- a/test/unit/WrappedMToken.t.sol +++ b/test/unit/WrappedMToken.t.sol @@ -149,7 +149,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken.enableEarning(); - _wrappedMToken.setAccountOf(_alice, 0, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 0, _EXP_SCALED_ONE, false); _mToken.setBalanceOf(_alice, 1_002); @@ -201,7 +201,7 @@ contract WrappedMTokenTests is Test { balance_ = uint240(bound(balance_, 0, _getMaxAmount(accountIndex_))); if (accountEarning_) { - _wrappedMToken.setAccountOf(_alice, balance_, accountIndex_); + _wrappedMToken.setAccountOf(_alice, balance_, accountIndex_, false); _wrappedMToken.setTotalEarningSupply(balance_); _wrappedMToken.setPrincipalOfTotalEarningSupply( @@ -259,7 +259,7 @@ contract WrappedMTokenTests is Test { balance_ = uint240(bound(balance_, 0, _getMaxAmount(accountIndex_))); if (accountEarning_) { - _wrappedMToken.setAccountOf(_alice, balance_, accountIndex_); + _wrappedMToken.setAccountOf(_alice, balance_, accountIndex_, false); _wrappedMToken.setTotalEarningSupply(balance_); _wrappedMToken.setPrincipalOfTotalEarningSupply( @@ -318,7 +318,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken.enableEarning(); - _wrappedMToken.setAccountOf(_alice, 999, _currentIndex); + _wrappedMToken.setAccountOf(_alice, 999, _currentIndex, false); vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 999, 1_000)); vm.prank(_alice); @@ -357,7 +357,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken.setPrincipalOfTotalEarningSupply(909); _wrappedMToken.setTotalEarningSupply(1_000); - _wrappedMToken.setAccountOf(_alice, 1_000, _currentIndex); + _wrappedMToken.setAccountOf(_alice, 1_000, _currentIndex, false); _mToken.setBalanceOf(address(_wrappedMToken), 1_000); @@ -398,7 +398,7 @@ contract WrappedMTokenTests is Test { balance_ = uint240(bound(balance_, 0, _getMaxAmount(accountIndex_))); if (accountEarning_) { - _wrappedMToken.setAccountOf(_alice, balance_, accountIndex_); + _wrappedMToken.setAccountOf(_alice, balance_, accountIndex_, false); _wrappedMToken.setTotalEarningSupply(balance_); _wrappedMToken.setPrincipalOfTotalEarningSupply( @@ -466,7 +466,7 @@ contract WrappedMTokenTests is Test { balance_ = uint240(bound(balance_, 0, _getMaxAmount(accountIndex_))); if (accountEarning_) { - _wrappedMToken.setAccountOf(_alice, balance_, accountIndex_); + _wrappedMToken.setAccountOf(_alice, balance_, accountIndex_, false); _wrappedMToken.setTotalEarningSupply(balance_); _wrappedMToken.setPrincipalOfTotalEarningSupply( @@ -517,7 +517,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken.enableEarning(); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE, false); assertEq(_wrappedMToken.balanceOf(_alice), 1_000); @@ -542,7 +542,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken.enableEarning(); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE, false); assertEq(_wrappedMToken.balanceOf(_alice), 1_000); @@ -579,7 +579,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken.setTotalEarningSupply(balance_); - _wrappedMToken.setAccountOf(_alice, balance_, accountIndex_); + _wrappedMToken.setAccountOf(_alice, balance_, accountIndex_, false); _mToken.setCurrentIndex(index_); @@ -661,7 +661,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken.enableEarning(); - _wrappedMToken.setAccountOf(_alice, 999, _currentIndex); + _wrappedMToken.setAccountOf(_alice, 999, _currentIndex, false); vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 999, 1_000)); vm.prank(_alice); @@ -728,7 +728,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken.setTotalNonEarningSupply(500); - _wrappedMToken.setAccountOf(_alice, 1_000, _currentIndex); + _wrappedMToken.setAccountOf(_alice, 1_000, _currentIndex, false); _wrappedMToken.setAccountOf(_bob, 500); vm.expectEmit(); @@ -771,7 +771,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken.setTotalNonEarningSupply(1_000); _wrappedMToken.setAccountOf(_alice, 1_000); - _wrappedMToken.setAccountOf(_bob, 500, _currentIndex); + _wrappedMToken.setAccountOf(_bob, 500, _currentIndex, false); vm.expectEmit(); emit IERC20.Transfer(_alice, _bob, 500); @@ -796,8 +796,8 @@ contract WrappedMTokenTests is Test { _wrappedMToken.setPrincipalOfTotalEarningSupply(1_363); _wrappedMToken.setTotalEarningSupply(1_500); - _wrappedMToken.setAccountOf(_alice, 1_000, _currentIndex); - _wrappedMToken.setAccountOf(_bob, 500, _currentIndex); + _wrappedMToken.setAccountOf(_alice, 1_000, _currentIndex, false); + _wrappedMToken.setAccountOf(_bob, 500, _currentIndex, false); vm.expectEmit(); emit IERC20.Transfer(_alice, _bob, 500); @@ -841,7 +841,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken.setPrincipalOfTotalEarningSupply(909); _wrappedMToken.setTotalEarningSupply(1_000); - _wrappedMToken.setAccountOf(_alice, 1_000, _currentIndex); + _wrappedMToken.setAccountOf(_alice, 1_000, _currentIndex, false); _mToken.setCurrentIndex((_currentIndex * 5) / 3); // 1_833333447838 @@ -880,7 +880,7 @@ contract WrappedMTokenTests is Test { aliceBalance_ = uint240(bound(aliceBalance_, 0, _getMaxAmount(aliceIndex_) / 4)); if (aliceEarning_) { - _wrappedMToken.setAccountOf(_alice, aliceBalance_, aliceIndex_); + _wrappedMToken.setAccountOf(_alice, aliceBalance_, aliceIndex_, false); _wrappedMToken.setTotalEarningSupply(aliceBalance_); _wrappedMToken.setPrincipalOfTotalEarningSupply( @@ -895,7 +895,7 @@ contract WrappedMTokenTests is Test { bobBalance_ = uint240(bound(bobBalance_, 0, _getMaxAmount(bobIndex_) / 4)); if (bobEarning_) { - _wrappedMToken.setAccountOf(_bob, bobBalance_, bobIndex_); + _wrappedMToken.setAccountOf(_bob, bobBalance_, bobIndex_, false); _wrappedMToken.setTotalEarningSupply(_wrappedMToken.totalEarningSupply() + bobBalance_); _wrappedMToken.setPrincipalOfTotalEarningSupply( @@ -1064,7 +1064,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken.setPrincipalOfTotalEarningSupply(909); _wrappedMToken.setTotalEarningSupply(1_000); - _wrappedMToken.setAccountOf(_alice, 999, _currentIndex); + _wrappedMToken.setAccountOf(_alice, 999, _currentIndex, false); vm.expectEmit(); emit IWrappedMToken.StoppedEarning(_alice); @@ -1089,7 +1089,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken.setTotalEarningSupply(balance_); - _wrappedMToken.setAccountOf(_alice, balance_, accountIndex_); + _wrappedMToken.setAccountOf(_alice, balance_, accountIndex_, false); _mToken.setCurrentIndex(index_); @@ -1107,6 +1107,38 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.totalEarningSupply(), 0); } + /* ============ setClaimRecipient ============ */ + function test_setClaimRecipient() external { + (, , , bool hasClaimRecipient_) = _wrappedMToken.getAccountOf(_alice); + + assertFalse(hasClaimRecipient_); + assertEq(_wrappedMToken.getInternalClaimRecipientOf(_alice), address(0)); + + vm.prank(_alice); + _wrappedMToken.setClaimRecipient(_alice); + + (, , , hasClaimRecipient_) = _wrappedMToken.getAccountOf(_alice); + + assertTrue(hasClaimRecipient_); + assertEq(_wrappedMToken.getInternalClaimRecipientOf(_alice), _alice); + + vm.prank(_alice); + _wrappedMToken.setClaimRecipient(_bob); + + (, , , hasClaimRecipient_) = _wrappedMToken.getAccountOf(_alice); + + assertTrue(hasClaimRecipient_); + assertEq(_wrappedMToken.getInternalClaimRecipientOf(_alice), _bob); + + vm.prank(_alice); + _wrappedMToken.setClaimRecipient(address(0)); + + (, , , hasClaimRecipient_) = _wrappedMToken.getAccountOf(_alice); + + assertFalse(hasClaimRecipient_); + assertEq(_wrappedMToken.getInternalClaimRecipientOf(_alice), address(0)); + } + /* ============ enableEarning ============ */ function test_enableEarning_notApprovedEarner() external { vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.NotApprovedEarner.selector, address(_wrappedMToken))); @@ -1198,11 +1230,11 @@ contract WrappedMTokenTests is Test { _wrappedMToken.enableEarning(); - _wrappedMToken.setAccountOf(_alice, 500, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 500, _EXP_SCALED_ONE, false); assertEq(_wrappedMToken.balanceOf(_alice), 500); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE, false); assertEq(_wrappedMToken.balanceOf(_alice), 1_000); @@ -1210,7 +1242,7 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.balanceOf(_alice), 1_000); - _wrappedMToken.setAccountOf(_alice, 1_000, 3 * _currentIndex); + _wrappedMToken.setAccountOf(_alice, 1_000, 3 * _currentIndex, false); assertEq(_wrappedMToken.balanceOf(_alice), 1_000); } @@ -1239,11 +1271,11 @@ contract WrappedMTokenTests is Test { _wrappedMToken.enableEarning(); - _wrappedMToken.setAccountOf(_alice, 500, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 500, _EXP_SCALED_ONE, false); assertEq(_wrappedMToken.balanceWithYieldOf(_alice), 550); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE, false); assertEq(_wrappedMToken.balanceWithYieldOf(_alice), 1_100); @@ -1251,7 +1283,7 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.balanceWithYieldOf(_alice), 2_200); - _wrappedMToken.setAccountOf(_alice, 1_000, 3 * _currentIndex); + _wrappedMToken.setAccountOf(_alice, 1_000, 3 * _currentIndex, false); assertEq(_wrappedMToken.balanceWithYieldOf(_alice), 1_000); } @@ -1280,11 +1312,11 @@ contract WrappedMTokenTests is Test { _wrappedMToken.enableEarning(); - _wrappedMToken.setAccountOf(_alice, 500, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 500, _EXP_SCALED_ONE, false); assertEq(_wrappedMToken.accruedYieldOf(_alice), 50); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE, false); assertEq(_wrappedMToken.accruedYieldOf(_alice), 100); @@ -1292,18 +1324,18 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.accruedYieldOf(_alice), 1_200); - _wrappedMToken.setAccountOf(_alice, 1_000, 3 * _currentIndex); + _wrappedMToken.setAccountOf(_alice, 1_000, 3 * _currentIndex, false); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); } /* ============ lastIndexOf ============ */ function test_lastIndexOf() external { - _wrappedMToken.setAccountOf(_alice, 0, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 0, _EXP_SCALED_ONE, false); assertEq(_wrappedMToken.lastIndexOf(_alice), _EXP_SCALED_ONE); - _wrappedMToken.setAccountOf(_alice, 0, 2 * _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 0, 2 * _EXP_SCALED_ONE, false); assertEq(_wrappedMToken.lastIndexOf(_alice), 2 * _EXP_SCALED_ONE); } @@ -1314,7 +1346,7 @@ contract WrappedMTokenTests is Test { assertFalse(_wrappedMToken.isEarning(_alice)); - _wrappedMToken.setAccountOf(_alice, 0, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 0, _EXP_SCALED_ONE, false); assertTrue(_wrappedMToken.isEarning(_alice)); } @@ -1353,16 +1385,37 @@ contract WrappedMTokenTests is Test { assertTrue(_wrappedMToken.wasEarningEnabled()); } - /* ============ claimOverrideRecipientFor ============ */ - function test_claimOverrideRecipientFor() external { - assertEq(_wrappedMToken.claimOverrideRecipientFor(_alice), address(0)); + /* ============ claimRecipientFor ============ */ + function test_claimRecipientFor() external view { + assertEq(_wrappedMToken.claimRecipientFor(_alice), _alice); + } + + function test_claimRecipientFor_hasClaimRecipient() external { + _wrappedMToken.setAccountOf(_alice, 0, 0, true); + _wrappedMToken.setInternalClaimRecipient(_alice, _bob); + + assertEq(_wrappedMToken.claimRecipientFor(_alice), _bob); + } + + function test_claimRecipientFor_hasClaimOverrideRecipient() external { + _registrar.set( + keccak256(abi.encode(_CLAIM_OVERRIDE_RECIPIENT_KEY_PREFIX, _alice)), + bytes32(uint256(uint160(_charlie))) + ); + + assertEq(_wrappedMToken.claimRecipientFor(_alice), _charlie); + } + + function test_claimRecipientFor_hasClaimRecipientAndOverrideRecipient() external { + _wrappedMToken.setAccountOf(_alice, 0, 0, true); + _wrappedMToken.setInternalClaimRecipient(_alice, _bob); _registrar.set( keccak256(abi.encode(_CLAIM_OVERRIDE_RECIPIENT_KEY_PREFIX, _alice)), bytes32(uint256(uint160(_charlie))) ); - assertEq(_wrappedMToken.claimOverrideRecipientFor(_alice), _charlie); + assertEq(_wrappedMToken.claimRecipientFor(_alice), _bob); } /* ============ totalSupply ============ */ diff --git a/test/utils/WrappedMTokenHarness.sol b/test/utils/WrappedMTokenHarness.sol index b285c83..da42999 100644 --- a/test/utils/WrappedMTokenHarness.sol +++ b/test/utils/WrappedMTokenHarness.sol @@ -20,12 +20,16 @@ contract WrappedMTokenHarness is WrappedMToken { _accounts[account_].lastIndex = uint128(index_); } - function setAccountOf(address account_, uint256 balance_, uint256 index_) external { - _accounts[account_] = Account(true, uint240(balance_), uint128(index_)); + function setAccountOf(address account_, uint256 balance_, uint256 index_, bool hasClaimRecipient_) external { + _accounts[account_] = Account(true, uint240(balance_), uint128(index_), hasClaimRecipient_); } function setAccountOf(address account_, uint256 balance_) external { - _accounts[account_] = Account(false, uint240(balance_), 0); + _accounts[account_] = Account(false, uint240(balance_), 0, false); + } + + function setInternalClaimRecipient(address account_, address claimRecipient_) external { + _claimRecipients[account_] = claimRecipient_; } function setTotalNonEarningSupply(uint256 totalNonEarningSupply_) external { @@ -40,8 +44,14 @@ contract WrappedMTokenHarness is WrappedMToken { principalOfTotalEarningSupply = uint112(principalOfTotalEarningSupply_); } - function getAccountOf(address account_) external view returns (bool isEarning_, uint240 balance_, uint128 index_) { + function getAccountOf( + address account_ + ) external view returns (bool isEarning_, uint240 balance_, uint128 index_, bool hasClaimRecipient_) { Account storage account = _accounts[account_]; - return (account.isEarning, account.balance, account.lastIndex); + return (account.isEarning, account.balance, account.lastIndex, account.hasClaimRecipient); + } + + function getInternalClaimRecipientOf(address account_) external view returns (address claimRecipient_) { + return _claimRecipients[account_]; } }