fix: Add LiquidStone totalAssets for a given token owner

credbull · Nov 26, 2024 · c9518fb · c9518fb
1 parent 968c1f3
commit c9518fb
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/packages/contracts/src/yield/LiquidContinuousMultiTokenVault.sol b/packages/contracts/src/yield/LiquidContinuousMultiTokenVault.sol
@@ -280,6 +280,19 @@ contract LiquidContinuousMultiTokenVault is
         return totalAssets_;
     }
 
+    /// @notice Total amount of `asset` held in the vault by the `owner`
+    // @dev - this is a heavy operation as the period of the vault increases
+    function totalAssets(address owner) public view returns (uint256 totalManagedAssets) {
+        uint256 totalAssets_ = 0;
+        uint256 currentPeriod_ = currentPeriod();
+
+        for (uint256 depositPeriod = 0; depositPeriod <= currentPeriod_; ++depositPeriod) {
+            totalAssets_ += convertToAssetsForDepositPeriod(balanceOf(owner, depositPeriod), depositPeriod);
+        }
+
+        return totalAssets_;
+    }
+
     /**
      * @notice Equivalent amount of shares for the given amount of assets
      * @param assets Amount of `asset` to convert

diff --git a/packages/contracts/test/src/yield/LiquidContinuousMultiTokenVaultTest.t.sol b/packages/contracts/test/src/yield/LiquidContinuousMultiTokenVaultTest.t.sol
@@ -353,6 +353,12 @@ contract LiquidContinuousMultiTokenVaultTest is LiquidContinuousMultiTokenVaultT
         );
 
         assertEq(assetsAtRequestRedeemPeriod, _liquidVault.totalAssets(), "totalAssets wrong at requestRedeem period");
+        assertEq(
+            assetsAtRequestRedeemPeriod,
+            _liquidVault.totalAssets(alice),
+            "totalAssets(user) wrong at requestRedeem period"
+        );
+        assertEq(0, _liquidVault.totalAssets(bob), "totalAssets(user) wrong at requestRedeem period");
 
         // -------------- redeem period --------------
         _warpToPeriod(_liquidVault, redeemPeriod);
@@ -365,6 +371,8 @@ contract LiquidContinuousMultiTokenVaultTest is LiquidContinuousMultiTokenVaultT
         assertEq(assetsAtRedeemPeriod, _liquidVault.convertToAssets(totalShares), "assets wrong at redeem period");
 
         assertEq(assetsAtRedeemPeriod, _liquidVault.totalAssets(), "totalAssets wrong at redeem period");
+        assertEq(assetsAtRedeemPeriod, _liquidVault.totalAssets(alice), "totalAssets(alice) wrong at redeem period");
+        assertEq(0, _liquidVault.totalAssets(bob), "totalAssets(bob) wrong at redeem period");
     }
 
     function test__LiquidContinuousMultiTokenVault__DepositCallerValidation() public {