[WIP] transferring yielddistributiontokens

smart-wallets/src/token/YieldDistributionToken.sol

@@ -67,18 +67,35 @@ abstract contract YieldDistributionToken is ERC20, Ownable, IYieldDistributionTo
     struct YieldDistributionTokenStorage {
         /// @dev CurrencyToken in which the yield is deposited and denominated
         IERC20 currencyToken;
+        /// @dev Current sum of all amountSeconds for all users
+        uint256 totalAmountSeconds;
+        /// @dev Timestamp of the last change in totalSupply()
+        uint256 lastSupplyTimestamp;
+        // uint8 decimals and uint8 registeredDEXCount are packed together in the same slot.
+        // We add a uint240 __gap to fill the rest of this slot, which allows for future upgrades without changing the
+        // storage layout.
         /// @dev Number of decimals of the YieldDistributionToken
         uint8 decimals;
+        /// @dev Counter for the number of registered DEXes
+        uint8 registeredDEXCount;
+        /// @dev Padding to fill the slot
+        uint240 __gap;
+        /// @dev Registered DEX addresses
+        address[] registeredDEXes;
         /// @dev URI for the YieldDistributionToken metadata
         string tokenURI;
         /// @dev History of deposits into the YieldDistributionToken
         DepositHistory depositHistory;
-        /// @dev Current sum of all amountSeconds for all users
-        uint256 totalAmountSeconds;
-        /// @dev Timestamp of the last change in totalSupply()
-        uint256 lastSupplyTimestamp;
         /// @dev State for each user
         mapping(address user => UserState userState) userStates;
+        /// @dev Mapping to track registered DEX addresses
+        mapping(address => bool) isDEX;
+        /// @dev Mapping to track tokens in open orders for each user on each DEX
+        mapping(address user => mapping(address dex => uint256 amount)) tokensInOpenOrders;
+        /// @dev Mapping to track total tokens held on all DEXes for each user
+        mapping(address user => uint256 totalTokensOnDEXs) tokensHeldOnDEXs;
+        /// @dev Mapping to  store index of each DEX to avoid looping in unregisterDex
+        mapping(address => uint256) dexIndex;
     }
 
     // keccak256(abi.encode(uint256(keccak256("plume.storage.YieldDistributionToken")) - 1)) & ~bytes32(uint256(0xff))
@@ -120,6 +137,25 @@ abstract contract YieldDistributionToken is ERC20, Ownable, IYieldDistributionTo
      */
     event YieldAccrued(address indexed user, uint256 currencyTokenAmount);
 
+    /**
+     * @notice Emitted when yield is accrued to a user for DEX tokens
+     * @param user Address of the user who accrued the yield
+     * @param currencyTokenAmount Amount of CurrencyToken accrued as yield
+     */
+    event YieldAccruedForDEXTokens(address indexed user, uint256 currencyTokenAmount);
+
+    /**
+     * @notice Emitted when a DEX is registered
+     * @param dex Address of the user who accrued the yield
+     */
+    event DEXRegistered(address indexed dex);
+
+    /**
+     * @notice Emitted when a DEX is unregistered
+     * @param dex Address of the user who accrued the yield
+     */
+    event DEXUnregistered(address indexed dex);
+
     // Errors
 
     /**
@@ -129,6 +165,24 @@ abstract contract YieldDistributionToken is ERC20, Ownable, IYieldDistributionTo
      */
     error TransferFailed(address user, uint256 currencyTokenAmount);
 
+    /// @notice Error thrown when a non-registered DEX attempts to perform a DEX-specific operation
+    error NotRegisteredDEX();
+
+    /// @notice Error thrown when attempting to unregister more tokens than are currently in open orders for a user on a
+    /// DEX
+    error InsufficientTokensInOpenOrders();
+
+    /// @notice Error thrown when attempting to register a DEX that is already registered
+    error DEXAlreadyRegistered();
+
+    /// @notice Error thrown when attempting to unregister a DEX that is not currently registered
+    error DEXNotRegistered();
+
+    // do we want to limit maximum Dexs?
+    /// @notice Error thrown when attempting to register a new DEX when the maximum number of DEXes has already been
+    /// reached
+    //error MaximumDEXesReached();
+
     // Constructor
 
     /**
@@ -188,6 +242,11 @@ abstract contract YieldDistributionToken is ERC20, Ownable, IYieldDistributionTo
             fromState.amount = balanceOf(from);
             fromState.lastBalanceTimestamp = timestamp;
             $.userStates[from] = fromState;
+
+            // Adjust balances if transferring to a DEX
+            if ($.isDEX[to]) {
+                _adjustUserDEXBalance(from, to, value, true);
+            }
         }
 
         if (to != address(0)) {
@@ -197,6 +256,11 @@ abstract contract YieldDistributionToken is ERC20, Ownable, IYieldDistributionTo
             toState.amount = balanceOf(to);
             toState.lastBalanceTimestamp = timestamp;
             $.userStates[to] = toState;
+
+            // Adjust balances if transferring from a DEX
+            if ($.isDEX[from]) {
+                _adjustUserDEXBalance(to, from, value, false);
+            }
         }
     }
 
@@ -245,6 +309,25 @@ abstract contract YieldDistributionToken is ERC20, Ownable, IYieldDistributionTo
         emit Deposited(msg.sender, timestamp, currencyTokenAmount);
     }
 
+    /**
+     * @notice Adjusts the balance of tokens held on a DEX for a user
+     * @dev This function is called when tokens are transferred to or from a DEX
+     * @param user Address of the user whose balance is being adjusted
+     * @param dex Address of the DEX where the tokens are held
+     * @param amount Amount of tokens to adjust
+     * @param increase Boolean indicating whether to increase (true) or decrease (false) the balance
+     */
+    function _adjustUserDEXBalance(address user, address dex, uint256 amount, bool increase) internal {
+        YieldDistributionTokenStorage storage $ = _getYieldDistributionTokenStorage();
+        if (increase) {
+            $.tokensInOpenOrders[user][dex] += amount;
+            $.tokensHeldOnDEXs[user] += amount;
+        } else {
+            $.tokensInOpenOrders[user][dex] -= amount;
+            $.tokensHeldOnDEXs[user] -= amount;
+        }
+    }
+
     // Admin Setter Functions
 
     /**
@@ -256,6 +339,59 @@ abstract contract YieldDistributionToken is ERC20, Ownable, IYieldDistributionTo
         _getYieldDistributionTokenStorage().tokenURI = tokenURI;
     }
 
+    /**
+     * @notice Register a DEX address
+     * @dev Only the owner can call this function
+     * @param dexAddress Address of the DEX to register
+     */
+    function registerDEX(address dexAddress) external onlyOwner {
+        YieldDistributionTokenStorage storage $ = _getYieldDistributionTokenStorage();
+        if ($.isDEX[dexAddress]) {
+            revert DEXAlreadyRegistered();
+        }
+
+        // Should we limit registeredDexs?
+        /*
+        if ($.registeredDEXCount >= 10) {
+            revert MaximumDEXesReached();
+        }
+        */
+
+        $.isDEX[dexAddress] = true;
+        $.dexIndex[dexAddress] = $.registeredDEXes.length;
+        $.registeredDEXes.push(dexAddress);
+        $.registeredDEXCount++;
+        emit DEXRegistered(dexAddress);
+    }
+
+    /**
+     * @notice Unregister a DEX address
+     * @dev Only the owner can call this function
+     * @param dexAddress Address of the DEX to unregister
+     */
+    function unregisterDEX(address dexAddress) external onlyOwner {
+        YieldDistributionTokenStorage storage $ = _getYieldDistributionTokenStorage();
+        if (!$.isDEX[dexAddress]) {
+            revert DEXNotRegistered();
+        }
+
+        uint256 index = $.dexIndex[dexAddress];
+        uint256 lastIndex = $.registeredDEXes.length - 1;
+
+        if (index != lastIndex) {
+            address lastDex = $.registeredDEXes[lastIndex];
+            $.registeredDEXes[index] = lastDex;
+            $.dexIndex[lastDex] = index;
+        }
+
+        $.registeredDEXes.pop();
+        delete $.isDEX[dexAddress];
+        delete $.dexIndex[dexAddress];
+        $.registeredDEXCount--;
+
+        emit DEXUnregistered(dexAddress);
+    }
+
     // Getter View Functions
 
     /// @notice CurrencyToken in which the yield is deposited and denominated
@@ -268,6 +404,34 @@ abstract contract YieldDistributionToken is ERC20, Ownable, IYieldDistributionTo
         return _getYieldDistributionTokenStorage().tokenURI;
     }
 
+    /**
+     * @notice Checks if an address is a registered DEX
+     * @param addr Address to check
+     * @return bool True if the address is a registered DEX, false otherwise
+     */
+    function isDEXRegistered(address addr) external view returns (bool) {
+        return _getYieldDistributionTokenStorage().isDEX[addr];
+    }
+
+    /**
+     * @notice Gets the amount of tokens a user has in open orders on a specific DEX
+     * @param user Address of the user
+     * @param dex Address of the DEX
+     * @return uint256 Amount of tokens in open orders
+     */
+    function tokensInOpenOrdersOnDEX(address user, address dex) public view returns (uint256) {
+        return _getYieldDistributionTokenStorage().tokensInOpenOrders[user][dex];
+    }
+
+    /**
+     * @notice Gets the total amount of tokens a user has held on all DEXes
+     * @param user Address of the user
+     * @return uint256 Total amount of tokens held on all DEXes
+     */
+    function totalTokensHeldOnDEXs(address user) public view returns (uint256) {
+        return _getYieldDistributionTokenStorage().tokensHeldOnDEXs[user];
+    }
+
     // Permissionless Functions
 
     /**
@@ -309,6 +473,10 @@ abstract contract YieldDistributionToken is ERC20, Ownable, IYieldDistributionTo
         uint256 depositTimestamp = depositHistory.lastTimestamp;
         uint256 lastBalanceTimestamp = userState.lastBalanceTimestamp;
 
+        // Check if the user has any tokens on DEXs
+        uint256 tokensOnDEXs = $.tokensHeldOnDEXs[user];
+        uint256 totalUserTokens = userState.amount + tokensOnDEXs;
+
         /**
          * There is a race condition in the current implementation that occurs when
          * we deposit yield, then accrue yield for some users, then deposit more yield
@@ -341,7 +509,7 @@ abstract contract YieldDistributionToken is ERC20, Ownable, IYieldDistributionTo
                  * There can be a sequence of deposits made while the user balance remains the same throughout.
                  * Subtract the amountSeconds in this interval to get the total amountSeconds at the previous deposit.
                  */
-                uint256 intervalAmountSeconds = userState.amount * (depositTimestamp - previousDepositTimestamp);
+                uint256 intervalAmountSeconds = totalUserTokens * (depositTimestamp - previousDepositTimestamp);
                 amountSeconds -= intervalAmountSeconds;
                 yieldAccrued += _BASE * depositAmount * intervalAmountSeconds / intervalTotalAmountSeconds;
             } else {
@@ -358,11 +526,56 @@ abstract contract YieldDistributionToken is ERC20, Ownable, IYieldDistributionTo
         }
 
         userState.lastDepositAmountSeconds = userState.amountSeconds;
-        userState.amountSeconds += userState.amount * (depositHistory.lastTimestamp - lastBalanceTimestamp);
+        userState.amountSeconds += totalUserTokens * (depositHistory.lastTimestamp - lastBalanceTimestamp);
         userState.lastBalanceTimestamp = depositHistory.lastTimestamp;
-        userState.yieldAccrued += yieldAccrued / _BASE;
-        $.userStates[user] = userState;
-        emit YieldAccrued(user, yieldAccrued / _BASE);
+
+        uint256 newYield = yieldAccrued / _BASE;
+
+        if (totalUserTokens > 0) {
+            // Calculate total yield for the user, including tokens on DEXs
+            userState.yieldAccrued += newYield;
+            $.userStates[user] = userState;
+
+            // Emit an event for the total yield accrued
+            emit YieldAccrued(user, newYield);
+
+            // If user has tokens on DEXs, emit an additional event to track this
+            if (tokensOnDEXs > 0) {
+                uint256 yieldForDEXTokens = (newYield * tokensOnDEXs) / totalUserTokens;
+                emit YieldAccrued(user, yieldForDEXTokens);
+            }
+        }
+    }
+
+    /**
+     * @notice Registers an open order for a user on a DEX
+     * @dev Can only be called by a registered DEX
+     * @param user Address of the user placing the order
+     * @param amount Amount of tokens in the order
+     */
+    function registerOpenOrder(address user, uint256 amount) external {
+        YieldDistributionTokenStorage storage $ = _getYieldDistributionTokenStorage();
+        if (!$.isDEX[msg.sender]) {
+            revert NotRegisteredDEX();
+        }
+        _adjustUserDEXBalance(user, msg.sender, amount, true);
+    }
+
+    /**
+     * @notice Unregisters an open order for a user on a DEX
+     * @dev Can only be called by a registered DEX
+     * @param user Address of the user whose order is being unregistered
+     * @param amount Amount of tokens to unregister
+     */
+    function unregisterOpenOrder(address user, uint256 amount) external {
+        YieldDistributionTokenStorage storage $ = _getYieldDistributionTokenStorage();
+        if (!$.isDEX[msg.sender]) {
+            revert NotRegisteredDEX();
+        }
+        if ($.tokensInOpenOrders[user][msg.sender] < amount) {
+            revert InsufficientTokensInOpenOrders();
+        }
+        _adjustUserDEXBalance(user, msg.sender, amount, false);
     }
 
 }