feat: setClaimRecipient

m0-foundation · Jan 17, 2025 · dfb97c4 · dfb97c4
1 parent d25d55a
commit dfb97c4
Show file tree

Hide file tree

Showing 5 changed files with 142 additions and 55 deletions.
diff --git a/src/WrappedMToken.sol b/src/WrappedMToken.sol
@@ -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 ============ */
@@ -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
@@ -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
@@ -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);