From a7b6bce53006b4144a62c271b643a6491d8005df Mon Sep 17 00:00:00 2001
From: Schlagonia <schlagonia@gmail.com>
Date: Thu, 10 Oct 2024 13:44:15 -0600
Subject: [PATCH] feat: add recplacement

---
 .../ProtocolAddressProvider.vy                | 20 +++++++++++++++
 src/interfaces/IProtocolAddressProvider.sol   |  4 +++
 src/test/TestAddressProvider.t.sol            | 25 +++++++++++++++++++
 3 files changed, 49 insertions(+)

diff --git a/src/addressProviders/ProtocolAddressProvider.vy b/src/addressProviders/ProtocolAddressProvider.vy
index 83ad336..d1305da 100644
--- a/src/addressProviders/ProtocolAddressProvider.vy
+++ b/src/addressProviders/ProtocolAddressProvider.vy
@@ -27,6 +27,7 @@ event UpdateGovernance:
 ROUTER: constant(bytes32) = keccak256("Router")
 KEEPER: constant(bytes32) = keccak256("Keeper")
 APR_ORACLE: constant(bytes32) = keccak256("APR Oracle")
+REPLACEMENT: constant(bytes32) = keccak256("Replacement")
 RELEASE_REGISTRY: constant(bytes32) = keccak256("Release Registry")
 BASE_FEE_PROVIDER: constant(bytes32) = keccak256("Base Fee Provider")
 COMMON_REPORT_TRIGGER: constant(bytes32) = keccak256("Common Report Trigger")
@@ -79,6 +80,16 @@ def getAddress(address_id: bytes32) -> address:
 def _get_address(address_id: bytes32) -> address:
     return self.addresses[address_id]
 
+@view
+@external
+def getReplacement() -> address:
+    """
+    @notice Get the replacement of the address provider if it exists.
+        This will return address 0 if this is the latest version.
+    @return The replacement address provider if any.
+    """
+    return self._get_address(REPLACEMENT)
+
 @view
 @external
 def getRouter() -> address:
@@ -214,6 +225,15 @@ def _set_address(address_id: bytes32, new_address: address):
 
     log UpdatedAddress(address_id, old_address, new_address)
 
+@external
+def setReplacement(new_address: address):
+    """
+    @notice Set the replacement address provider.
+    @dev Must be called by the governance.
+    @param new_address of the replacement address provider.
+    """
+    self._set_address(REPLACEMENT, new_address)
+
 @external
 def setRouter(new_address: address):
     """
diff --git a/src/interfaces/IProtocolAddressProvider.sol b/src/interfaces/IProtocolAddressProvider.sol
index 303feb1..575c6f0 100644
--- a/src/interfaces/IProtocolAddressProvider.sol
+++ b/src/interfaces/IProtocolAddressProvider.sol
@@ -18,6 +18,8 @@ interface IProtocolAddressProvider {
 
     function getAddress(bytes32 _id) external view returns (address);
 
+    function setReplacement(address _replacement) external;
+
     function setRouter(address _router) external;
 
     function setKeeper(address _keeper) external;
@@ -42,6 +44,8 @@ interface IProtocolAddressProvider {
 
     function setRoleManagerFactory(address _roleManagerFactory) external;
 
+    function getReplacement() external view returns (address);
+
     function getRouter() external view returns (address);
 
     function getKeeper() external view returns (address);
diff --git a/src/test/TestAddressProvider.t.sol b/src/test/TestAddressProvider.t.sol
index 9eadf98..0602ce9 100644
--- a/src/test/TestAddressProvider.t.sol
+++ b/src/test/TestAddressProvider.t.sol
@@ -12,6 +12,7 @@ contract TestAddressProvider is Setup {
         assertEq(addressProvider.name(), "Yearn V3 Protocol Address Provider");
         assertEq(addressProvider.governance(), daddy);
         assertEq(addressProvider.pendingGovernance(), ZERO_ADDRESS);
+        assertEq(addressProvider.getReplacement(), ZERO_ADDRESS);
         assertEq(addressProvider.getRouter(), ZERO_ADDRESS);
         assertEq(addressProvider.getKeeper(), ZERO_ADDRESS);
         assertEq(addressProvider.getAprOracle(), ZERO_ADDRESS);
@@ -46,6 +47,29 @@ contract TestAddressProvider is Setup {
         assertEq(addressProvider.getAddress(id), newAddress);
     }
 
+    function test__set_replacement() public {
+        bytes32 id = AddressIds.REPLACEMENT;
+        address newAddress = address(0x456);
+
+        assertEq(addressProvider.getAddress(id), ZERO_ADDRESS);
+        assertEq(addressProvider.getReplacement(), ZERO_ADDRESS);
+
+        vm.prank(user);
+        vm.expectRevert("!governance");
+        addressProvider.setReplacement(newAddress);
+
+        assertEq(addressProvider.getAddress(id), ZERO_ADDRESS);
+        assertEq(addressProvider.getReplacement(), ZERO_ADDRESS);
+
+        vm.prank(daddy);
+        vm.expectEmit(true, true, true, true, address(addressProvider));
+        emit UpdatedAddress(id, ZERO_ADDRESS, newAddress);
+        addressProvider.setReplacement(newAddress);
+
+        assertEq(addressProvider.getAddress(id), newAddress);
+        assertEq(addressProvider.getReplacement(), newAddress);
+    }
+
     function test__set_router() public {
         bytes32 id = AddressIds.ROUTER;
         address newAddress = address(0x123);
@@ -407,6 +431,7 @@ library AddressIds {
     bytes32 constant ROUTER = keccak256("Router");
     bytes32 constant KEEPER = keccak256("Keeper");
     bytes32 constant APR_ORACLE = keccak256("APR Oracle");
+    bytes32 constant REPLACEMENT = keccak256("Replacement");
     bytes32 constant RELEASE_REGISTRY = keccak256("Release Registry");
     bytes32 constant BASE_FEE_PROVIDER = keccak256("Base Fee Provider");
     bytes32 constant COMMON_REPORT_TRIGGER = keccak256("Common Report Trigger");