handle pubdata using the normal rollback mechanism

src/callframe.rs

@@ -21,8 +21,6 @@ pub struct Callframe {
     pub gas: u32,
     pub stipend: u32,
 
-    pub total_pubdata_spent: i32,
-
     near_calls: Vec<NearCallFrame>,
 
     pub(crate) program: Program,
@@ -102,7 +100,6 @@ impl Callframe {
             exception_handler,
             near_calls: vec![],
             world_before_this_frame,
-            total_pubdata_spent: 0,
         }
     }
 
@@ -132,7 +129,6 @@ impl Callframe {
                 program_counter: f.call_instruction,
                 exception_handler: f.exception_handler,
                 snapshot: f.world_before_this_frame,
-                total_pubdata_spent: 0,
             }
         })
     }
@@ -163,5 +159,4 @@ pub(crate) struct FrameRemnant {
     pub(crate) program_counter: u16,
     pub(crate) exception_handler: u16,
     pub(crate) snapshot: Snapshot,
-    pub(crate) total_pubdata_spent: i32,
 }

src/instruction_handlers/context.rs

@@ -69,21 +69,26 @@ impl ContextOp for SP {
     }
 }
 
-struct Meta;
-impl ContextOp for Meta {
-    fn get(state: &State) -> U256 {
-        VmMetaParameters {
-            heap_size: state.current_frame.heap_size,
-            aux_heap_size: state.current_frame.aux_heap_size,
+fn context_meta(
+    vm: &mut VirtualMachine,
+    instruction: *const Instruction,
+    world: &mut dyn World,
+) -> InstructionResult {
+    instruction_boilerplate(vm, instruction, world, |vm, args, _| {
+        let result = VmMetaParameters {
+            heap_size: vm.state.current_frame.heap_size,
+            aux_heap_size: vm.state.current_frame.aux_heap_size,
             this_shard_id: 0, // TODO properly implement shards
             caller_shard_id: 0,
             code_shard_id: 0,
             // This field is actually pubdata!
             // TODO PLA-893: This should be zero when not in kernel mode
-            aux_field_0: state.current_frame.total_pubdata_spent as u32,
+            aux_field_0: vm.world_diff.pubdata.0 as u32,
         }
-        .to_u256()
-    }
+        .to_u256();
+
+        Register1::set(args, &mut vm.state, result);
+    })
 }
 
 fn set_context_u128(
@@ -145,7 +150,10 @@ impl Instruction {
         Self::from_context::<SP>(out, arguments)
     }
     pub fn from_context_meta(out: Register1, arguments: Arguments) -> Self {
-        Self::from_context::<Meta>(out, arguments)
+        Self {
+            handler: context_meta,
+            arguments: arguments.write_destination(&out),
+        }
     }
     pub fn from_set_context_u128(src: Register1, arguments: Arguments) -> Self {
         Self {

src/instruction_handlers/precompiles.rs

@@ -37,7 +37,7 @@ fn precompile_call(
         let Ok(()) = vm.state.use_gas(aux_data.extra_ergs_cost) else {
             return Ok(&PANIC);
         };
-        vm.state.current_frame.total_pubdata_spent += aux_data.extra_pubdata_cost as i32;
+        vm.world_diff.pubdata.0 += aux_data.extra_pubdata_cost as i32;
 
         let mut abi = PrecompileCallABI::from_u256(Register1::get(args, &mut vm.state));
         if abi.memory_page_to_read == 0 {

src/instruction_handlers/ret.rs

@@ -41,13 +41,11 @@ 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, total_pubdata_spent) = if let Some(FrameRemnant {
+    let (pc, snapshot, leftover_gas) = if let Some(FrameRemnant {
         program_counter,
         exception_handler,
         snapshot,
-        total_pubdata_spent,
-    }) =
-        vm.state.current_frame.pop_near_call()
+    }) = vm.state.current_frame.pop_near_call()
     {
         (
             if TO_LABEL {
@@ -59,7 +57,6 @@ 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 {
@@ -88,7 +85,6 @@ 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()
@@ -130,14 +126,11 @@ fn ret<const RETURN_TYPE: u8, const TO_LABEL: bool>(
             },
             snapshot,
             leftover_gas,
-            total_pubdata_spent,
         )
     };
 
     if return_type.is_failure() {
         vm.world_diff.rollback(snapshot);
-    } else {
-        vm.state.current_frame.total_pubdata_spent += total_pubdata_spent;
     }
 
     vm.state.flags = Flags::new(return_type == ReturnType::Panic, false, false);

src/instruction_handlers/storage.rs

@@ -27,15 +27,13 @@ fn sstore(
             let key = Register1::get(args, &mut vm.state);
             let value = Register2::get(args, &mut vm.state);
 
-            let (refund, pubdata_change) =
+            let refund =
                 vm.world_diff
                     .write_storage(world, 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
         },
     )

src/rollback.rs

@@ -117,3 +117,21 @@ impl<T> AsRef<[T]> for RollbackableLog<T> {
         &self.entries
     }
 }
+
+/// Rollbackable Plain Old Data simply stores copies of itself in snapshots.
+#[derive(Default, Copy, Clone)]
+pub struct RollbackablePod<T: Copy>(pub T);
+
+impl<T: Copy> Rollback for RollbackablePod<T> {
+    type Snapshot = T;
+
+    fn snapshot(&self) -> Self::Snapshot {
+        self.0
+    }
+
+    fn rollback(&mut self, snapshot: Self::Snapshot) {
+        self.0 = snapshot
+    }
+
+    fn delete_history(&mut self) {}
+}

src/vm.rs

@@ -240,7 +240,6 @@ impl VirtualMachine {
                     exception_handler,
                     world_before_this_frame,
                     stack,
-                    total_pubdata_spent,
                     ..
                 } = frame;
 
@@ -255,7 +254,6 @@ impl VirtualMachine {
                     program_counter,
                     exception_handler,
                     snapshot: world_before_this_frame,
-                    total_pubdata_spent,
                 }
             })
     }

src/world_diff.rs

@@ -1,7 +1,7 @@
 use std::collections::BTreeMap;
 
 use crate::{
-    rollback::{Rollback, RollbackableLog, RollbackableMap, RollbackableSet},
+    rollback::{Rollback, RollbackableLog, RollbackableMap, RollbackablePod, RollbackableSet},
     World,
 };
 use u256::{H160, U256};
@@ -20,6 +20,7 @@ pub struct WorldDiff {
     transient_storage_changes: RollbackableMap<(H160, U256), U256>,
     events: RollbackableLog<Event>,
     l2_to_l1_logs: RollbackableLog<L2ToL1Log>,
+    pub pubdata: RollbackablePod<i32>,
 
     // The fields below are only rolled back when the whole VM is rolled back.
     pub(crate) decommitted_hashes: RollbackableSet<U256>,
@@ -81,34 +82,18 @@ impl WorldDiff {
         (value, refund)
     }
 
-    pub(crate) fn read_transient_storage(&mut self, contract: H160, key: U256) -> U256 {
-        let value = self
-            .transient_storage_changes
-            .as_ref()
-            .get(&(contract, key))
-            .cloned()
-            .unwrap_or_default();
-
-        value
-    }
-
-    pub(crate) fn write_transient_storage(&mut self, contract: H160, key: U256, value: U256) {
-        self.transient_storage_changes
-            .insert((contract, key), value);
-    }
-
     /// Returns the refund based the hot/cold status of the storage slot and the change in pubdata.
     pub(crate) fn write_storage(
         &mut self,
         world: &mut dyn World,
         contract: H160,
         key: U256,
         value: U256,
-    ) -> (u32, i32) {
+    ) -> u32 {
         self.storage_changes.insert((contract, key), value);
 
         if world.is_free_storage_slot(&contract, &key) {
-            return (WARM_WRITE_REFUND, 0);
+            return WARM_WRITE_REFUND;
         }
 
         let update_cost = world.cost_of_writing_storage(contract, key, value);
@@ -134,7 +119,9 @@ impl WorldDiff {
             }
         };
 
-        (refund, (update_cost as i32) - (prepaid as i32))
+        self.pubdata.0 += (update_cost as i32) - (prepaid as i32);
+
+        refund
     }
 
     pub fn get_storage_state(&self) -> &BTreeMap<(H160, U256), U256> {
@@ -152,6 +139,22 @@ impl WorldDiff {
         self.storage_changes.changes_after(snapshot.storage_changes)
     }
 
+    pub(crate) fn read_transient_storage(&mut self, contract: H160, key: U256) -> U256 {
+        let value = self
+            .transient_storage_changes
+            .as_ref()
+            .get(&(contract, key))
+            .cloned()
+            .unwrap_or_default();
+
+        value
+    }
+
+    pub(crate) fn write_transient_storage(&mut self, contract: H160, key: U256, value: U256) {
+        self.transient_storage_changes
+            .insert((contract, key), value);
+    }
+
     pub(crate) fn record_event(&mut self, event: Event) {
         self.events.push(event);
     }
@@ -186,29 +189,32 @@ impl WorldDiff {
     pub fn snapshot(&self) -> Snapshot {
         Snapshot {
             storage_changes: self.storage_changes.snapshot(),
+            paid_changes: self.paid_changes.snapshot(),
             events: self.events.snapshot(),
             l2_to_l1_logs: self.l2_to_l1_logs.snapshot(),
-            paid_changes: self.paid_changes.snapshot(),
             transient_storage_changes: self.transient_storage_changes.snapshot(),
+            pubdata: self.pubdata.snapshot(),
         }
     }
 
     pub(crate) fn rollback(
         &mut self,
         Snapshot {
             storage_changes,
+            paid_changes,
             events,
             l2_to_l1_logs,
-            paid_changes,
             transient_storage_changes,
+            pubdata,
         }: Snapshot,
     ) {
         self.storage_changes.rollback(storage_changes);
+        self.paid_changes.rollback(paid_changes);
         self.events.rollback(events);
         self.l2_to_l1_logs.rollback(l2_to_l1_logs);
-        self.paid_changes.rollback(paid_changes);
         self.transient_storage_changes
             .rollback(transient_storage_changes);
+        self.pubdata.rollback(pubdata);
     }
 
     /// This function must only be called during the initial frame
@@ -259,10 +265,11 @@ impl WorldDiff {
 #[derive(Clone, PartialEq, Debug)]
 pub struct Snapshot {
     storage_changes: <RollbackableMap<(H160, U256), U256> as Rollback>::Snapshot,
+    paid_changes: <RollbackableMap<(H160, U256), u32> as Rollback>::Snapshot,
     events: <RollbackableLog<Event> as Rollback>::Snapshot,
     l2_to_l1_logs: <RollbackableLog<L2ToL1Log> as Rollback>::Snapshot,
-    paid_changes: <RollbackableMap<(H160, U256), u32> as Rollback>::Snapshot,
     transient_storage_changes: <RollbackableMap<(H160, U256), U256> as Rollback>::Snapshot,
+    pubdata: <RollbackablePod<i32> as Rollback>::Snapshot,
 }
 
 const WARM_READ_REFUND: u32 = STORAGE_ACCESS_COLD_READ_COST - STORAGE_ACCESS_WARM_READ_COST;