feat: minor optimization & avoid recreat existed fork

axieinfinity · Dec 18, 2023 · 9d88627 · 9d88627
1 parent b143e9a
commit 9d88627
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 9 deletions.
diff --git a/script/configs/NetworkConfig.sol b/script/configs/NetworkConfig.sol
@@ -82,14 +82,13 @@ abstract contract NetworkConfig is INetworkConfig {
       console.log(StdStyle.yellow("NetworkConfig: fork mode disabled, no active fork"));
       return NULL_FORK_ID;
     }
-    if (chainId == block.chainid) {
-      console.log(
-        StdStyle.yellow(string.concat("NetworkConfig: ", chainAlias, " is already created and active at forkId:")),
-        currentFork
-      );
-      return currentFork;
-    }
+
+    if (chainId == block.chainid) return currentFork;
     if (!_isForkModeEnabled) return NULL_FORK_ID;
+    if (_networkDataMap[_networkMap[chainId]].forkId != NULL_FORK_ID) {
+      return _networkDataMap[_networkMap[chainId]].forkId;
+    }
+
     try vm.createFork(vm.rpcUrl(chainAlias)) returns (uint256 forkId) {
       console.log(StdStyle.blue(string.concat("NetworkConfig: ", chainAlias, " fork created with forkId:")), forkId);
       return forkId;

diff --git a/script/extensions/ScriptExtended.s.sol b/script/extensions/ScriptExtended.s.sol
@@ -17,13 +17,19 @@ abstract contract ScriptExtended is Script, StdAssertions, IScriptExtended {
   IGeneralConfig public constant CONFIG = IGeneralConfig(LibSharedAddress.CONFIG);
 
   modifier logFn(string memory fnName) {
-    console.log("> ", StdStyle.blue(fnName), "...");
+    _logFn(fnName);
     _;
   }
 
   modifier onlyOn(TNetwork networkType) {
-    require(network() == networkType, string.concat("ScriptExtended: Only allowed on ", CONFIG.getAlias(networkType)));
+    _requireOn(networkType);
+    _;
+  }
+
+  modifier onNetwork(TNetwork networkType) {
+    TNetwork currentNetwork = _before(networkType);
     _;
+    _after(currentNetwork);
   }
 
   function setUp() public virtual {
@@ -86,4 +92,22 @@ abstract contract ScriptExtended is Script, StdAssertions, IScriptExtended {
     assertTrue(success, "ScriptExtended: Failed to create runtime bytecode.");
     vm.etch(where, runtimeBytecode);
   }
+
+  function _logFn(string memory fnName) private view {
+    console.log("> ", StdStyle.blue(fnName), "...");
+  }
+
+  function _requireOn(TNetwork networkType) private view {
+    require(network() == networkType, string.concat("ScriptExtended: Only allowed on ", CONFIG.getAlias(networkType)));
+  }
+
+  function _before(TNetwork networkType) private returns (TNetwork currentNetwork) {
+    currentNetwork = network();
+    CONFIG.createFork(networkType);
+    CONFIG.switchTo(networkType);
+  }
+
+  function _after(TNetwork currentNetwork) private {
+    CONFIG.switchTo(currentNetwork);
+  }
 }