From 9328e7211fb5a82944e3d32eed606b50590d35e1 Mon Sep 17 00:00:00 2001 From: Alex Ostrovski Date: Thu, 26 Sep 2024 17:12:13 +0300 Subject: [PATCH] Test storage changes more thoroughly --- crates/vm2/src/world_diff.rs | 245 +++++++++++++++++++++++------------ 1 file changed, 164 insertions(+), 81 deletions(-) diff --git a/crates/vm2/src/world_diff.rs b/crates/vm2/src/world_diff.rs index 363a0bc..65fb00e 100644 --- a/crates/vm2/src/world_diff.rs +++ b/crates/vm2/src/world_diff.rs @@ -380,105 +380,188 @@ mod tests { use super::*; use crate::StorageSlot; + fn test_storage_changes( + initial_values: &BTreeMap<(H160, U256), StorageSlot>, + first_changes: BTreeMap<(H160, U256), U256>, + second_changes: BTreeMap<(H160, U256), U256>, + ) { + let mut world_diff = WorldDiff { + storage_initial_values: initial_values.clone(), + ..WorldDiff::default() + }; + + let checkpoint1 = world_diff.snapshot(); + for (key, value) in &first_changes { + world_diff.write_storage(&mut NoWorld, &mut (), key.0, key.1, *value); + } + let actual_changes = world_diff + .get_storage_changes_after(&checkpoint1) + .collect::>(); + let expected_changes = first_changes + .iter() + .map(|(key, value)| { + let before = initial_values + .get(key) + .map_or_else(U256::zero, |slot| slot.value); + let is_initial = initial_values + .get(key) + .map_or(true, |slot| slot.is_write_initial); + ( + *key, + StorageChange { + before, + after: *value, + is_initial, + }, + ) + }) + .collect(); + assert_eq!(actual_changes, expected_changes); + + let checkpoint2 = world_diff.snapshot(); + for (key, value) in &second_changes { + world_diff.write_storage(&mut NoWorld, &mut (), key.0, key.1, *value); + } + let actual_changes = world_diff + .get_storage_changes_after(&checkpoint2) + .collect::>(); + let expected_changes = second_changes + .iter() + .map(|(key, value)| { + let before = first_changes + .get(key) + .or(initial_values.get(key).map(|slot| &slot.value)) + .copied() + .unwrap_or_default(); + let is_initial = initial_values + .get(key) + .map_or(true, |slot| slot.is_write_initial); + ( + *key, + StorageChange { + before, + after: *value, + is_initial, + }, + ) + }) + .collect(); + assert_eq!(actual_changes, expected_changes); + + let mut combined = first_changes + .into_iter() + .filter_map(|(key, value)| { + let initial = initial_values + .get(&key) + .copied() + .unwrap_or(StorageSlot::EMPTY); + (initial.value != value).then_some(( + key, + StorageChange { + before: initial.value, + after: value, + is_initial: initial.is_write_initial, + }, + )) + }) + .collect::>(); + for (key, value) in second_changes { + let initial = initial_values + .get(&key) + .copied() + .unwrap_or(StorageSlot::EMPTY); + if initial.value == value { + combined.remove(&key); + } else { + combined.insert( + key, + StorageChange { + before: initial.value, + after: value, + is_initial: initial.is_write_initial, + }, + ); + } + } + + assert_eq!(combined, world_diff.get_storage_changes().collect()); + } + proptest! { #[test] - fn test_storage_changes( - initial_values in arbitrary_storage_changes(), + fn storage_changes_work_as_expected( + initial_values in arbitrary_initial_storage(), first_changes in arbitrary_storage_changes(), second_changes in arbitrary_storage_changes(), ) { - let storage_initial_values = initial_values - .iter() - .map(|(key, &value)| (*key, StorageSlot { - value, - is_write_initial: false, - })) - .collect(); - let mut world_diff = WorldDiff { - storage_initial_values, - ..WorldDiff::default() - }; - - let checkpoint1 = world_diff.snapshot(); - for (key, value) in &first_changes { - world_diff.write_storage(&mut NoWorld, &mut (), key.0, key.1, *value); - } - assert_eq!( - world_diff - .get_storage_changes_after(&checkpoint1) - .collect::>(), - first_changes - .iter() - .map(|(key, value)| ( - *key, - StorageChange { - before: initial_values.get(key).copied().unwrap_or_default(), - after: *value, - is_initial: !initial_values.contains_key(key), - } - )) - .collect() - ); + test_storage_changes(&initial_values, first_changes, second_changes); + } - let checkpoint2 = world_diff.snapshot(); - for (key, value) in &second_changes { - world_diff.write_storage(&mut NoWorld, &mut (), key.0, key.1, *value); - } - assert_eq!( - world_diff - .get_storage_changes_after(&checkpoint2) - .collect::>(), - second_changes - .iter() - .map(|(key, value)| ( - *key, - StorageChange { - before: first_changes.get(key).or(initial_values.get(key)).copied().unwrap_or_default(), - after: *value, - is_initial: !initial_values.contains_key(key), - } - )) - .collect() - ); - - let mut combined = first_changes - .into_iter() - .filter_map(|(key, value)| { - let initial = initial_values.get(&key).copied(); - (initial.unwrap_or_default() != value).then_some((key, StorageChange { - before: initial.unwrap_or_default(), - after: value, - is_initial: initial.is_none(), - })) + #[test] + fn storage_changes_work_with_constrained_changes( + initial_values in constrained_initial_storage(), + first_changes in constrained_storage_changes(), + second_changes in constrained_storage_changes(), + ) { + test_storage_changes(&initial_values, first_changes, second_changes); + } + } + + fn arbitrary_initial_storage() -> impl Strategy> { + any::>().prop_map(|vec| { + vec.into_iter() + .map(|(contract, key, value, is_write_initial)| { + ( + (H160::from(contract), U256::from(key)), + StorageSlot { + value: U256::from(value), + is_write_initial, + }, + ) }) - .collect::>(); - for (key, value) in second_changes { - let initial = initial_values.get(&key).copied(); - if initial.unwrap_or_default() == value { - combined.remove(&key); - } else { - combined.insert(key, StorageChange { - before: initial.unwrap_or_default(), - after: value, - is_initial: initial.is_none(), - }); - } - } + .collect() + }) + } - assert_eq!(combined, world_diff.get_storage_changes().collect()); - } + fn constrained_initial_storage() -> impl Strategy> { + any::>().prop_map(|vec| { + vec.into_iter() + .map(|(contract, key, value, is_write_initial)| { + ( + (H160::repeat_byte(contract), U256::from(key)), + StorageSlot { + value: U256::from(value), + is_write_initial, + }, + ) + }) + .collect() + }) } fn arbitrary_storage_changes() -> impl Strategy> { - any::>().prop_map(|vec| { + any::>().prop_map(|vec| { vec.into_iter() - .map(|((contract, key), value)| { + .map(|(contract, key, value)| { ((H160::from(contract), U256::from(key)), U256::from(value)) }) .collect() }) } + fn constrained_storage_changes() -> impl Strategy> { + any::>().prop_map(|vec| { + vec.into_iter() + .map(|(contract, key, value)| { + ( + (H160::repeat_byte(contract), U256::from(key)), + U256::from(value), + ) + }) + .collect() + }) + } + struct NoWorld; impl StorageInterface for NoWorld {