Quick refactor

dl-solarity · Oct 25, 2023 · 1f214f1 · 1f214f1
1 parent c3db10e
commit 1f214f1
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/contracts/finance/Vesting.sol b/contracts/finance/Vesting.sol
@@ -118,11 +118,13 @@ abstract contract VestingWallet is Initializable {
     }
 
     function totalAllocation() public view virtual returns (uint256) {
-        return address(this).balance + totalReleased();
+        return totalAllocation(ETH);
     }
 
     function totalAllocation(address token_) public view virtual returns (uint256) {
-        return IERC20(token_).balanceOf(address(this)) + totalReleased(token_);
+        return
+            (token_ == ETH ? address(this).balance : IERC20(token_).balanceOf(address(this))) +
+            totalReleased(token_);
     }
 
     function totalReleased() public view virtual returns (uint256) {