From 8083077964a1e5cd9ce8ec40ea7535a057ff1740 Mon Sep 17 00:00:00 2001
From: Vladislav Volosnikov <Volosnikov.apmath@gmail.com>
Date: Tue, 27 Aug 2024 09:32:56 +0200
Subject: [PATCH] Replace keccak with two-layers data structure in
 EvmGasManager

---
 system-contracts/contracts/EvmGasManager.sol | 91 ++++++++++++++++----
 1 file changed, 75 insertions(+), 16 deletions(-)

diff --git a/system-contracts/contracts/EvmGasManager.sol b/system-contracts/contracts/EvmGasManager.sol
index f4ea2ba21..60e94c797 100644
--- a/system-contracts/contracts/EvmGasManager.sol
+++ b/system-contracts/contracts/EvmGasManager.sol
@@ -78,16 +78,36 @@ 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)
+            }
         }
     }
 
@@ -95,29 +115,68 @@ contract EvmGasManager {
         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)
+                }
+            }
         }
     }