[EVM-Equivalence-YUL] Optimize EvmGasManager

system-contracts/contracts/EvmGasManager.sol

@@ -2,6 +2,7 @@
 pragma solidity ^0.8.0;
 
 import "./EvmConstants.sol";
+import "./libraries/Utils.sol";
 
 import {ACCOUNT_CODE_STORAGE_SYSTEM_CONTRACT} from "./Constants.sol";
 import {ISystemContract} from "./interfaces/ISystemContract.sol";
@@ -13,12 +14,11 @@ uint160 constant PRECOMPILES_END = 0xffff;
 // Denotes that passGas has been consumed
 uint256 constant INF_PASS_GAS = 0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff;
 
-contract EvmGasManager {
-    // We need strust to use `storage` pointers
-    struct WarmAccountInfo {
-        bool isWarm;
-    }
+uint256 constant IS_ACCOUNT_EVM_PREFIX = 1 << 255;
+uint256 constant IS_ACCOUNT_WARM_PREFIX = 1 << 254;
+uint256 constant IS_SLOT_WARM_PREFIX = 1 << 253;
 
+contract EvmGasManager {
     struct SlotInfo {
         bool warm;
         uint256 originalValue;
@@ -32,56 +32,32 @@ contract EvmGasManager {
 
     // The following storage variables are not used anywhere explicitly and are just used to obtain the storage pointers
     // to use the transient storage with.
-    mapping(address => WarmAccountInfo) private warmAccounts;
+    mapping(address => bool) private warmAccounts;
     mapping(address => mapping(uint256 => SlotInfo)) private warmSlots;
     EVMStackFrameInfo[] private evmStackFrames;
 
-    function tstoreWarmAccount(address account, bool isWarm) internal {
-        WarmAccountInfo storage ptr = warmAccounts[account];
-
-        assembly {
-            tstore(ptr.slot, isWarm)
-        }
-    }
-
-    function tloadWarmAccount(address account) internal returns (bool isWarm) {
-        WarmAccountInfo storage ptr = warmAccounts[account];
-
-        assembly {
-            isWarm := tload(ptr.slot)
-        }
-    }
-
-    function tstoreWarmSlot(address _account, uint256 _key, SlotInfo memory info) internal {
-        SlotInfo storage ptr = warmSlots[_account][_key];
-
-        bool warm = info.warm;
-        uint256 originalValue = info.originalValue;
-
+    modifier onlySystemEvm() {
+        // cache use is safe since we do not support SELFDESTRUCT
+        uint256 transient_slot = IS_ACCOUNT_EVM_PREFIX | uint256(uint160(msg.sender));
+        bool isEVM;
         assembly {
-            tstore(ptr.slot, warm)
-            tstore(add(ptr.slot, 1), originalValue)
+            isEVM := tload(transient_slot)
         }
-    }
-
-    function tloadWarmSlot(address _account, uint256 _key) internal view returns (SlotInfo memory info) {
-        SlotInfo storage ptr = warmSlots[_account][_key];
-
-        bool isWarm;
-        uint256 originalValue;
 
-        assembly {
-            isWarm := tload(ptr.slot)
-            originalValue := tload(add(ptr.slot, 1))
+        if (!isEVM) {
+            bytes32 bytecodeHash = ACCOUNT_CODE_STORAGE_SYSTEM_CONTRACT.getRawCodeHash(msg.sender);
+            isEVM = Utils.isCodeHashEVM(bytecodeHash);
+            if (isEVM) {
+                if (!Utils.isContractConstructing(bytecodeHash)) {
+                    assembly {
+                        tstore(transient_slot, isEVM)
+                    }
+                }
+            }
         }
 
-        info.warm = isWarm;
-        info.originalValue = originalValue;
-    }
-
-    modifier onlySystemEvm() {
-        require(ACCOUNT_CODE_STORAGE_SYSTEM_CONTRACT.isAccountEVM(msg.sender), "only system evm");
-        require(SystemContractHelper.isSystemCall(), "This method require system call flag");
+        require(isEVM, "only system evm");
+        require(SystemContractHelper.isSystemCall(), "This method requires system call flag");
         _;
     }
 
@@ -91,27 +67,61 @@ contract EvmGasManager {
     function warmAccount(address account) external payable onlySystemEvm returns (bool wasWarm) {
         if (uint160(account) < PRECOMPILES_END) return true;
 
-        wasWarm = tloadWarmAccount(account);
-        if (!wasWarm) tstoreWarmAccount(account, true);
-    }
+        uint256 transient_slot = IS_ACCOUNT_WARM_PREFIX | uint256(uint160(account));
 
-    function isSlotWarm(uint256 _slot) external view returns (bool) {
-        return tloadWarmSlot(msg.sender, _slot).warm;
+        assembly {
+            wasWarm := tload(transient_slot)
+        }
+
+        if (!wasWarm) {
+            assembly {
+                tstore(transient_slot, 1)
+            }
+        }
     }
 
-    function warmSlot(uint256 _slot, uint256 _currentValue) external payable onlySystemEvm returns (bool, uint256) {
-        SlotInfo memory info = tloadWarmSlot(msg.sender, _slot);
+    function isSlotWarm(uint256 _slot) external view returns (bool isWarm) {
+        uint256 prefix = IS_SLOT_WARM_PREFIX | uint256(uint160(msg.sender));
+        uint256 transient_slot;
+        assembly {
+            mstore(0, prefix)
+            mstore(0x20, _slot)
+            transient_slot := keccak256(0, 64)
+        }
 
-        if (info.warm) {
-            return (true, info.originalValue);
+        assembly {
+            isWarm := tload(transient_slot)
         }
+    }
 
-        info.warm = true;
-        info.originalValue = _currentValue;
+    function warmSlot(
+        uint256 _slot,
+        uint256 _currentValue
+    ) external payable onlySystemEvm returns (bool isWarm, uint256 originalValue) {
+        uint256 prefix = IS_SLOT_WARM_PREFIX | uint256(uint160(msg.sender));
+        uint256 transient_slot;
+        assembly {
+            mstore(0, prefix)
+            mstore(0x20, _slot)
+            transient_slot := keccak256(0, 64)
+        }
 
-        tstoreWarmSlot(msg.sender, _slot, info);
+        assembly {
+            isWarm := tload(transient_slot)
+        }
 
-        return (false, _currentValue);
+        if (isWarm) {
+            assembly {
+                originalValue := tload(add(transient_slot, 1))
+            }
+        } else {
+            originalValue = _currentValue;
+
+            assembly {
+                tstore(transient_slot, 1)
+                tstore(add(transient_slot, 1), originalValue)
+            }
+        }
     }
 
     /*
@@ -135,13 +145,13 @@ contract EvmGasManager {
     function consumeEvmFrame() external onlySystemEvm returns (uint256 passGas, bool isStatic) {
         if (evmStackFrames.length == 0) return (INF_PASS_GAS, false);
 
-        EVMStackFrameInfo memory frameInfo = evmStackFrames[evmStackFrames.length - 1];
+        EVMStackFrameInfo storage frameInfo = evmStackFrames[evmStackFrames.length - 1];
 
         passGas = frameInfo.passGas;
         isStatic = frameInfo.isStatic;
 
         // Mark as used
-        evmStackFrames[evmStackFrames.length - 1].passGas = INF_PASS_GAS;
+        frameInfo.passGas = INF_PASS_GAS;
     }
 
     function popEVMFrame() external onlySystemEvm {