Merge pull request #6 from Layr-Labs/jb/fix-state

More state refactoring

src/templates/EOADeployer.sol

@@ -58,7 +58,7 @@ 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));
-        ZEnvHelpers.state().updatedContracts[name] = deployedTo;
+        ZEnvHelpers.state().__updateContract(name, deployedTo);
     }
 
     function deployInstance(address deployedTo, string memory name) internal {
@@ -70,7 +70,7 @@ abstract contract EOADeployer is ZeusScript {
         uint256 count = state.deployedInstanceCount(name);
         string memory env = string.concat(name, "_", vm.toString(count));
 
-        state.updatedContracts[env] = deployedTo;
+        state.__updateContract(env, deployedTo);
     }
 
     function deploys() public view returns (Deployment[] memory) {

src/utils/ScriptHelpers.sol

@@ -12,4 +12,8 @@ library ScriptHelpers {
     function proxy(string memory name) internal pure returns (string memory) {
         return string.concat(name, PROXY_SUFFIX);
     }
+
+    function instance(string memory name, string memory index) internal pure returns (string memory) {
+        return string.concat(name, "_", index);
+    }
 }

src/utils/ZEnvHelpers.sol

@@ -15,6 +15,13 @@ enum EnvironmentVariableType {
     UINT_8
 }
 
+enum Cleanliness {
+    UNCHANGED, // this key has not been touched previously
+    UPTODATE, // this key has been asserted since its last change
+    DIRTY // this key has a pending unasserted change.
+
+}
+
 struct State {
     mapping(string => address) updatedContracts;
     mapping(string => EnvironmentVariableType) updatedTypes;
@@ -26,6 +33,10 @@ struct State {
     mapping(string => uint16) updatedUInt16s;
     mapping(string => uint8) updatedUInt8s;
     mapping(string => bool) updatedBools;
+
+    ////////////////////////////////////
+    mapping(string => Cleanliness) _dirty;
+    string[] _modifiedKeys;
 }
 
 library ZEnvHelpers {
@@ -47,6 +58,16 @@ library ZEnvHelpers {
         }
     }
 
+    /**
+     * NOTE: do not use this directly.
+     *
+     * please use deploySingleton / deployInstance from the EOADeployer.
+     */
+    function __updateContract(State storage s, string memory name, address deployedTo) internal {
+        __markDirty(s, name);
+        s.updatedContracts[name] = deployedTo;
+    }
+
     /**
      * @notice Returns the address of a proxy contract based on the provided key, querying the envvars injected by Zeus.
      * @param name The key to look up the address for. Should be the contract name, e.g. `type(DelegationManager).name`
@@ -212,4 +233,55 @@ library ZEnvHelpers {
         string memory envvar = string.concat(ENV_PREFIX, key);
         return vm.envBool(envvar);
     }
+
+    /**
+     * Asserts that a contract was deployed, via deploySingleton/deployInstance/deployContract.
+     */
+    function assertDeployed(State storage s, string[] memory contractNames) public onlyTest {
+        for (uint256 i = 0; i < contractNames.length; i++) {
+            clean(s, contractNames[i]);
+        }
+    }
+
+    /**
+     * Asserts that an environment variable was updated, i.e via `zUpdate*(...)`
+     */
+    function assertUpdated(State storage s, string[] memory environmentParameters) public onlyTest {
+        for (uint256 i = 0; i < environmentParameters.length; i++) {
+            clean(s, environmentParameters[i]);
+        }
+    }
+
+    /**
+     * Asserts that there are no;
+     *      - un-asserted changes to the state,
+     *      - un-asserted deployments.
+     */
+    function assertClean(State storage s) public onlyTest {
+        for (uint256 i = 0; i < s._modifiedKeys.length; i++) {
+            string memory message = string.concat(s._modifiedKeys[i], ": key was not asserted");
+            require(uint256(s._dirty[s._modifiedKeys[i]]) == uint256(Cleanliness.UPTODATE), message);
+        }
+
+        delete s._modifiedKeys;
+    }
+
+    ///////////////////////////////////////////////////// private methods
+
+    function __markDirty(State storage s, string memory key) internal {
+        if (s._dirty[key] == Cleanliness.UNCHANGED) {
+            s._modifiedKeys.push(key);
+        }
+        s._dirty[key] = Cleanliness.DIRTY;
+    }
+
+    function clean(State storage s, string memory key) private {
+        require(s._dirty[key] == Cleanliness.DIRTY, string.concat(key, ": key was unchanged."));
+        s._dirty[key] = Cleanliness.UPTODATE;
+    }
+
+    modifier onlyTest() {
+        require(vm.envBool("ZEUS_TEST"), "not a zeus test");
+        _;
+    }
 }

src/utils/ZeusScript.sol

@@ -27,28 +27,6 @@ abstract contract ZeusScript is Script, Test {
     event ZeusDeploy(string name, address addr, bool singleton);
     event ZeusMultisigExecute(address to, uint256 value, bytes data, EncGnosisSafe.Operation op);
 
-    enum Cleanliness {
-        UNCHANGED, // this key has not been touched previously
-        UPTODATE, // this key has been asserted since its last change
-        DIRTY // this key has a pending unasserted change.
-
-    }
-
-    mapping(string => Cleanliness) internal _dirty; // 1 if dirty, else 0.
-    string[] internal _modifiedKeys;
-
-    function _markDirty(string memory key) private {
-        if (_dirty[key] == Cleanliness.UNCHANGED) {
-            _modifiedKeys.push(key);
-        }
-        _dirty[key] = Cleanliness.DIRTY;
-    }
-
-    function _clean(string memory key) private {
-        require(_dirty[key] == Cleanliness.DIRTY, "Asserted key was unchanged.");
-        _dirty[key] = Cleanliness.UPTODATE;
-    }
-
     /**
      * Environment manipulation - update variables in the current environment's configuration *****
      */
@@ -64,6 +42,7 @@ abstract contract ZeusScript is Script, Test {
         );
         state.updatedTypes[key] = EnvironmentVariableType.STRING;
         state.updatedStrings[key] = key;
+        state.__markDirty(key);
         emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.STRING, abi.encode(value));
         return value;
     }
@@ -77,6 +56,7 @@ abstract contract ZeusScript is Script, Test {
         );
         state.updatedTypes[key] = EnvironmentVariableType.ADDRESS;
         state.updatedAddresses[key] = value;
+        state.__markDirty(key);
         emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.ADDRESS, abi.encode(value));
         return value;
     }
@@ -90,6 +70,7 @@ abstract contract ZeusScript is Script, Test {
         );
         state.updatedTypes[key] = EnvironmentVariableType.UINT_256;
         state.updatedUInt256s[key] = value;
+        state.__markDirty(key);
         emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.UINT_256, abi.encode(value));
         return value;
     }
@@ -104,6 +85,7 @@ abstract contract ZeusScript is Script, Test {
         state.updatedTypes[key] = EnvironmentVariableType.UINT_64;
         state.updatedUInt64s[key] = value;
         emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.UINT_64, abi.encode(value));
+        state.__markDirty(key);
         return value;
     }
 
@@ -117,6 +99,7 @@ abstract contract ZeusScript is Script, Test {
         state.updatedTypes[key] = EnvironmentVariableType.UINT_32;
         state.updatedUInt32s[key] = value;
         emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.UINT_32, abi.encode(value));
+        state.__markDirty(key);
         return value;
     }
 
@@ -130,6 +113,7 @@ abstract contract ZeusScript is Script, Test {
         state.updatedTypes[key] = EnvironmentVariableType.UINT_16;
         state.updatedUInt16s[key] = value;
         emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.UINT_16, abi.encode(value));
+        state.__markDirty(key);
         return value;
     }
 
@@ -143,6 +127,7 @@ abstract contract ZeusScript is Script, Test {
         state.updatedTypes[key] = EnvironmentVariableType.UINT_8;
         state.updatedUInt8s[key] = value;
         emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.UINT_8, abi.encode(value));
+        state.__markDirty(key);
         return value;
     }
 
@@ -156,36 +141,7 @@ abstract contract ZeusScript is Script, Test {
         state.updatedTypes[key] = EnvironmentVariableType.BOOL;
         state.updatedBools[key] = value;
         emit ZeusEnvironmentUpdate(key, EnvironmentVariableType.BOOL, abi.encode(value));
+        state.__markDirty(key);
         return value;
     }
-    ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
-
-    function zAssertDeployed(string[] memory contractNames) public {
-        zAssertTest();
-        for (uint256 i = 0; i < contractNames.length; i++) {
-            _clean(contractNames[i]);
-        }
-    }
-
-    function zAssertUpdated(string[] memory environmentParameters) public {
-        zAssertTest();
-        for (uint256 i = 0; i < environmentParameters.length; i++) {
-            _clean(environmentParameters[i]);
-        }
-    }
-
-    function zAssertClean() public {
-        zAssertTest();
-
-        for (uint256 i = 0; i < _modifiedKeys.length; i++) {
-            string memory message = string.concat(_modifiedKeys[i], ": key was not asserted");
-            require(uint256(_dirty[_modifiedKeys[i]]) == uint256(Cleanliness.UPTODATE), message);
-        }
-
-        delete _modifiedKeys;
-    }
-
-    function zAssertTest() public view {
-        require(vm.envBool("ZEUS_TEST"), "not a zeus test");
-    }
 }