feat: track pubdata price (#14)

Pubdata is tallied. Storage slots are considered both warm and free in terms of pubdata if they are one of the free slots. --------- Co-authored-by: Joonatan Saarhelo <[email protected]>
matter-labs · May 8, 2024 · 5d4d1c0 · 5d4d1c0
1 parent 69947e4
commit 5d4d1c0
Show file tree

Hide file tree

Showing 8 changed files with 78 additions and 32 deletions.
diff --git a/src/callframe.rs b/src/callframe.rs
@@ -20,6 +20,8 @@ pub struct Callframe {
     pub gas: u32,
     pub stipend: u32,
 
+    pub total_pubdata_spent: i32,
+
     near_calls: Vec<NearCallFrame>,
 
     pub(crate) program: Program,
@@ -83,6 +85,7 @@ impl Callframe {
             exception_handler,
             near_calls: vec![],
             world_before_this_frame,
+            total_pubdata_spent: 0,
         }
     }
 
@@ -112,6 +115,7 @@ impl Callframe {
                 program_counter: f.call_instruction,
                 exception_handler: f.exception_handler,
                 snapshot: f.world_before_this_frame,
+                total_pubdata_spent: 0,
             }
         })
     }
@@ -142,4 +146,5 @@ pub(crate) struct FrameRemnant {
     pub(crate) program_counter: u16,
     pub(crate) exception_handler: u16,
     pub(crate) snapshot: Snapshot,
+    pub(crate) total_pubdata_spent: i32,
 }
diff --git a/src/instruction_handlers/ret.rs b/src/instruction_handlers/ret.rs
@@ -40,11 +40,13 @@ fn ret<const RETURN_TYPE: u8, const TO_LABEL: bool>(
     let mut return_type = ReturnType::from_u8(RETURN_TYPE);
     let near_call_leftover_gas = vm.state.current_frame.gas;
 
-    let (pc, snapshot, leftover_gas) = if let Some(FrameRemnant {
+    let (pc, snapshot, leftover_gas, total_pubdata_spent) = if let Some(FrameRemnant {
         program_counter,
         exception_handler,
         snapshot,
-    }) = vm.state.current_frame.pop_near_call()
+        total_pubdata_spent,
+    }) =
+        vm.state.current_frame.pop_near_call()
     {
         (
             if TO_LABEL {
@@ -56,6 +58,7 @@ fn ret<const RETURN_TYPE: u8, const TO_LABEL: bool>(
             },
             snapshot,
             near_call_leftover_gas,
+            total_pubdata_spent,
         )
     } else {
         let return_value_or_panic = if return_type == ReturnType::Panic {
@@ -84,6 +87,7 @@ fn ret<const RETURN_TYPE: u8, const TO_LABEL: bool>(
             program_counter,
             exception_handler,
             snapshot,
+            total_pubdata_spent,
         }) = vm.pop_frame(
             return_value_or_panic
                 .as_ref()
@@ -125,12 +129,16 @@ fn ret<const RETURN_TYPE: u8, const TO_LABEL: bool>(
             },
             snapshot,
             leftover_gas,
+            total_pubdata_spent,
         )
     };
 
     if return_type.is_failure() {
         vm.world.rollback(snapshot);
+    } else {
+        vm.state.current_frame.total_pubdata_spent += total_pubdata_spent;
     }
+
     vm.state.flags = Flags::new(return_type == ReturnType::Panic, false, false);
     vm.state.current_frame.gas += leftover_gas;
 

diff --git a/src/instruction_handlers/storage.rs b/src/instruction_handlers/storage.rs
@@ -12,22 +12,22 @@ use crate::{
 
 fn sstore(vm: &mut VirtualMachine, instruction: *const Instruction) -> InstructionResult {
     instruction_boilerplate_with_panic(vm, instruction, |vm, args, continue_normally| {
-        let key = Register1::get(args, &mut vm.state);
-        let value = Register2::get(args, &mut vm.state);
-
-        // TODO: pubdata cost
-
         if vm.state.current_frame.is_static {
             return Ok(&PANIC);
         }
 
-        let refund = vm
-            .world
-            .write_storage(vm.state.current_frame.address, key, value);
+        let key = Register1::get(args, &mut vm.state);
+        let value = Register2::get(args, &mut vm.state);
+
+        let (refund, pubdata_change) =
+            vm.world
+                .write_storage(vm.state.current_frame.address, key, value);
 
         assert!(refund <= SSTORE_COST);
         vm.state.current_frame.gas += refund;
 
+        vm.state.current_frame.total_pubdata_spent += pubdata_change;
+
         continue_normally
     })
 }

diff --git a/src/lib.rs b/src/lib.rs
@@ -35,4 +35,10 @@ pub trait World {
 
     /// There is no write_storage; [ModifiedWorld::get_storage_changes] gives a list of all storage changes.
     fn read_storage(&mut self, contract: H160, key: U256) -> U256;
+
+    /// Computes the cost of writing a storage slot.
+    fn cost_of_writing_storage(&mut self, contract: H160, key: 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;
 }
diff --git a/src/modified_world.rs b/src/modified_world.rs
@@ -19,6 +19,7 @@ pub struct ModifiedWorld {
     storage_changes: RollbackableMap<(H160, U256), U256>,
     events: RollbackableLog<Event>,
     l2_to_l1_logs: RollbackableLog<L2ToL1Log>,
+    paid_changes: RollbackableMap<(H160, U256), u32>,
 
     // The fields below are only rolled back when the whole VM is rolled back.
     pub(crate) decommitted_hashes: RollbackableSet<U256>,
@@ -27,11 +28,7 @@ pub struct ModifiedWorld {
 }
 
 pub struct ExternalSnapshot {
-    storage_changes: <RollbackableMap<(H160, U256), U256> as Rollback>::Snapshot,
-    events: <RollbackableLog<Event> as Rollback>::Snapshot,
-    l2_to_l1_logs: <RollbackableLog<L2ToL1Log> as Rollback>::Snapshot,
-
-    // The field below are only rolled back when the whole VM is rolled back.
+    internal_snapshot: Snapshot,
     pub(crate) decommitted_hashes: <RollbackableMap<U256, ()> as Rollback>::Snapshot,
     read_storage_slots: <RollbackableMap<(H160, U256), ()> as Rollback>::Snapshot,
     written_storage_slots: <RollbackableMap<(H160, U256), ()> as Rollback>::Snapshot,
@@ -67,19 +64,22 @@ impl ModifiedWorld {
             decommitted_hashes: Default::default(),
             read_storage_slots: Default::default(),
             written_storage_slots: Default::default(),
+            paid_changes: Default::default(),
         }
     }
 
     /// Returns the storage slot's value and a refund based on its hot/cold status.
-    pub fn read_storage(&mut self, contract: H160, key: U256) -> (U256, u32) {
+    pub(crate) fn read_storage(&mut self, contract: H160, key: U256) -> (U256, u32) {
         let value = self
             .storage_changes
             .as_ref()
             .get(&(contract, key))
             .cloned()
             .unwrap_or_else(|| self.world.read_storage(contract, key));
 
-        let refund = if self.read_storage_slots.contains(&(contract, key)) {
+        let refund = if self.world.is_free_storage_slot(&contract, &key)
+            || self.read_storage_slots.contains(&(contract, key))
+        {
             WARM_READ_REFUND
         } else {
             self.read_storage_slots.add((contract, key));
@@ -89,11 +89,21 @@ impl ModifiedWorld {
         (value, refund)
     }
 
-    /// Returns the refund based the hot/cold status of the storage slot.
-    pub fn write_storage(&mut self, contract: H160, key: U256, value: U256) -> u32 {
+    /// Returns the refund based the hot/cold status of the storage slot and the change in pubdata.
+    pub(crate) fn write_storage(&mut self, contract: H160, key: U256, value: U256) -> (u32, i32) {
         self.storage_changes.insert((contract, key), value);
 
-        if self
+        if self.world.is_free_storage_slot(&contract, &key) {
+            return (WARM_WRITE_REFUND, 0);
+        }
+
+        let update_cost = self.world.cost_of_writing_storage(contract, key, value);
+        let prepaid = self
+            .paid_changes
+            .insert((contract, key), update_cost)
+            .unwrap_or(0);
+
+        let refund = if self
             .written_storage_slots
             .as_ref()
             .contains_key(&(contract, key))
@@ -108,7 +118,9 @@ impl ModifiedWorld {
                 self.read_storage_slots.add((contract, key));
                 0
             }
-        }
+        };
+
+        (refund, (update_cost as i32) - (prepaid as i32))
     }
 
     pub fn get_storage_changes(&self) -> &BTreeMap<(H160, U256), U256> {
@@ -136,32 +148,30 @@ impl ModifiedWorld {
             self.storage_changes.snapshot(),
             self.events.snapshot(),
             self.l2_to_l1_logs.snapshot(),
+            self.paid_changes.snapshot(),
         )
     }
 
-    pub(crate) fn rollback(&mut self, (storage, events, l2_to_l1_logs): Snapshot) {
+    pub(crate) fn rollback(&mut self, (storage, events, l2_to_l1_logs, paid_changes): Snapshot) {
         self.storage_changes.rollback(storage);
         self.events.rollback(events);
         self.l2_to_l1_logs.rollback(l2_to_l1_logs);
+        self.paid_changes.rollback(paid_changes);
     }
 
     /// This function must only be called during the initial frame
     /// because otherwise internal rollbacks can roll back past the external snapshot.
     pub(crate) fn external_snapshot(&self) -> ExternalSnapshot {
         ExternalSnapshot {
-            storage_changes: self.storage_changes.snapshot(),
-            events: self.events.snapshot(),
-            l2_to_l1_logs: self.l2_to_l1_logs.snapshot(),
+            internal_snapshot: self.snapshot(),
             decommitted_hashes: self.decommitted_hashes.snapshot(),
             read_storage_slots: self.read_storage_slots.snapshot(),
             written_storage_slots: self.written_storage_slots.snapshot(),
         }
     }
 
     pub(crate) fn external_rollback(&mut self, snapshot: ExternalSnapshot) {
-        self.storage_changes.rollback(snapshot.storage_changes);
-        self.events.rollback(snapshot.events);
-        self.l2_to_l1_logs.rollback(snapshot.l2_to_l1_logs);
+        self.rollback(snapshot.internal_snapshot);
         self.decommitted_hashes
             .rollback(snapshot.decommitted_hashes);
         self.read_storage_slots
@@ -187,6 +197,7 @@ pub(crate) type Snapshot = (
     <RollbackableMap<(H160, U256), U256> as Rollback>::Snapshot,
     <RollbackableLog<Event> as Rollback>::Snapshot,
     <RollbackableLog<L2ToL1Log> as Rollback>::Snapshot,
+    <RollbackableMap<(H160, U256), u32> as Rollback>::Snapshot,
 );
 
 const WARM_READ_REFUND: u32 = STORAGE_ACCESS_COLD_READ_COST - STORAGE_ACCESS_WARM_READ_COST;

diff --git a/src/rollback.rs b/src/rollback.rs
@@ -14,10 +14,11 @@ pub struct RollbackableMap<K: Ord, V> {
     old_entries: Vec<(K, Option<V>)>,
 }
 
-impl<K: Ord + Clone, V> RollbackableMap<K, V> {
-    pub fn insert(&mut self, key: K, value: V) {
-        self.old_entries
-            .push((key.clone(), self.map.insert(key, value)));
+impl<K: Ord + Clone, V: Clone> RollbackableMap<K, V> {
+    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
+        let old_value = self.map.insert(key.clone(), value);
+        self.old_entries.push((key.clone(), old_value.clone()));
+        old_value
     }
 }
 

diff --git a/src/testworld.rs b/src/testworld.rs
@@ -61,4 +61,17 @@ impl World for TestWorld {
             0.into()
         }
     }
+
+    fn cost_of_writing_storage(
+        &mut self,
+        _contract: u256::H160,
+        _key: U256,
+        _new_value: U256,
+    ) -> u32 {
+        50
+    }
+
+    fn is_free_storage_slot(&self, _contract: &u256::H160, _key: &U256) -> bool {
+        false
+    }
 }
diff --git a/src/vm.rs b/src/vm.rs
@@ -247,6 +247,7 @@ impl VirtualMachine {
                     exception_handler,
                     world_before_this_frame,
                     stack,
+                    total_pubdata_spent,
                     ..
                 } = frame;
 
@@ -261,6 +262,7 @@ impl VirtualMachine {
                     program_counter,
                     exception_handler,
                     snapshot: world_before_this_frame,
+                    total_pubdata_spent,
                 }
             })
     }