From bddd3202e4603251d76ecceb6bfacf40b26ec46e Mon Sep 17 00:00:00 2001 From: Alexey Shekhirin Date: Mon, 21 Oct 2024 20:04:09 +0100 Subject: [PATCH] test(trie): narrow the range of keys for sparse trie fuzz (#11937) --- crates/trie/sparse/src/trie.rs | 73 ++++++++++++++++++++-------------- 1 file changed, 44 insertions(+), 29 deletions(-) diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 0b9ffb5c0..6680d7e9b 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -265,15 +265,15 @@ impl RevealedSparseTrie { } /// Remove leaf node from the trie. - pub fn remove_leaf(&mut self, path: Nibbles) -> SparseTrieResult<()> { + pub fn remove_leaf(&mut self, path: &Nibbles) -> SparseTrieResult<()> { self.prefix_set.insert(path.clone()); - let existing = self.values.remove(&path); + let existing = self.values.remove(path); if existing.is_none() { // trie structure unchanged, return immediately return Ok(()) } - let mut removed_nodes = self.take_nodes_for_path(&path)?; + let mut removed_nodes = self.take_nodes_for_path(path)?; debug!(target: "trie::sparse", ?path, ?removed_nodes, "Removed nodes for path"); // Pop the first node from the stack which is the leaf node we want to remove. let mut child = removed_nodes.pop().expect("leaf exists"); @@ -282,7 +282,7 @@ impl RevealedSparseTrie { let mut child_path = child.path.clone(); let SparseNode::Leaf { key, .. } = &child.node else { panic!("expected leaf node") }; child_path.extend_from_slice_unchecked(key); - assert_eq!(child_path, path); + assert_eq!(&child_path, path); } // If we don't have any other removed nodes, insert an empty node at the root. @@ -727,6 +727,7 @@ mod tests { use super::*; use alloy_primitives::U256; use itertools::Itertools; + use prop::sample::SizeRange; use proptest::prelude::*; use rand::seq::IteratorRandom; use reth_testing_utils::generators; @@ -998,7 +999,7 @@ mod tests { ]) ); - sparse.remove_leaf(Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3])).unwrap(); + sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3])).unwrap(); // Extension (Key = 5) // └── Branch (Mask = 1001) @@ -1049,7 +1050,7 @@ mod tests { ]) ); - sparse.remove_leaf(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1])).unwrap(); + sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1])).unwrap(); // Extension (Key = 5) // └── Branch (Mask = 1001) @@ -1085,7 +1086,7 @@ mod tests { ]) ); - sparse.remove_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2])).unwrap(); + sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2])).unwrap(); // Extension (Key = 5) // └── Branch (Mask = 1001) @@ -1118,7 +1119,7 @@ mod tests { ]) ); - sparse.remove_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0])).unwrap(); + sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0])).unwrap(); // Extension (Key = 5) // └── Branch (Mask = 1001) @@ -1140,7 +1141,7 @@ mod tests { ]) ); - sparse.remove_leaf(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3])).unwrap(); + sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3])).unwrap(); // Leaf (Key = 53302) pretty_assertions::assert_eq!( @@ -1151,7 +1152,7 @@ mod tests { ),]) ); - sparse.remove_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2])).unwrap(); + sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2])).unwrap(); // Empty pretty_assertions::assert_eq!( @@ -1162,34 +1163,31 @@ mod tests { #[test] fn sparse_trie_fuzz() { - proptest!(ProptestConfig::with_cases(10), |(updates: Vec>)| { + // Having only the first 3 nibbles set, we narrow down the range of keys + // to 4096 different hashes. It allows us to generate collisions more likely + // to test the sparse trie updates. + const KEY_NIBBLES_LEN: usize = 3; + + fn test(updates: Vec>>) { let mut rng = generators::rng(); let mut state = BTreeMap::default(); - let mut unpacked_state = BTreeMap::default(); let mut sparse = RevealedSparseTrie::default(); for update in updates { let keys_to_delete_len = update.len() / 2; - let unpacked_update = update.iter().map(|(key, value)| ( - Nibbles::unpack(key), - alloy_rlp::encode_fixed_size(value).to_vec() - )); - // Insert state updates into the sparse trie and calculate the root - for (key, value) in unpacked_update.clone() { + for (key, value) in update.clone() { sparse.update_leaf(key, value).unwrap(); } let sparse_root = sparse.root(); // Insert state updates into the hash builder and calculate the root - unpacked_state.extend(unpacked_update); state.extend(update); - let keys = state.keys().map(Nibbles::unpack).collect::>(); let (hash_builder_root, hash_builder_proof_nodes) = hash_builder_root_with_proofs( - unpacked_state.clone(), - keys, + state.clone(), + state.keys().cloned().collect::>(), ); // Assert that the sparse trie root matches the hash builder root @@ -1204,20 +1202,18 @@ mod tests { .keys() .choose_multiple(&mut rng, keys_to_delete_len) .into_iter() - .copied() + .cloned() .collect::>(); for key in keys_to_delete { state.remove(&key).unwrap(); - unpacked_state.remove(&Nibbles::unpack(key)).unwrap(); - sparse.remove_leaf(Nibbles::unpack(key)).unwrap(); + sparse.remove_leaf(&key).unwrap(); } let sparse_root = sparse.root(); - let keys = state.keys().map(Nibbles::unpack).collect::>(); let (hash_builder_root, hash_builder_proof_nodes) = hash_builder_root_with_proofs( - unpacked_state.clone(), - keys, + state.clone(), + state.keys().cloned().collect::>(), ); // Assert that the sparse trie root matches the hash builder root @@ -1225,6 +1221,25 @@ mod tests { // Assert that the sparse trie nodes match the hash builder proof nodes assert_eq_sparse_trie_proof_nodes(&sparse, hash_builder_proof_nodes); } - }); + } + + /// Pad nibbles of length [`KEY_NIBBLES_LEN`] with zeros to the length of a B256 hash. + fn pad_nibbles(nibbles: Nibbles) -> Nibbles { + let mut base = + Nibbles::from_nibbles_unchecked([0; { B256::len_bytes() / 2 - KEY_NIBBLES_LEN }]); + base.extend_from_slice_unchecked(&nibbles); + base + } + + proptest!(ProptestConfig::with_cases(10), |( + updates in proptest::collection::vec( + proptest::collection::hash_map( + any_with::(SizeRange::new(KEY_NIBBLES_LEN..=KEY_NIBBLES_LEN)).prop_map(pad_nibbles), + any::>(), + 1..100, + ), + 1..100, + ) + )| { test(updates) }); } }