diff --git a/src/callframe.rs b/src/callframe.rs index d84434a3..98e364fb 100644 --- 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, 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 index 844b5c3f..90897565 100644 --- a/src/instruction_handlers/ret.rs +++ b/src/instruction_handlers/ret.rs @@ -40,11 +40,13 @@ fn ret( 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( }, 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( program_counter, exception_handler, snapshot, + total_pubdata_spent, }) = vm.pop_frame( return_value_or_panic .as_ref() @@ -125,12 +129,16 @@ fn ret( }, 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 index f3266e0b..72ef87c8 100644 --- 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 index dac37577..f90cc056 100644 --- 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 index 7d9e5578..02a45f5e 100644 --- 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, l2_to_l1_logs: RollbackableLog, + 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, @@ -27,11 +28,7 @@ pub struct ModifiedWorld { } pub struct ExternalSnapshot { - storage_changes: as Rollback>::Snapshot, - events: as Rollback>::Snapshot, - l2_to_l1_logs: as Rollback>::Snapshot, - - // The field below are only rolled back when the whole VM is rolled back. + internal_snapshot: Snapshot, pub(crate) decommitted_hashes: as Rollback>::Snapshot, read_storage_slots: as Rollback>::Snapshot, written_storage_slots: as Rollback>::Snapshot, @@ -67,11 +64,12 @@ 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() @@ -79,7 +77,9 @@ impl ModifiedWorld { .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,22 +148,22 @@ 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(), @@ -159,9 +171,7 @@ impl ModifiedWorld { } 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 = ( as Rollback>::Snapshot, as Rollback>::Snapshot, as Rollback>::Snapshot, + 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 index 0a65469b..6d69908b 100644 --- a/src/rollback.rs +++ b/src/rollback.rs @@ -14,10 +14,11 @@ pub struct RollbackableMap { old_entries: Vec<(K, Option)>, } -impl RollbackableMap { - pub fn insert(&mut self, key: K, value: V) { - self.old_entries - .push((key.clone(), self.map.insert(key, value))); +impl RollbackableMap { + pub fn insert(&mut self, key: K, value: V) -> Option { + 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 index 8e72ab98..c4eca8b8 100644 --- 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 index dd991e12..e23ebdb8 100644 --- 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, } }) }