refactor: return address instead of IERC20 for better ux

bgd-labs · Sep 26, 2024 · 90a214a · 90a214a
1 parent c5ef87e
commit 90a214a
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 9 deletions.
diff --git a/src/contracts/extensions/static-a-token/ERC4626StataTokenUpgradeable.sol b/src/contracts/extensions/static-a-token/ERC4626StataTokenUpgradeable.sol
@@ -115,9 +115,9 @@ abstract contract ERC4626StataTokenUpgradeable is ERC4626Upgradeable, IERC4626St
   }
 
   ///@inheritdoc IERC4626StataToken
-  function aToken() public view returns (IERC20) {
+  function aToken() public view returns (address) {
     ERC4626StataTokenStorage storage $ = _getERC4626StataTokenStorage();
-    return $._aToken;
+    return address($._aToken);
   }
 
   ///@inheritdoc IERC4626

diff --git a/src/contracts/extensions/static-a-token/StataTokenV2.sol b/src/contracts/extensions/static-a-token/StataTokenV2.sol
@@ -68,10 +68,10 @@ contract StataTokenV2 is
   function maxRescue(
     address asset
   ) public view override(IRescuableBase, RescuableBase) returns (uint256) {
-    IERC20 cachedAToken = aToken();
-    if (asset == address(cachedAToken)) {
+    address cachedAToken = aToken();
+    if (asset == cachedAToken) {
       uint256 requiredBacking = _convertToAssets(totalSupply(), Math.Rounding.Ceil);
-      uint256 balance = cachedAToken.balanceOf(address(this));
+      uint256 balance = IERC20(cachedAToken).balanceOf(address(this));
       return balance > requiredBacking ? balance - requiredBacking : 0;
     }
     return type(uint256).max;

diff --git a/src/contracts/extensions/static-a-token/interfaces/IERC4626StataToken.sol b/src/contracts/extensions/static-a-token/interfaces/IERC4626StataToken.sol
@@ -1,8 +1,6 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.10;
 
-import {IERC20} from 'openzeppelin-contracts/contracts/interfaces/IERC20.sol';
-
 interface IERC4626StataToken {
   struct SignatureParams {
     uint8 v;
@@ -54,9 +52,9 @@ interface IERC4626StataToken {
 
   /**
    * @notice The aToken used inside the 4626 vault.
-   * @return IERC20 The aToken IERC20.
+   * @return address The aToken address.
    */
-  function aToken() external view returns (IERC20);
+  function aToken() external view returns (address);
 
   /**
    * @notice Returns the current asset price of the stataToken.

diff --git a/tests/extensions/static-a-token/StataTokenV2Getters.sol b/tests/extensions/static-a-token/StataTokenV2Getters.sol
@@ -5,6 +5,7 @@ import {Initializable} from 'openzeppelin-contracts-upgradeable/contracts/proxy/
 import {IERC20Metadata, IERC20} from 'openzeppelin-contracts/contracts/token/ERC20/extensions/IERC20Metadata.sol';
 import {AToken} from '../../../src/contracts/protocol/tokenization/AToken.sol';
 import {StataTokenV2} from '../../../src/contracts/extensions/static-a-token/StataTokenV2.sol'; // TODO: change import to isolate to 4626
+import {DataTypes} from '../../../src/contracts/protocol/libraries/types/DataTypes.sol';
 import {BaseTest} from './TestBase.sol';
 
 contract StataTokenV2GettersTest is BaseTest {
@@ -24,6 +25,9 @@ contract StataTokenV2GettersTest is BaseTest {
     address underlyingAddress = address(stataTokenV2.asset());
     assertEq(underlyingAddress, underlying);
 
+    DataTypes.ReserveDataLegacy memory data = contracts.poolProxy.getReserveData(underlyingAddress);
+    assertEq(stataTokenV2.aToken(), data.aTokenAddress);
+
     IERC20Metadata underlying = IERC20Metadata(underlyingAddress);
     assertEq(stataTokenV2.decimals(), underlying.decimals());