From b6c8d3c011f3f6fde181b7677263c0aa5f228e43 Mon Sep 17 00:00:00 2001
From: Joonatan Saarhelo <joon.saar@gmail.com>
Date: Sat, 31 Aug 2024 01:31:41 +0200
Subject: [PATCH] count storage cycles on the fly because order matters

---
 src/rollback.rs   |  4 ----
 src/world_diff.rs | 54 ++++++++++++++++-------------------------------
 2 files changed, 18 insertions(+), 40 deletions(-)

diff --git a/src/rollback.rs b/src/rollback.rs
index e96b2e52..16443001 100644
--- a/src/rollback.rs
+++ b/src/rollback.rs
@@ -83,10 +83,6 @@ impl<T: Ord + Clone> RollbackableSet<T> {
     pub fn contains(&self, key: &T) -> bool {
         self.map.contains_key(key)
     }
-
-    pub(crate) fn added_after(&self, snapshot: <Self as Rollback>::Snapshot) -> &[T] {
-        &self.old_entries[snapshot..]
-    }
 }
 
 impl<K: Ord> Rollback for RollbackableSet<K> {
diff --git a/src/world_diff.rs b/src/world_diff.rs
index 92df704f..49d91e33 100644
--- a/src/world_diff.rs
+++ b/src/world_diff.rs
@@ -32,6 +32,8 @@ pub struct WorldDiff {
 
     // This is never rolled back. It is just a cache to avoid asking these from DB every time.
     storage_initial_values: BTreeMap<(H160, U256), Option<U256>>,
+
+    pub storage_application_cycles: u32,
 }
 
 #[derive(Debug)]
@@ -103,9 +105,12 @@ impl WorldDiff {
             .copied()
             .unwrap_or_else(|| world.read_storage(contract, key).unwrap_or_default());
 
-        let refund = if self.read_storage_slots.add((contract, key))
-            || world.is_free_storage_slot(&contract, &key)
-        {
+        let already_accessed = self.read_storage_slots.add((contract, key));
+        if !already_accessed {
+            self.storage_application_cycles += STORAGE_READ_STORAGE_APPLICATION_CYCLES;
+        }
+
+        let refund = if already_accessed || world.is_free_storage_slot(&contract, &key) {
             WARM_READ_REFUND
         } else {
             0
@@ -130,7 +135,9 @@ impl WorldDiff {
             .or_insert_with(|| world.read_storage(contract, key));
 
         if world.is_free_storage_slot(&contract, &key) {
-            self.written_storage_slots.add((contract, key));
+            if !self.written_storage_slots.add((contract, key)) {
+                self.storage_application_cycles += STORAGE_WRITE_STORAGE_APPLICATION_CYCLES;
+            }
             self.read_storage_slots.add((contract, key));
 
             self.storage_refunds.push(WARM_WRITE_REFUND);
@@ -146,10 +153,14 @@ impl WorldDiff {
 
         let refund = if self.written_storage_slots.add((contract, key)) {
             WARM_WRITE_REFUND
-        } else if self.read_storage_slots.add((contract, key)) {
-            COLD_WRITE_AFTER_WARM_READ_REFUND
         } else {
-            0
+            self.storage_application_cycles += STORAGE_WRITE_STORAGE_APPLICATION_CYCLES;
+
+            if self.read_storage_slots.add((contract, key)) {
+                COLD_WRITE_AFTER_WARM_READ_REFUND
+            } else {
+                0
+            }
         };
 
         let pubdata_cost = (update_cost as i32) - (prepaid as i32);
@@ -331,35 +342,6 @@ impl WorldDiff {
     pub(crate) fn clear_transient_storage(&mut self) {
         self.transient_storage_changes = Default::default();
     }
-
-    pub fn storage_application_cycles_snapshot(&self) -> StorageApplicationCyclesSnapshot {
-        StorageApplicationCyclesSnapshot {
-            written_storage_slots: self.written_storage_slots.snapshot(),
-            read_storage_slots: self.read_storage_slots.snapshot(),
-        }
-    }
-
-    pub fn storage_application_cycles_after(
-        &self,
-        snapshot: &StorageApplicationCyclesSnapshot,
-    ) -> usize {
-        self.written_storage_slots
-            .added_after(snapshot.written_storage_slots)
-            .len()
-            * STORAGE_WRITE_STORAGE_APPLICATION_CYCLES as usize
-            + self
-                .read_storage_slots
-                .added_after(snapshot.read_storage_slots)
-                .iter()
-                .filter(|slot| !self.written_storage_slots.contains(slot))
-                .count()
-                * STORAGE_READ_STORAGE_APPLICATION_CYCLES as usize
-    }
-}
-
-pub struct StorageApplicationCyclesSnapshot {
-    written_storage_slots: <RollbackableSet<(H160, U256)> as Rollback>::Snapshot,
-    read_storage_slots: <RollbackableSet<(H160, U256)> as Rollback>::Snapshot,
 }
 
 #[derive(Clone, PartialEq, Debug)]