Skip to content

Commit

Permalink
feat: setClaimRecipient
Browse files Browse the repository at this point in the history
  • Loading branch information
@deluca-mike
deluca-mike committed Jan 13, 2025
1 parent e635cb7 commit 869827e
Show file tree
Hide file tree
Showing 5 changed files with 135 additions and 44 deletions.
33 changes: 23 additions & 10 deletions src/WrappedMToken.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,16 +34,18 @@ 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
bool isEarning;
uint240 balance;
// Second slot
uint128 lastIndex;
bool hasClaimRecipient;
}

/* ============ Variables ============ */
Expand Down Expand Up @@ -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 ============ */

/**
Expand Down Expand Up @@ -210,6 +214,13 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
}
}

/// @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
Expand Down Expand Up @@ -247,13 +258,15 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
}

/// @inheritdoc IWrappedMToken
function claimOverrideRecipientFor(address account_) public view returns (address recipient_) {
function claimRecipientFor(address account_) public view returns (address recipient_) {
return
address(
uint160(
uint256(_getFromRegistrar(keccak256(abi.encode(CLAIM_OVERRIDE_RECIPIENT_KEY_PREFIX, account_))))
)
);
(_accounts[account_].hasClaimRecipient)
? _claimRecipients[account_]
: address(
uint160(
uint256(_getFromRegistrar(keccak256(abi.encode(CLAIM_OVERRIDE_RECIPIENT_KEY_PREFIX, account_))))
)
);
}

/// @inheritdoc IWrappedMToken
Expand Down Expand Up @@ -446,7 +459,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
totalEarningSupply += yield_;
}

address claimOverrideRecipient_ = claimOverrideRecipientFor(account_);
address claimOverrideRecipient_ = claimRecipientFor(account_);
address claimRecipient_ = claimOverrideRecipient_ == address(0) ? account_ : claimOverrideRecipient_;

// Emit the appropriate `Claimed` and `Transfer` events, depending on the claim override recipient
Expand Down
15 changes: 14 additions & 1 deletion src/interfaces/IWrappedMToken.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -174,6 +181,12 @@ interface IWrappedMToken is IMigratable, IERC20Extended {
*/
function stopEarningFor(address[] calldata accounts) 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 ============ */

/**
Expand Down Expand Up @@ -222,7 +235,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);
Expand Down
2 changes: 1 addition & 1 deletion test/integration/UniswapV3.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
109 changes: 82 additions & 27 deletions test/unit/WrappedMToken.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,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);

Expand Down Expand Up @@ -193,7 +193,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(
Expand Down Expand Up @@ -251,7 +251,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(
Expand Down Expand Up @@ -310,7 +310,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);
Expand Down Expand Up @@ -349,7 +349,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);

Expand Down Expand Up @@ -390,7 +390,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(
Expand Down Expand Up @@ -457,7 +457,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(
Expand Down Expand Up @@ -508,7 +508,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);

Expand All @@ -533,7 +533,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);

Expand Down Expand Up @@ -570,7 +570,7 @@ contract WrappedMTokenTests is Test {

_wrappedMToken.setTotalEarningSupply(balance_);

_wrappedMToken.setAccountOf(_alice, balance_, accountIndex_);
_wrappedMToken.setAccountOf(_alice, balance_, accountIndex_, false);

_mToken.setCurrentIndex(index_);

Expand Down Expand Up @@ -652,7 +652,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);
Expand Down Expand Up @@ -719,7 +719,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();
Expand Down Expand Up @@ -762,7 +762,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);
Expand All @@ -787,8 +787,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);
Expand Down Expand Up @@ -832,7 +832,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

Expand Down Expand Up @@ -871,7 +871,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(
Expand All @@ -886,7 +886,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(
Expand Down Expand Up @@ -1095,7 +1095,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);
Expand All @@ -1120,7 +1120,7 @@ contract WrappedMTokenTests is Test {

_wrappedMToken.setTotalEarningSupply(balance_);

_wrappedMToken.setAccountOf(_alice, balance_, accountIndex_);
_wrappedMToken.setAccountOf(_alice, balance_, accountIndex_, false);

_mToken.setCurrentIndex(index_);

Expand All @@ -1138,6 +1138,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));
}

/* ============ stopEarningFor batch ============ */
function test_stopEarningFor_batch_isApprovedEarner() external {
_registrar.setListContains(_EARNERS_LIST_NAME, _bob, true);
Expand All @@ -1155,8 +1187,8 @@ contract WrappedMTokenTests is Test {

_wrappedMToken.enableEarning();

_wrappedMToken.setAccountOf(_alice, 0, _currentIndex);
_wrappedMToken.setAccountOf(_bob, 0, _currentIndex);
_wrappedMToken.setAccountOf(_alice, 0, _currentIndex, false);
_wrappedMToken.setAccountOf(_bob, 0, _currentIndex, false);

address[] memory accounts_ = new address[](2);
accounts_[0] = _alice;
Expand Down Expand Up @@ -1254,7 +1286,7 @@ 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);

Expand All @@ -1267,16 +1299,39 @@ contract WrappedMTokenTests is Test {
assertEq(_wrappedMToken.balanceOf(_alice), 1_000);
}

/* ============ claimOverrideRecipientFor ============ */
function test_claimOverrideRecipientFor() external {
assertEq(_wrappedMToken.claimOverrideRecipientFor(_alice), address(0));
/* ============ claimRecipientFor ============ */
function test_claimRecipientFor_hasClaimRecipient() external {
assertEq(_wrappedMToken.claimRecipientFor(_alice), address(0));

_wrappedMToken.setAccountOf(_alice, 0, 0, true);
_wrappedMToken.setInternalClaimRecipient(_alice, _bob);

assertEq(_wrappedMToken.claimRecipientFor(_alice), _bob);
}

function test_claimRecipientFor_hasClaimOverrideRecipient() external {
assertEq(_wrappedMToken.claimRecipientFor(_alice), address(0));

_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 {
assertEq(_wrappedMToken.claimRecipientFor(_alice), address(0));

_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 ============ */
Expand Down
Loading

0 comments on commit 869827e

Please sign in to comment.