From 90a214a03da2dcbd8f2afaf86ca972b8cf9843b7 Mon Sep 17 00:00:00 2001
From: Lukas <lukasstrassel@googlemail.com>
Date: Thu, 26 Sep 2024 15:04:43 +0200
Subject: [PATCH] refactor: return address instead of IERC20 for better ux

---
 .../static-a-token/ERC4626StataTokenUpgradeable.sol         | 4 ++--
 src/contracts/extensions/static-a-token/StataTokenV2.sol    | 6 +++---
 .../static-a-token/interfaces/IERC4626StataToken.sol        | 6 ++----
 tests/extensions/static-a-token/StataTokenV2Getters.sol     | 4 ++++
 4 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/src/contracts/extensions/static-a-token/ERC4626StataTokenUpgradeable.sol b/src/contracts/extensions/static-a-token/ERC4626StataTokenUpgradeable.sol
index 060fc135..f61784a5 100644
--- 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
index f2af30b7..27083407 100644
--- 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
index 3cc4e9ca..d1ed6bbd 100644
--- 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
index 17d71dbf..d1ff188c 100644
--- 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());