fix: don't repeatedly get initial values (#35)

Previously multivm had to touch DB a lot to figure out whether writes
are initial etc. This patch makes all that unnecessary.

Also, the previous version of `get_storage_changes` would have worked
incorrectly, had history been deleted. `get_storage_changes_after` does
not suffer from the same problem, as the snapshot it consumes guarantees
that history after the snapshot exists.

src/decommit.rs

@@ -78,8 +78,9 @@ impl WorldDiff {
 pub fn initial_decommit(world: &mut impl World, address: H160) -> Program {
     let deployer_system_contract_address =
         Address::from_low_u64_be(DEPLOYER_SYSTEM_CONTRACT_ADDRESS_LOW as u64);
-    let code_info =
-        world.read_storage(deployer_system_contract_address, address_into_u256(address));
+    let code_info = world
+        .read_storage(deployer_system_contract_address, address_into_u256(address))
+        .unwrap_or_default();
 
     let mut code_info_bytes = [0; 32];
     code_info.to_big_endian(&mut code_info_bytes);

src/lib.rs

@@ -36,10 +36,16 @@ pub trait World {
     fn decommit(&mut self, hash: U256) -> Program;
 
     /// There is no write_storage; [WorldDiff::get_storage_changes] gives a list of all storage changes.
-    fn read_storage(&mut self, contract: H160, key: U256) -> U256;
+    fn read_storage(&mut self, contract: H160, key: U256) -> Option<U256>;
 
     /// Computes the cost of writing a storage slot.
-    fn cost_of_writing_storage(&mut self, contract: H160, key: U256, new_value: U256) -> u32;
+    fn cost_of_writing_storage(
+        &mut self,
+        contract: H160,
+        key: U256,
+        initial_value: Option<U256>,
+        new_value: U256,
+    ) -> u32;
 
     /// Returns if the storage slot is free both in terms of gas and pubdata.
     fn is_free_storage_slot(&self, contract: &H160, key: &U256) -> bool;

src/testworld.rs

@@ -48,24 +48,27 @@ impl World for TestWorld {
         }
     }
 
-    fn read_storage(&mut self, contract: u256::H160, key: u256::U256) -> u256::U256 {
+    fn read_storage(&mut self, contract: u256::H160, key: u256::U256) -> Option<U256> {
         let deployer_system_contract_address =
             Address::from_low_u64_be(DEPLOYER_SYSTEM_CONTRACT_ADDRESS_LOW as u64);
 
         if contract == deployer_system_contract_address {
-            self.address_to_hash
-                .get(&key)
-                .copied()
-                .unwrap_or(U256::zero())
+            Some(
+                self.address_to_hash
+                    .get(&key)
+                    .copied()
+                    .unwrap_or(U256::zero()),
+            )
         } else {
-            0.into()
+            None
         }
     }
 
     fn cost_of_writing_storage(
         &mut self,
         _contract: u256::H160,
         _key: U256,
+        _initial_value: Option<U256>,
         _new_value: U256,
     ) -> u32 {
         50

src/world_diff.rs

@@ -26,6 +26,9 @@ pub struct WorldDiff {
     pub(crate) decommitted_hashes: RollbackableSet<U256>,
     read_storage_slots: RollbackableSet<(H160, U256)>,
     written_storage_slots: RollbackableSet<(H160, U256)>,
+
+    // 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 struct ExternalSnapshot {
@@ -69,7 +72,7 @@ impl WorldDiff {
             .get(&(contract, key))
             .cloned()
             .map(|v| v.1)
-            .unwrap_or_else(|| world.read_storage(contract, key));
+            .unwrap_or_else(|| world.read_storage(contract, key).unwrap_or_default());
 
         let refund = if world.is_free_storage_slot(&contract, &key)
             || self.read_storage_slots.contains(&(contract, key))
@@ -99,7 +102,12 @@ impl WorldDiff {
             return WARM_WRITE_REFUND;
         }
 
-        let update_cost = world.cost_of_writing_storage(contract, key, value);
+        let initial_value = self
+            .storage_initial_values
+            .entry((contract, key))
+            .or_insert_with(|| world.read_storage(contract, key));
+
+        let update_cost = world.cost_of_writing_storage(contract, key, *initial_value, value);
         let prepaid = self
             .paid_changes
             .insert((contract, key), update_cost)
@@ -131,19 +139,34 @@ impl WorldDiff {
         self.storage_changes.as_ref()
     }
 
-    #[allow(clippy::type_complexity)]
-    pub fn get_storage_changes(
-        &self,
-    ) -> BTreeMap<(H160, U256), (Option<(u16, U256)>, (u16, U256))> {
-        self.storage_changes.changes_after(0)
+    pub fn get_storage_changes(&self) -> BTreeMap<(H160, U256), (u16, Option<U256>, U256)> {
+        let mut result = BTreeMap::new();
+        for (key, &(tx_number, value)) in self.storage_changes.as_ref() {
+            if self.storage_initial_values[key] != Some(value) {
+                result.insert(*key, (tx_number, self.storage_initial_values[key], value));
+            }
+        }
+        result
     }
 
-    #[allow(clippy::type_complexity)]
     pub fn get_storage_changes_after(
         &self,
         snapshot: &Snapshot,
-    ) -> BTreeMap<(H160, U256), (Option<(u16, U256)>, (u16, U256))> {
-        self.storage_changes.changes_after(snapshot.storage_changes)
+    ) -> BTreeMap<(H160, U256), (u16, Option<U256>, U256)> {
+        self.storage_changes
+            .changes_after(snapshot.storage_changes)
+            .into_iter()
+            .map(|(key, (before, (tx_id, after)))| {
+                (
+                    key,
+                    (
+                        tx_id,
+                        before.map(|x| x.1).or(self.storage_initial_values[&key]),
+                        after,
+                    ),
+                )
+            })
+            .collect()
     }
 
     pub(crate) fn read_transient_storage(&mut self, contract: H160, key: U256) -> U256 {