From 9d06c2919998241fc79953c03598e17c6b3627d5 Mon Sep 17 00:00:00 2001
From: johnnyonline <johnnyonline.eth@gmail.com>
Date: Mon, 30 Sep 2024 20:52:27 +0900
Subject: [PATCH] feat: finish updateTokenStakingNodesBalances

---
 src/ynEIGEN/EigenStrategyManager.sol | 104 ++++++++++++---------------
 1 file changed, 47 insertions(+), 57 deletions(-)

diff --git a/src/ynEIGEN/EigenStrategyManager.sol b/src/ynEIGEN/EigenStrategyManager.sol
index 59cdbffc2..29e311144 100644
--- a/src/ynEIGEN/EigenStrategyManager.sol
+++ b/src/ynEIGEN/EigenStrategyManager.sol
@@ -116,10 +116,11 @@ contract EigenStrategyManager is
     IRedemptionAssetsVaultExt public redemptionAssetsVault;
     IWrapper public wrapper;
 
-    // @todo - struct
-    mapping(IStrategy => uint256) public strategiesBalance;
-    mapping(IStrategy => uint256) public strategiesWithdrawalQueueBalance;
-    mapping(IStrategy => uint256) public strategiesWithdrawnBalance;
+    struct StrategyBalance {
+        uint256 stakedBalance;
+        uint256 withdrawnBalance;
+    }
+    mapping(IStrategy => StrategyBalance) public strategiesBalance;
 
     //--------------------------------------------------------------------------------------
     //----------------------------------  INITIALIZATION  ----------------------------------
@@ -196,6 +197,46 @@ contract EigenStrategyManager is
         for (uint256 i = 0; i < assetsLength; i++) updateTokenStakingNodesBalances(assets[i], IStrategy(address(0)));
     }
 
+    //--------------------------------------------------------------------------------------
+    //------------------------------------ ACCOUNTING  ----------------------------------------
+    //--------------------------------------------------------------------------------------
+
+    /// @notice Updates the staked balances for all nodes across a strategies.
+    /// @dev Should be called atomically after any node-balance-changing operation.
+    /// @dev On a slashing events, users will have an incentive to call this function, to decrease the exchange rate.
+    /// @param asset The asset for which the balances are to be updated.
+    /// @param strategy The strategy for which the balances are to be updated. If not provided, we search for the strategy associated with the asset.
+    function updateTokenStakingNodesBalances(IERC20 asset, IStrategy strategy) public {
+
+        ITokenStakingNode[] memory nodes = tokenStakingNodesManager.getAllNodes();
+        uint256 nodesCount = nodes.length;
+
+        uint256 _strategiesBalance;
+        uint256 _strategiesWithdrawalQueueBalance;
+        uint256 _strategiesWithdrawnBalance;
+        if (address(strategy) == address(0)) strategy = strategies[asset];
+        for (uint256 i; i < nodesCount; i++ ) {
+            ITokenStakingNode node = nodes[i];
+
+            _strategiesBalance += strategy.userUnderlyingView((address(node)));
+
+            (uint256 queuedShares, uint256 strategyWithdrawnBalance) = node.getQueuedSharesAndWithdrawn(strategy, asset);
+
+            if (queuedShares > 0) {
+                _strategiesWithdrawalQueueBalance += strategy.sharesToUnderlyingView(queuedShares);
+            }
+
+            _strategiesWithdrawnBalance += strategyWithdrawnBalance;
+        }
+
+        StrategyBalance memory _strategyBalance = StrategyBalance({
+            stakedBalance: _strategiesBalance + _strategiesWithdrawalQueueBalance,
+            withdrawnBalance: _strategiesWithdrawnBalance
+        });
+
+        strategiesBalance[strategy] = _strategyBalance;
+    }
+
     //--------------------------------------------------------------------------------------
     //------------------------------------ DEPOSIT  ----------------------------------------
     //--------------------------------------------------------------------------------------
@@ -353,72 +394,21 @@ contract EigenStrategyManager is
         uint256[] memory stakedBalances = DynamicArrayLib.malloc(assets.length);
         // Add balances contained in each TokenStakingNode, including those managed by strategies.
 
-        ITokenStakingNode[] memory nodes = tokenStakingNodesManager.getAllNodes();
-        uint256 nodesCount = nodes.length;
         uint256 assetsCount = assets.length;
         for (uint256 j = 0; j < assetsCount; j++) {      
 
-            // uint256 strategiesBalance;
-            // uint256 strategiesWithdrawalQueueBalance;
-            // uint256 strategiesWithdrawnBalance;
             IERC20 asset = assets[j];
             IStrategy strategy = strategies[asset];
-            // for (uint256 i; i < nodesCount; i++ ) {
-            //     ITokenStakingNode node = nodes[i];
-
-            //     strategiesBalance += strategy.userUnderlyingView((address(node)));
-
-            //     (uint256 queuedShares, uint256 strategyWithdrawnBalance) = node.getQueuedSharesAndWithdrawn(strategy, asset);
-
-            //     if (queuedShares > 0) {
-            //         strategiesWithdrawalQueueBalance += strategy.sharesToUnderlyingView(queuedShares);
-            //     }
-
-            //     strategiesWithdrawnBalance += strategyWithdrawnBalance;
-            // }
-            // strategiesBalance
-            // strategiesWithdrawalQueueBalance
-            // strategiesWithdrawnBalance
-
             DynamicArrayLib.set(
                 stakedBalances,
                 j,
-                wrapper.toUserAssetAmount(
-                    asset,
-                    strategiesBalance[strategy] + strategiesWithdrawalQueueBalance[strategy]
-                )
-                + strategiesWithdrawnBalance[strategy]
+                wrapper.toUserAssetAmount(asset, strategiesBalance[strategy].stakedBalance)
+                + strategiesBalance[strategy].withdrawnBalance
             );
         }
 
         return stakedBalances;
     }
-    function updateTokenStakingNodesBalances(IERC20 asset, IStrategy strategy) public { // @todo - doesnt have to iterate over all nodes everytime
-
-        if (address(strategy) == address(0)) strategy = strategies[asset];
-        ITokenStakingNode[] memory nodes = tokenStakingNodesManager.getAllNodes();
-        uint256 nodesCount = nodes.length;
-        uint256 _strategiesBalance;
-        uint256 _strategiesWithdrawalQueueBalance;
-        uint256 _strategiesWithdrawnBalance;
-        for (uint256 i; i < nodesCount; i++ ) {
-            ITokenStakingNode node = nodes[i];
-
-            _strategiesBalance += strategy.userUnderlyingView((address(node)));
-
-            (uint256 queuedShares, uint256 strategyWithdrawnBalance) = node.getQueuedSharesAndWithdrawn(strategy, asset);
-
-            if (queuedShares > 0) {
-                _strategiesWithdrawalQueueBalance += strategy.sharesToUnderlyingView(queuedShares);
-            }
-
-            _strategiesWithdrawnBalance += strategyWithdrawnBalance;
-        }
-
-        strategiesBalance[strategy] = _strategiesBalance;
-        strategiesWithdrawalQueueBalance[strategy] = _strategiesWithdrawalQueueBalance;
-        strategiesWithdrawnBalance[strategy] = _strategiesWithdrawnBalance;
-    }
 
     /**
      * @notice Retrieves the total staked balance of a specific asset across all nodes.