Replace keccak with two-layers data structure in EvmGasManager

matter-labs · Aug 27, 2024 · 8083077 · 8083077
1 parent 1c5c124
commit 8083077
Showing 1 changed file with 75 additions and 16 deletions.
diff --git a/system-contracts/contracts/EvmGasManager.sol b/system-contracts/contracts/EvmGasManager.sol
@@ -78,46 +78,105 @@ contract EvmGasManager {
     }
 
     function isSlotWarm(uint256 _slot) external view returns (bool isWarm) {
-        uint256 prefix = IS_SLOT_WARM_PREFIX | uint256(uint160(msg.sender));
+        address prefix = msg.sender;
         uint256 transient_slot;
+
+        /*if (_slot > 2**96 - 1) {
+            assembly {
+                mstore(0, prefix)
+                mstore(0x20, _slot)
+                transient_slot := keccak256(0, 64)
+            }
+        } else {
+            assembly {
+                transient_slot := or(shl(160, _slot), prefix)
+            }
+        }*/
+
+        uint256 transient_slot_1;
+        uint256 mask = (1 << 32 - 1) << 224;
+
+        uint256 counter;
         assembly {
-            mstore(0, prefix)
-            mstore(0x20, _slot)
-            transient_slot := keccak256(0, 64)
+            transient_slot_1 := or(prefix, shr(32, and(mask, _slot)))
+            counter := tload(transient_slot_1)
         }
 
-        assembly {
-            isWarm := tload(transient_slot)
+        if (counter != 0) {
+            assembly {
+                transient_slot := and(not(mask), _slot)
+                transient_slot := or(transient_slot, shl(224, counter))
+                isWarm := tload(transient_slot)
+            }
         }
     }
 
     function warmSlot(
         uint256 _slot,
         uint256 _currentValue
     ) external payable onlySystemEvm returns (bool isWarm, uint256 originalValue) {
-        uint256 prefix = IS_SLOT_WARM_PREFIX | uint256(uint160(msg.sender));
+        address prefix = msg.sender;
         uint256 transient_slot;
-        assembly {
-            mstore(0, prefix)
-            mstore(0x20, _slot)
-            transient_slot := keccak256(0, 64)
-        }
+        /*if (_slot > 2**96 - 1) {
+            assembly {
+                mstore(0, prefix)
+                mstore(0x20, _slot)
+                transient_slot := keccak256(0, 64)
+            }
+        } else {
+            assembly {
+                transient_slot := or(shl(160, _slot), prefix)
+            }
+        }*/
 
+        uint256 transient_slot_1;
+        uint256 mask = (1 << 32 - 1) << 224;
+
+        uint256 counter;
         assembly {
-            isWarm := tload(transient_slot)
+            transient_slot_1 := or(prefix, shr(32, and(mask, _slot)))
+            counter := tload(transient_slot_1)
         }
 
-        if (isWarm) {
+        if (counter == 0) {
+            // isWarm = false
             assembly {
-                originalValue := tload(add(transient_slot, 1))
+                counter := add(tload(IS_SLOT_WARM_PREFIX), 1)
+                tstore(IS_SLOT_WARM_PREFIX, counter)
+                tstore(transient_slot_1, counter)
             }
-        } else {
+
+            assembly {
+                transient_slot := and(not(mask), _slot)
+                transient_slot := or(transient_slot, shl(224, counter))
+            }
+
             originalValue = _currentValue;
 
             assembly {
                 tstore(transient_slot, 1)
                 tstore(add(transient_slot, 1), originalValue)
             }
+
+        } else {
+            assembly {
+                transient_slot := and(not(mask), _slot)
+                transient_slot := or(transient_slot, shl(224, counter))
+                isWarm := tload(transient_slot)
+            }
+
+            if (isWarm) {
+                assembly {
+                    originalValue := tload(add(transient_slot, 1))
+                }
+            } else {
+                originalValue = _currentValue;
+
+                assembly {
+                    tstore(transient_slot, 1)
+                    tstore(add(transient_slot, 1), originalValue)
+                }
+            }
         }
     }