zksync_vm2/
world_diff.rs

1use std::collections::BTreeMap;
2
3use primitive_types::{H160, U256};
4use zk_evm_abstractions::{aux::Timestamp, queries::LogQuery};
5use zkevm_opcode_defs::system_params::{
6    STORAGE_ACCESS_COLD_READ_COST, STORAGE_ACCESS_COLD_WRITE_COST, STORAGE_ACCESS_WARM_READ_COST,
7    STORAGE_ACCESS_WARM_WRITE_COST, STORAGE_AUX_BYTE,
8};
9use zksync_vm2_interface::{CycleStats, Event, HeapId, L2ToL1Log, Tracer};
10
11use crate::{
12    rollback::{Rollback, RollbackableLog, RollbackableMap, RollbackablePod, RollbackableSet},
13    StorageInterface, StorageSlot,
14};
15
16/// Pending modifications to the global state that are executed at the end of a block.
17/// In other words, side effects.
18#[derive(Debug, Default)]
19pub struct WorldDiff {
20    // These are rolled back on revert or panic (and when the whole VM is rolled back).
21    storage_changes: RollbackableMap<(H160, U256), U256>,
22    paid_changes: RollbackableMap<(H160, U256), u32>,
23    transient_storage_changes: RollbackableMap<(H160, U256), U256>,
24    events: RollbackableLog<Event>,
25    l2_to_l1_logs: RollbackableLog<L2ToL1Log>,
26    pub(crate) pubdata: RollbackablePod<i32>,
27    storage_refunds: RollbackableLog<u32>,
28    pubdata_costs: RollbackableLog<i32>,
29    storage_logs: Vec<LogQuery>,
30    rollback_storage_logs: Vec<LogQuery>,
31
32    // The fields below are only rolled back when the whole VM is rolled back.
33    /// Tracks decommit visibility state for each bytecode hash.
34    ///
35    /// Besides successful decommits, we also retain far-call decommit attempts that failed with
36    /// out-of-gas in `pay_for_decommit()`. Legacy VM includes those hashes into "used contracts"
37    /// output, and shadow-mode compares that output (`CurrentExecutionState.used_contract_hashes`).
38    ///
39    /// This field is rolled back only by external VM snapshots.
40    pub(crate) decommitted_hashes: RollbackableMap<U256, DecommitState>,
41    /// Reverse index for `decommitted_hashes` entries that carry materialized heap pages.
42    ///
43    /// This is used to quickly check whether a heap page is globally pinned by decommitment reuse
44    /// semantics.
45    ///
46    /// This follows external snapshot / rollback semantics together with `decommitted_hashes`.
47    decommit_pinned_pages: RollbackableSet<u32>,
48    read_storage_slots: RollbackableSet<(H160, U256)>,
49    written_storage_slots: RollbackableSet<(H160, U256)>,
50
51    // This is never rolled back. It is just a cache to avoid asking these from DB every time.
52    storage_initial_values: BTreeMap<(H160, U256), StorageSlot>,
53}
54
55#[derive(Debug)]
56pub(crate) struct ExternalSnapshot {
57    internal_snapshot: Snapshot,
58    pub(crate) decommitted_hashes: <RollbackableMap<U256, DecommitState> as Rollback>::Snapshot,
59    decommit_pinned_pages: <RollbackableSet<u32> as Rollback>::Snapshot,
60    read_storage_slots: <RollbackableMap<(H160, U256), ()> as Rollback>::Snapshot,
61    written_storage_slots: <RollbackableMap<(H160, U256), ()> as Rollback>::Snapshot,
62    storage_refunds: <RollbackableLog<u32> as Rollback>::Snapshot,
63    pubdata_costs: <RollbackableLog<i32> as Rollback>::Snapshot,
64}
65
66#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
67pub(crate) enum DecommitState {
68    /// A far-call decommit attempt ran out of gas before materialization.
69    ///
70    /// We preserve this state for legacy compatibility: old VM exposes these hashes as used
71    /// contracts. This state is observable via `decommitted_hashes()`, but it must not make
72    /// future decommits free.
73    ///
74    /// Note that `log.decommit` out-of-gas is not represented by this state because that opcode
75    /// exits before decommit bookkeeping.
76    #[default]
77    Unsuccessful,
78    /// A bytecode hash was successfully decommitted and has an assigned reusable heap page.
79    Succeeded(u32),
80}
81
82impl WorldDiff {
83    /// Returns the storage slot's value and a refund based on its hot/cold status.
84    pub(crate) fn read_storage(
85        &mut self,
86        world: &mut impl StorageInterface,
87        tracer: &mut impl Tracer,
88        contract: H160,
89        key: U256,
90        tx_number_in_block: u16,
91    ) -> (U256, u32) {
92        let (value, newly_added) =
93            self.read_storage_inner(world, tracer, contract, key, tx_number_in_block);
94        let refund = if !newly_added || world.is_free_storage_slot(&contract, &key) {
95            WARM_READ_REFUND
96        } else {
97            0
98        };
99        self.storage_refunds.push(refund);
100        (value, refund)
101    }
102
103    /// Same as [`Self::read_storage()`], but without recording the refund value (which is important
104    /// because the storage is read not only from the `sload` op handler, but also from the `farcall` op handler;
105    /// the latter must not record a refund as per previous VM versions).
106    pub(crate) fn read_storage_without_refund(
107        &mut self,
108        world: &mut impl StorageInterface,
109        tracer: &mut impl Tracer,
110        contract: H160,
111        key: U256,
112        tx_number_in_block: u16,
113    ) -> U256 {
114        self.read_storage_inner(world, tracer, contract, key, tx_number_in_block)
115            .0
116    }
117
118    fn read_storage_inner(
119        &mut self,
120        world: &mut impl StorageInterface,
121        tracer: &mut impl Tracer,
122        contract: H160,
123        key: U256,
124        tx_number_in_block: u16,
125    ) -> (U256, bool) {
126        let newly_added = self.read_storage_slots.add((contract, key));
127        if newly_added {
128            tracer.on_extra_prover_cycles(CycleStats::StorageRead);
129        }
130
131        self.pubdata_costs.push(0);
132        let value = self.just_read_storage(world, contract, key);
133        // Note: timestamp logic does not match `zk_evm`, we're only ensuring that the timestamps
134        // are unique. This is fine as long as we don't need to actually generate witness based on these logs.
135        self.storage_logs.push(LogQuery {
136            timestamp: Timestamp(
137                u32::try_from(self.storage_logs.len()).expect("Too many storage logs"),
138            ),
139            tx_number_in_block,
140            aux_byte: STORAGE_AUX_BYTE,
141            shard_id: 0,
142            address: contract,
143            key,
144            read_value: value,
145            written_value: value,
146            rw_flag: false,
147            rollback: false,
148            is_service: false,
149        });
150        (value, newly_added)
151    }
152
153    /// Reads the value of a storage slot without any extra bookkeeping.
154    /// Should only be used for tracers.
155    pub(crate) fn just_read_storage(
156        &self,
157        world: &mut impl StorageInterface,
158        contract: H160,
159        key: U256,
160    ) -> U256 {
161        self.storage_changes
162            .as_ref()
163            .get(&(contract, key))
164            .copied()
165            .unwrap_or_else(|| world.read_storage_value(contract, key))
166    }
167
168    /// Returns the refund based the hot/cold status of the storage slot and the change in pubdata.
169    pub(crate) fn write_storage(
170        &mut self,
171        world: &mut impl StorageInterface,
172        tracer: &mut impl Tracer,
173        contract: H160,
174        key: U256,
175        value: U256,
176        tx_number_in_block: u16,
177    ) -> u32 {
178        let read_value = self.just_read_storage(world, contract, key);
179        let log_query = LogQuery {
180            timestamp: Timestamp(u32::try_from(self.storage_logs.len()).unwrap_or(u32::MAX)),
181            tx_number_in_block,
182            aux_byte: STORAGE_AUX_BYTE,
183            shard_id: 0,
184            address: contract,
185            key,
186            read_value,
187            written_value: value,
188            rw_flag: true,
189            rollback: false,
190            is_service: false,
191        };
192        self.storage_logs.push(log_query);
193        self.rollback_storage_logs.push(LogQuery {
194            rollback: true,
195            ..log_query
196        });
197
198        self.storage_changes.insert((contract, key), value);
199
200        let initial_value = self
201            .storage_initial_values
202            .entry((contract, key))
203            .or_insert_with(|| world.read_storage(contract, key));
204
205        if world.is_free_storage_slot(&contract, &key) {
206            if self.written_storage_slots.add((contract, key)) {
207                tracer.on_extra_prover_cycles(CycleStats::StorageWrite);
208            }
209            self.read_storage_slots.add((contract, key));
210
211            self.storage_refunds.push(WARM_WRITE_REFUND);
212            self.pubdata_costs.push(0);
213            return WARM_WRITE_REFUND;
214        }
215
216        let update_cost = world.cost_of_writing_storage(*initial_value, value);
217        let prepaid = self
218            .paid_changes
219            .insert((contract, key), update_cost)
220            .unwrap_or(0);
221
222        let refund = if self.written_storage_slots.add((contract, key)) {
223            tracer.on_extra_prover_cycles(CycleStats::StorageWrite);
224
225            if self.read_storage_slots.add((contract, key)) {
226                0
227            } else {
228                COLD_WRITE_AFTER_WARM_READ_REFUND
229            }
230        } else {
231            WARM_WRITE_REFUND
232        };
233
234        #[allow(clippy::cast_possible_wrap)]
235        {
236            let pubdata_cost = (update_cost as i32) - (prepaid as i32);
237            self.pubdata.0 += pubdata_cost;
238            self.storage_refunds.push(refund);
239            self.pubdata_costs.push(pubdata_cost);
240        }
241        refund
242    }
243
244    pub(crate) fn pubdata(&self) -> i32 {
245        self.pubdata.0
246    }
247
248    /// Returns recorded refunds for all storage operations.
249    pub fn storage_refunds(&self) -> &[u32] {
250        self.storage_refunds.as_ref()
251    }
252
253    /// Returns recorded pubdata costs for all storage operations.
254    pub fn pubdata_costs(&self) -> &[i32] {
255        self.pubdata_costs.as_ref()
256    }
257
258    /// Returns all recorded storage log queries.
259    ///
260    /// These logs are sufficient for vm2 state-transition checks and diagnostics.
261    // TODO: We don't fill all the `zk_evm` witness metadata, so this is not suitable for
262    // generating EraVM prover witness data. This is not the goal, however, as we only need
263    // to emit enough data to verify the correctness of the state transition.
264    pub fn storage_log_queries(&self) -> &[LogQuery] {
265        &self.storage_logs
266    }
267
268    /// Returns storage log queries recorded after the specified `snapshot` was created.
269    pub fn storage_log_queries_after(&self, snapshot: &Snapshot) -> &[LogQuery] {
270        &self.storage_logs[snapshot.storage_logs_len..]
271    }
272
273    #[doc(hidden)] // duplicates `StateInterface::get_storage_state()`, but we use random access in some places
274    pub fn get_storage_state(&self) -> &BTreeMap<(H160, U256), U256> {
275        self.storage_changes.as_ref()
276    }
277
278    /// Gets changes for all touched storage slots.
279    pub fn get_storage_changes(&self) -> impl Iterator<Item = ((H160, U256), StorageChange)> + '_ {
280        self.storage_changes
281            .as_ref()
282            .iter()
283            .filter_map(|(key, &value)| {
284                let initial_slot = &self.storage_initial_values[key];
285                if initial_slot.value == value {
286                    None
287                } else {
288                    Some((
289                        *key,
290                        StorageChange {
291                            before: initial_slot.value,
292                            after: value,
293                            is_initial: initial_slot.is_write_initial,
294                        },
295                    ))
296                }
297            })
298    }
299
300    /// Gets changes for storage slots touched after the specified `snapshot` was created.
301    pub fn get_storage_changes_after(
302        &self,
303        snapshot: &Snapshot,
304    ) -> impl Iterator<Item = ((H160, U256), StorageChange)> + '_ {
305        self.storage_changes
306            .changes_after(snapshot.storage_changes)
307            .into_iter()
308            .map(|(key, (before, after))| {
309                let initial = self.storage_initial_values[&key];
310                (
311                    key,
312                    StorageChange {
313                        before: before.unwrap_or(initial.value),
314                        after,
315                        is_initial: initial.is_write_initial,
316                    },
317                )
318            })
319    }
320
321    pub(crate) fn read_transient_storage(&mut self, contract: H160, key: U256) -> U256 {
322        self.pubdata_costs.push(0);
323        self.transient_storage_changes
324            .as_ref()
325            .get(&(contract, key))
326            .copied()
327            .unwrap_or_default()
328    }
329
330    pub(crate) fn write_transient_storage(&mut self, contract: H160, key: U256, value: U256) {
331        self.pubdata_costs.push(0);
332        self.transient_storage_changes
333            .insert((contract, key), value);
334    }
335
336    pub(crate) fn get_transient_storage_state(&self) -> &BTreeMap<(H160, U256), U256> {
337        self.transient_storage_changes.as_ref()
338    }
339
340    pub(crate) fn record_event(&mut self, event: Event) {
341        self.events.push(event);
342    }
343
344    pub(crate) fn events(&self) -> &[Event] {
345        self.events.as_ref()
346    }
347
348    /// Returns events emitted after the specified `snapshot` was created.
349    pub fn events_after(&self, snapshot: &Snapshot) -> &[Event] {
350        self.events.logs_after(snapshot.events)
351    }
352
353    pub(crate) fn record_l2_to_l1_log(&mut self, log: L2ToL1Log) {
354        self.l2_to_l1_logs.push(log);
355    }
356
357    pub(crate) fn l2_to_l1_logs(&self) -> &[L2ToL1Log] {
358        self.l2_to_l1_logs.as_ref()
359    }
360
361    /// Returns L2-to-L1 logs emitted after the specified `snapshot` was created.
362    pub fn l2_to_l1_logs_after(&self, snapshot: &Snapshot) -> &[L2ToL1Log] {
363        self.l2_to_l1_logs.logs_after(snapshot.l2_to_l1_logs)
364    }
365
366    /// Returns hashes of contract bytecodes that were observed by decommit bookkeeping in no
367    /// particular order.
368    ///
369    /// This includes successful decommits and far-call out-of-gas attempts recorded as
370    /// [`DecommitState::Unsuccessful`] for legacy `used_contract_hashes` compatibility.
371    pub fn decommitted_hashes(&self) -> impl Iterator<Item = U256> + '_ {
372        self.decommitted_hashes.as_ref().keys().copied()
373    }
374
375    pub(crate) fn decommit_page(&self, code_hash: U256) -> Option<HeapId> {
376        self.decommitted_hashes
377            .as_ref()
378            .get(&code_hash)
379            .and_then(|state| {
380                if let DecommitState::Succeeded(page) = state {
381                    Some(HeapId::from_u32_unchecked(*page))
382                } else {
383                    None
384                }
385            })
386    }
387
388    pub(crate) fn is_decommit_page_pinned(&self, page: HeapId) -> bool {
389        self.decommit_pinned_pages.as_ref().contains(&page.as_u32())
390    }
391
392    pub(crate) fn set_decommit_page(&mut self, code_hash: U256, page: HeapId) {
393        self.decommitted_hashes
394            .insert(code_hash, DecommitState::Succeeded(page.as_u32()));
395        self.decommit_pinned_pages.add(page.as_u32());
396    }
397
398    /// Get a snapshot for selecting which logs & co. to output using [`Self::events_after()`] and other methods.
399    pub fn snapshot(&self) -> Snapshot {
400        Snapshot {
401            storage_changes: self.storage_changes.snapshot(),
402            paid_changes: self.paid_changes.snapshot(),
403            events: self.events.snapshot(),
404            l2_to_l1_logs: self.l2_to_l1_logs.snapshot(),
405            transient_storage_changes: self.transient_storage_changes.snapshot(),
406            pubdata: self.pubdata.snapshot(),
407            storage_logs_len: self.storage_logs.len(),
408            rollback_storage_logs_len: self.rollback_storage_logs.len(),
409        }
410    }
411
412    /// Appends rollback storage logs recorded after `snapshot` to `storage_logs`.
413    ///
414    /// This is needed for failed frame returns (revert / panic) where rolled-back writes
415    /// must remain observable in the storage log stream.
416    pub(crate) fn append_rollback_logs(&mut self, snapshot: &Snapshot) {
417        if self.rollback_storage_logs.len() > snapshot.rollback_storage_logs_len {
418            let rollback_logs = self
419                .rollback_storage_logs
420                .split_off(snapshot.rollback_storage_logs_len);
421            for log in rollback_logs.into_iter().rev() {
422                self.storage_logs.push(log);
423            }
424        }
425    }
426
427    #[allow(clippy::needless_pass_by_value)] // intentional: we require a snapshot to be rolled back to no more than once
428    pub(crate) fn rollback(&mut self, snapshot: Snapshot) {
429        self.storage_changes.rollback(snapshot.storage_changes);
430        self.paid_changes.rollback(snapshot.paid_changes);
431        self.events.rollback(snapshot.events);
432        self.l2_to_l1_logs.rollback(snapshot.l2_to_l1_logs);
433        self.transient_storage_changes
434            .rollback(snapshot.transient_storage_changes);
435        self.pubdata.rollback(snapshot.pubdata);
436    }
437
438    /// This function must only be called during the initial frame
439    /// because otherwise internal rollbacks can roll back past the external snapshot.
440    pub(crate) fn external_snapshot(&self) -> ExternalSnapshot {
441        // Rolling back to this snapshot will clear transient storage even though it is not empty
442        // after a transaction. This is ok because the next instruction in the bootloader
443        // (IncrementTxNumber) clears the transient storage anyway.
444        // This is necessary because clear_transient_storage cannot be undone.
445        ExternalSnapshot {
446            internal_snapshot: Snapshot {
447                transient_storage_changes: 0,
448                ..self.snapshot()
449            },
450            decommitted_hashes: self.decommitted_hashes.snapshot(),
451            decommit_pinned_pages: self.decommit_pinned_pages.snapshot(),
452            read_storage_slots: self.read_storage_slots.snapshot(),
453            written_storage_slots: self.written_storage_slots.snapshot(),
454            storage_refunds: self.storage_refunds.snapshot(),
455            pubdata_costs: self.pubdata_costs.snapshot(),
456        }
457    }
458
459    pub(crate) fn external_rollback(&mut self, snapshot: ExternalSnapshot) {
460        let storage_logs_len = snapshot.internal_snapshot.storage_logs_len;
461        let rollback_storage_logs_len = snapshot.internal_snapshot.rollback_storage_logs_len;
462
463        self.rollback(snapshot.internal_snapshot);
464        self.storage_refunds.rollback(snapshot.storage_refunds);
465        self.pubdata_costs.rollback(snapshot.pubdata_costs);
466        self.decommitted_hashes
467            .rollback(snapshot.decommitted_hashes);
468        self.decommit_pinned_pages
469            .rollback(snapshot.decommit_pinned_pages);
470        self.read_storage_slots
471            .rollback(snapshot.read_storage_slots);
472        self.written_storage_slots
473            .rollback(snapshot.written_storage_slots);
474        self.storage_logs.truncate(storage_logs_len);
475        self.rollback_storage_logs
476            .truncate(rollback_storage_logs_len);
477    }
478
479    pub(crate) fn delete_history(&mut self) {
480        self.storage_changes.delete_history();
481        self.paid_changes.delete_history();
482        self.transient_storage_changes.delete_history();
483        self.events.delete_history();
484        self.l2_to_l1_logs.delete_history();
485        self.pubdata.delete_history();
486        self.storage_refunds.delete_history();
487        self.pubdata_costs.delete_history();
488        self.decommitted_hashes.delete_history();
489        self.decommit_pinned_pages.delete_history();
490        self.read_storage_slots.delete_history();
491        self.written_storage_slots.delete_history();
492    }
493
494    pub(crate) fn clear_transient_storage(&mut self) {
495        self.transient_storage_changes = RollbackableMap::default();
496    }
497}
498
499/// Opaque snapshot of a [`WorldDiff`] output by its [eponymous method](WorldDiff::snapshot()).
500/// Can be provided to [`WorldDiff::events_after()`] etc. to get data after the snapshot was created.
501#[derive(Clone, PartialEq, Debug)]
502pub struct Snapshot {
503    storage_changes: <RollbackableMap<(H160, U256), U256> as Rollback>::Snapshot,
504    paid_changes: <RollbackableMap<(H160, U256), u32> as Rollback>::Snapshot,
505    events: <RollbackableLog<Event> as Rollback>::Snapshot,
506    l2_to_l1_logs: <RollbackableLog<L2ToL1Log> as Rollback>::Snapshot,
507    transient_storage_changes: <RollbackableMap<(H160, U256), U256> as Rollback>::Snapshot,
508    pubdata: <RollbackablePod<i32> as Rollback>::Snapshot,
509    storage_logs_len: usize,
510    rollback_storage_logs_len: usize,
511}
512
513/// Change in a single storage slot.
514#[derive(Debug, PartialEq)]
515pub struct StorageChange {
516    /// Value before the slot was written to.
517    pub before: U256,
518    /// Value written to the slot.
519    pub after: U256,
520    /// `true` if the slot is not set in the [`World`](crate::World).
521    /// A write may be initial even if it isn't the first write to a slot!
522    pub is_initial: bool,
523}
524
525const WARM_READ_REFUND: u32 = STORAGE_ACCESS_COLD_READ_COST - STORAGE_ACCESS_WARM_READ_COST;
526const WARM_WRITE_REFUND: u32 = STORAGE_ACCESS_COLD_WRITE_COST - STORAGE_ACCESS_WARM_WRITE_COST;
527const COLD_WRITE_AFTER_WARM_READ_REFUND: u32 = STORAGE_ACCESS_COLD_READ_COST;
528
529#[cfg(test)]
530mod tests {
531    use proptest::{bits, collection::btree_map, prelude::*};
532
533    use super::*;
534    use crate::StorageSlot;
535
536    fn test_storage_changes(
537        initial_values: &BTreeMap<(H160, U256), StorageSlot>,
538        first_changes: BTreeMap<(H160, U256), U256>,
539        second_changes: BTreeMap<(H160, U256), U256>,
540    ) {
541        let mut world_diff = WorldDiff {
542            storage_initial_values: initial_values.clone(),
543            ..WorldDiff::default()
544        };
545
546        let checkpoint1 = world_diff.snapshot();
547        for (key, value) in &first_changes {
548            world_diff.write_storage(&mut NoWorld, &mut (), key.0, key.1, *value, 0);
549        }
550        let actual_changes = world_diff
551            .get_storage_changes_after(&checkpoint1)
552            .collect::<BTreeMap<_, _>>();
553        let expected_changes = first_changes
554            .iter()
555            .map(|(key, value)| {
556                let before = initial_values
557                    .get(key)
558                    .map_or_else(U256::zero, |slot| slot.value);
559                let is_initial = initial_values
560                    .get(key)
561                    .is_none_or(|slot| slot.is_write_initial);
562                (
563                    *key,
564                    StorageChange {
565                        before,
566                        after: *value,
567                        is_initial,
568                    },
569                )
570            })
571            .collect();
572        assert_eq!(actual_changes, expected_changes);
573
574        let checkpoint2 = world_diff.snapshot();
575        for (key, value) in &second_changes {
576            world_diff.write_storage(&mut NoWorld, &mut (), key.0, key.1, *value, 0);
577        }
578        let actual_changes = world_diff
579            .get_storage_changes_after(&checkpoint2)
580            .collect::<BTreeMap<_, _>>();
581        let expected_changes = second_changes
582            .iter()
583            .map(|(key, value)| {
584                let before = first_changes
585                    .get(key)
586                    .or(initial_values.get(key).map(|slot| &slot.value))
587                    .copied()
588                    .unwrap_or_default();
589                let is_initial = initial_values
590                    .get(key)
591                    .is_none_or(|slot| slot.is_write_initial);
592                (
593                    *key,
594                    StorageChange {
595                        before,
596                        after: *value,
597                        is_initial,
598                    },
599                )
600            })
601            .collect();
602        assert_eq!(actual_changes, expected_changes);
603
604        let mut combined = first_changes
605            .into_iter()
606            .filter_map(|(key, value)| {
607                let initial = initial_values
608                    .get(&key)
609                    .copied()
610                    .unwrap_or(StorageSlot::EMPTY);
611                (initial.value != value).then_some((
612                    key,
613                    StorageChange {
614                        before: initial.value,
615                        after: value,
616                        is_initial: initial.is_write_initial,
617                    },
618                ))
619            })
620            .collect::<BTreeMap<_, _>>();
621        for (key, value) in second_changes {
622            let initial = initial_values
623                .get(&key)
624                .copied()
625                .unwrap_or(StorageSlot::EMPTY);
626            if initial.value == value {
627                combined.remove(&key);
628            } else {
629                combined.insert(
630                    key,
631                    StorageChange {
632                        before: initial.value,
633                        after: value,
634                        is_initial: initial.is_write_initial,
635                    },
636                );
637            }
638        }
639
640        assert_eq!(combined, world_diff.get_storage_changes().collect());
641    }
642
643    proptest! {
644        #[test]
645        fn storage_changes_work_as_expected(
646            initial_values in arbitrary_initial_storage(),
647            first_changes in arbitrary_storage_changes(),
648            second_changes in arbitrary_storage_changes(),
649        ) {
650            test_storage_changes(&initial_values, first_changes, second_changes);
651        }
652
653        #[test]
654        fn storage_changes_work_with_constrained_changes(
655            initial_values in constrained_initial_storage(),
656            first_changes in constrained_storage_changes(),
657            second_changes in constrained_storage_changes(),
658        ) {
659            test_storage_changes(&initial_values, first_changes, second_changes);
660        }
661    }
662
663    /// Max items in generated initial storage / changes.
664    const MAX_ITEMS: usize = 5;
665    /// Bit mask for bytes in constrained `U256` / `H160` values.
666    const BIT_MASK: u8 = 0b_1111;
667
668    fn arbitrary_initial_storage() -> impl Strategy<Value = BTreeMap<(H160, U256), StorageSlot>> {
669        btree_map(
670            any::<([u8; 20], [u8; 32])>()
671                .prop_map(|(contract, key)| (H160::from(contract), U256::from(key))),
672            any::<([u8; 32], bool)>().prop_map(|(value, is_write_initial)| StorageSlot {
673                value: U256::from(value),
674                is_write_initial,
675            }),
676            0..=MAX_ITEMS,
677        )
678    }
679
680    fn constrained_initial_storage() -> impl Strategy<Value = BTreeMap<(H160, U256), StorageSlot>> {
681        btree_map(
682            (bits::u8::masked(BIT_MASK), bits::u8::masked(BIT_MASK))
683                .prop_map(|(contract, key)| (H160::repeat_byte(contract), U256::from(key))),
684            (bits::u8::masked(BIT_MASK), any::<bool>()).prop_map(|(value, is_write_initial)| {
685                StorageSlot {
686                    value: U256::from(value),
687                    is_write_initial,
688                }
689            }),
690            0..=MAX_ITEMS,
691        )
692    }
693
694    fn arbitrary_storage_changes() -> impl Strategy<Value = BTreeMap<(H160, U256), U256>> {
695        btree_map(
696            any::<([u8; 20], [u8; 32])>()
697                .prop_map(|(contract, key)| (H160::from(contract), U256::from(key))),
698            any::<[u8; 32]>().prop_map(U256::from),
699            0..=MAX_ITEMS,
700        )
701    }
702
703    fn constrained_storage_changes() -> impl Strategy<Value = BTreeMap<(H160, U256), U256>> {
704        btree_map(
705            (bits::u8::masked(BIT_MASK), bits::u8::masked(BIT_MASK))
706                .prop_map(|(contract, key)| (H160::repeat_byte(contract), U256::from(key))),
707            bits::u8::masked(BIT_MASK).prop_map(U256::from),
708            0..=MAX_ITEMS,
709        )
710    }
711
712    struct NoWorld;
713
714    impl StorageInterface for NoWorld {
715        fn read_storage(&mut self, _: H160, _: U256) -> StorageSlot {
716            StorageSlot::EMPTY
717        }
718
719        fn cost_of_writing_storage(&mut self, _: StorageSlot, _: U256) -> u32 {
720            0
721        }
722
723        fn is_free_storage_slot(&self, _: &H160, _: &U256) -> bool {
724            false
725        }
726    }
727
728    #[derive(Default)]
729    struct TestWorld {
730        values: BTreeMap<(H160, U256), U256>,
731    }
732
733    impl StorageInterface for TestWorld {
734        fn read_storage(&mut self, contract: H160, key: U256) -> StorageSlot {
735            let value = self
736                .values
737                .get(&(contract, key))
738                .copied()
739                .unwrap_or_default();
740            StorageSlot {
741                value,
742                is_write_initial: !self.values.contains_key(&(contract, key)),
743            }
744        }
745
746        fn cost_of_writing_storage(&mut self, _: StorageSlot, _: U256) -> u32 {
747            0
748        }
749
750        fn is_free_storage_slot(&self, _: &H160, _: &U256) -> bool {
751            false
752        }
753    }
754
755    #[test]
756    fn storage_logs_include_reads_writes_and_rollbacks() {
757        let mut world_diff = WorldDiff::default();
758        let mut world = TestWorld::default();
759        let contract = H160::zero();
760        let key = U256::from(1);
761
762        let (value, _) = world_diff.read_storage(&mut world, &mut (), contract, key, 0);
763        assert_eq!(value, U256::zero());
764
765        world_diff.write_storage(&mut world, &mut (), contract, key, U256::from(10), 0);
766        let snapshot = world_diff.snapshot();
767        world_diff.write_storage(&mut world, &mut (), contract, key, U256::from(20), 0);
768        world_diff.append_rollback_logs(&snapshot);
769        world_diff.rollback(snapshot);
770
771        let logs = world_diff.storage_log_queries();
772        assert_eq!(logs.len(), 4);
773        assert!(!logs[0].rw_flag);
774        assert!(logs[1].rw_flag && !logs[1].rollback);
775        assert!(logs[2].rw_flag && !logs[2].rollback);
776        assert!(logs[3].rw_flag && logs[3].rollback);
777        assert_eq!(logs[3].read_value, logs[2].read_value);
778        assert_eq!(logs[3].written_value, logs[2].written_value);
779    }
780
781    #[test]
782    fn rollback_without_append_keeps_storage_log_stream_unchanged() {
783        let mut world_diff = WorldDiff::default();
784        let mut world = TestWorld::default();
785        let contract = H160::zero();
786        let key = U256::from(1);
787
788        world_diff.write_storage(&mut world, &mut (), contract, key, U256::from(10), 0);
789        let snapshot = world_diff.snapshot();
790        world_diff.write_storage(&mut world, &mut (), contract, key, U256::from(20), 0);
791        world_diff.rollback(snapshot);
792
793        let logs = world_diff.storage_log_queries();
794        assert_eq!(logs.len(), 2);
795        assert!(logs.iter().all(|log| log.rw_flag && !log.rollback));
796        assert_eq!(world_diff.rollback_storage_logs.len(), 2);
797    }
798
799    #[test]
800    fn external_rollback_truncates_storage_logs_to_internal_snapshot() {
801        let mut world_diff = WorldDiff::default();
802        let mut world = TestWorld::default();
803        let contract = H160::zero();
804        let key = U256::from(1);
805
806        world_diff.write_storage(&mut world, &mut (), contract, key, U256::from(10), 0);
807        let snapshot = world_diff.external_snapshot();
808        world_diff.write_storage(&mut world, &mut (), contract, key, U256::from(20), 0);
809
810        world_diff.external_rollback(snapshot);
811
812        let logs = world_diff.storage_log_queries();
813        assert_eq!(logs.len(), 1);
814        assert!(logs[0].rw_flag && !logs[0].rollback);
815        assert_eq!(world_diff.rollback_storage_logs.len(), 1);
816    }
817
818    #[test]
819    fn storage_read_log_sets_written_value_to_read_value() {
820        let mut world_diff = WorldDiff::default();
821        let mut world = TestWorld::default();
822        let contract = H160::repeat_byte(1);
823        let key = U256::from(7);
824        let value = U256::from(33);
825        world.values.insert((contract, key), value);
826
827        let (read_value, _) = world_diff.read_storage(&mut world, &mut (), contract, key, 0);
828        assert_eq!(read_value, value);
829
830        let logs = world_diff.storage_log_queries();
831        assert_eq!(logs.len(), 1);
832        assert!(!logs[0].rw_flag);
833        assert_eq!(logs[0].read_value, value);
834        assert_eq!(logs[0].written_value, value);
835    }
836}
zksync_vm2/world_diff.rs

zksync_vm2/
world_diff.rs
