Skip to content

Commit

Permalink
Test storage changes more thoroughly
Browse files Browse the repository at this point in the history
  • Loading branch information
@slowli
slowli committed Sep 26, 2024
1 parent ac6d184 commit 9328e72
Showing 1 changed file with 164 additions and 81 deletions.
245 changes: 164 additions & 81 deletions crates/vm2/src/world_diff.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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::<BTreeMap<_, _>>();
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::<BTreeMap<_, _>>();
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::<BTreeMap<_, _>>();
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::<BTreeMap<_, _>>(),
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::<BTreeMap<_, _>>(),
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<Value = BTreeMap<(H160, U256), StorageSlot>> {
any::<Vec<([u8; 20], [u8; 32], [u8; 32], bool)>>().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::<BTreeMap<_, _>>();
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<Value = BTreeMap<(H160, U256), StorageSlot>> {
any::<Vec<(u8, u8, u8, bool)>>().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<Value = BTreeMap<(H160, U256), U256>> {
any::<Vec<(([u8; 20], [u8; 32]), [u8; 32])>>().prop_map(|vec| {
any::<Vec<([u8; 20], [u8; 32], [u8; 32])>>().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<Value = BTreeMap<(H160, U256), U256>> {
any::<Vec<(u8, u8, u8)>>().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 {
Expand Down

0 comments on commit 9328e72

Please sign in to comment.