update deployed contracts

Layr-Labs · Nov 20, 2024 · e0ba0f1 · e0ba0f1
1 parent 1fc4608
commit e0ba0f1
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/src/templates/EOADeployer.sol b/src/templates/EOADeployer.sol
@@ -8,7 +8,6 @@ import {ZeusScript} from "../utils/ZeusScript.sol";
  * @notice Template for an Externally Owned Account (EOA) deploy script.
  */
 abstract contract EOADeployer is ZeusScript {
-
     Deployment[] private _deployments;
 
     /**
@@ -40,11 +39,12 @@ abstract contract EOADeployer is ZeusScript {
     function deploySingleton(address deployedTo, string memory name) internal {
         emit ZeusDeploy(name, deployedTo, true /* singleton */ );
         _deployments.push(Deployment(deployedTo, name, true));
+        updatedContracts[name] = deployedTo;
     }
 
     function deployInstance(address deployedTo, string memory name) internal {
         emit ZeusDeploy(name, deployedTo, false /* singleton */ );
-        _deployments.push(Deployment(deployedTo, name, true));
+        _deployments.push(Deployment(deployedTo, name, false));
     }
 
     function deploys() public view returns (Deployment[] memory) {

diff --git a/src/utils/ZeusScript.sol b/src/utils/ZeusScript.sol
@@ -33,6 +33,7 @@ abstract contract ZeusScript is Script {
      */
     function zeusTest() public virtual;
 
+    mapping(string => address) internal updatedContracts;
     mapping(string => EnvironmentVariableType) updatedTypes;
     mapping(string => string) updatedStrings;
     mapping(string => address) updatedAddresses;
@@ -129,18 +130,27 @@ abstract contract ZeusScript is Script {
     function zDeployedContract(string memory key) public view returns (address) {
         //                     ZEUS_DEPLOYED_ + key
         string memory envvar = addressPrefix.concat(key);
+        if (updatedContracts[envvar] != address(0)) {
+            return updatedContracts[envvar];
+        }
         return vm.envAddress(envvar);
     }
 
     function zDeployedProxy(string memory key) public view returns (address) {
         //                     ZEUS_DEPLOYED_ + key_Proxy
         string memory envvar = addressPrefix.concat(this.proxy(key));
+        if (updatedContracts[envvar] != address(0)) {
+            return updatedContracts[envvar];
+        }
         return vm.envAddress(envvar);
     }
 
     function zDeployedImpl(string memory key) public view returns (address) {
         //                     ZEUS_DEPLOYED_ + key_Impl
         string memory envvar = addressPrefix.concat(this.impl(key));
+        if (updatedContracts[envvar] != address(0)) {
+            return updatedContracts[envvar];
+        }
         return vm.envAddress(envvar);
     }