update zUpdate to be optimistic

Layr-Labs · Nov 13, 2024 · cc9e865 · cc9e865
1 parent 06be0c7
commit cc9e865
Showing 1 changed file with 93 additions and 6 deletions.
diff --git a/src/utils/ZeusScript.sol b/src/utils/ZeusScript.sol
@@ -8,9 +8,13 @@ abstract contract ZeusScript is Script {
     using StringUtils for string;
 
     enum EnvironmentVariableType {
-        INT_256,
+        UNMODIFIED,
+        UINT_256,
+        UINT_32,
+        UINT_64,
         ADDRESS,
-        STRING
+        STRING,
+        BOOL
     }
 
     event ZeusEnvironmentUpdate(string key, EnvironmentVariableType internalType, bytes value);
@@ -23,22 +27,85 @@ abstract contract ZeusScript is Script {
      */
     function zeusTest() public virtual;
 
+    mapping(string => EnvironmentVariableType) updatedTypes;
+    mapping(string => string) updatedStrings;
+    mapping(string => address) updatedAddresses;
+    mapping(string => uint256) updatedUInt256s;
+    mapping(string => uint64) updatedUInt64s;
+    mapping(string => uint32) updatedUInt32s;
+    mapping(string => bool) updatedBools;
+
     /**
      * Environment manipulation - update variables in the current environment's configuration *****
      */
     // NOTE: you do not need to use these for contract addresses, which are tracked and injected automatically.
     // NOTE: do not use `.update()` during a vm.broadcast() segment.
-    function update(string memory key, string memory value) public {
+    ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+    function zUpdate(string memory key, string memory value) public returns (string memory) {
+        require(
+            updatedTypes[key] == EnvironmentVariableType.UNMODIFIED
+                || updatedTypes[key] == EnvironmentVariableType.STRING
+        );
+        updatedTypes[key] = EnvironmentVariableType.STRING;
+        updatedStrings[key] = key;
         emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.STRING, abi.encode(value));
+        return value;
     }
 
-    function update(string memory key, address value) public {
+    function zUpdate(string memory key, address value) public returns (address) {
+        require(
+            updatedTypes[key] == EnvironmentVariableType.UNMODIFIED
+                || updatedTypes[key] == EnvironmentVariableType.ADDRESS
+        );
+        updatedTypes[key] = EnvironmentVariableType.ADDRESS;
+        updatedAddresses[key] = value;
         emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.ADDRESS, abi.encode(value));
+        return value;
+    }
+
+    function zUpdate(string memory key, uint256 value) public returns (uint256) {
+        require(
+            updatedTypes[key] == EnvironmentVariableType.UNMODIFIED
+                || updatedTypes[key] == EnvironmentVariableType.UINT_256
+        );
+        updatedTypes[key] = EnvironmentVariableType.UINT_256;
+        updatedUInt256s[key] = value;
+        emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.UINT_256, abi.encode(value));
+        return value;
+    }
+
+    function zUpdate(string memory key, uint64 value) public returns (uint64) {
+        require(
+            updatedTypes[key] == EnvironmentVariableType.UNMODIFIED
+                || updatedTypes[key] == EnvironmentVariableType.UINT_64
+        );
+        updatedTypes[key] = EnvironmentVariableType.UINT_64;
+        updatedUInt64s[key] = value;
+        emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.UINT_64, abi.encode(value));
+        return value;
+    }
+
+    function zUpdate(string memory key, uint32 value) public returns (uint32) {
+        require(
+            updatedTypes[key] == EnvironmentVariableType.UNMODIFIED
+                || updatedTypes[key] == EnvironmentVariableType.UINT_32
+        );
+        updatedTypes[key] = EnvironmentVariableType.UINT_32;
+        updatedUInt32s[key] = value;
+        emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.UINT_32, abi.encode(value));
+        return value;
     }
 
-    function update(string memory key, uint256 value) public {
-        emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.INT_256, abi.encode(value));
+    function zUpdate(string memory key, bool value) public returns (bool) {
+        require(
+            updatedTypes[key] == EnvironmentVariableType.UNMODIFIED || updatedTypes[key] == EnvironmentVariableType.BOOL
+        );
+        updatedTypes[key] = EnvironmentVariableType.BOOL;
+        updatedBools[key] = value;
+        emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.BOOL, abi.encode(value));
+        return value;
     }
+    ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
     /**
      * @notice Returns the address of a contract based on the provided key, querying the envvars injected by Zeus. This is typically the name of the contract.
@@ -56,6 +123,10 @@ abstract contract ZeusScript is Script {
      * @param key The environment key. Corresponds to a ZEUS_* env variable.
      */
     function zAddress(string memory key) public view returns (address) {
+        if (updatedTypes[key] != EnvironmentVariableType.UNMODIFIED) {
+            return updatedAddresses[key];
+        }
+
         return vm.envAddress(key);
     }
 
@@ -64,6 +135,10 @@ abstract contract ZeusScript is Script {
      * @param key The environment key. Corresponds to a ZEUS_* env variable.
      */
     function zUint32(string memory key) public view returns (uint32) {
+        if (updatedTypes[key] != EnvironmentVariableType.UNMODIFIED) {
+            return updatedUInt32s[key];
+        }
+
         string memory envvar = envPrefix.concat(key);
         return uint32(vm.envUint(envvar));
     }
@@ -73,6 +148,10 @@ abstract contract ZeusScript is Script {
      * @param key The environment key. Corresponds to a ZEUS_* env variable.
      */
     function zUint64(string memory key) public view returns (uint64) {
+        if (updatedTypes[key] != EnvironmentVariableType.UNMODIFIED) {
+            return updatedUInt64s[key];
+        }
+
         string memory envvar = envPrefix.concat(key);
         return uint64(vm.envUint(envvar));
     }
@@ -82,6 +161,10 @@ abstract contract ZeusScript is Script {
      * @param key The environment key. Corresponds to a ZEUS_* env variable.
      */
     function zBool(string memory key) public view returns (bool) {
+        if (updatedTypes[key] != EnvironmentVariableType.UNMODIFIED) {
+            return updatedBools[key];
+        }
+
         string memory envvar = envPrefix.concat(key);
         return bool(vm.envBool(envvar));
     }
@@ -91,6 +174,10 @@ abstract contract ZeusScript is Script {
      * @param key The environment key. Corresponds to a ZEUS_* env variable.
      */
     function zString(string memory key) public view returns (string memory) {
+        if (updatedTypes[key] != EnvironmentVariableType.UNMODIFIED) {
+            return updatedStrings[key];
+        }
+
         string memory envvar = envPrefix.concat(key);
         return vm.envString(envvar);
     }