fix(trie): prefix set extension (#11605)

This commit is contained in:
Roman Krasiuk
2024-10-09 13:20:58 +02:00
committed by GitHub
parent 3c4a94fb76
commit 5ad6484386
2 changed files with 19 additions and 6 deletions

View File

@@ -20,9 +20,9 @@ pub struct TriePrefixSetsMut {
impl TriePrefixSetsMut {
/// Extends prefix sets with contents of another prefix set.
pub fn extend(&mut self, other: Self) {
self.account_prefix_set.extend(other.account_prefix_set.keys);
self.account_prefix_set.extend(other.account_prefix_set);
for (hashed_address, prefix_set) in other.storage_prefix_sets {
self.storage_prefix_sets.entry(hashed_address).or_default().extend(prefix_set.keys);
self.storage_prefix_sets.entry(hashed_address).or_default().extend(prefix_set);
}
self.destroyed_accounts.extend(other.destroyed_accounts);
}
@@ -115,12 +115,18 @@ impl PrefixSetMut {
self.keys.push(nibbles);
}
/// Extend prefix set with contents of another prefix set.
pub fn extend(&mut self, other: Self) {
self.all |= other.all;
self.keys.extend(other.keys);
}
/// Extend prefix set keys with contents of provided iterator.
pub fn extend<I>(&mut self, nibbles_iter: I)
pub fn extend_keys<I>(&mut self, keys: I)
where
I: IntoIterator<Item = Nibbles>,
{
self.keys.extend(nibbles_iter);
self.keys.extend(keys);
}
/// Returns the number of elements in the set.
@@ -270,4 +276,11 @@ mod tests {
assert_eq!(prefix_set.keys.len(), 3); // Length should be 3 (excluding duplicate)
assert_eq!(prefix_set.keys.capacity(), 3); // Capacity should be 3 after shrinking
}
#[test]
fn test_prefix_set_all_extend() {
let mut prefix_set_mut = PrefixSetMut::default();
prefix_set_mut.extend(PrefixSetMut::all());
assert!(prefix_set_mut.all);
}
}

View File

@@ -96,7 +96,7 @@ where
// Create the walker.
let mut prefix_set = self.prefix_sets.account_prefix_set.clone();
prefix_set.extend(targets.keys().map(Nibbles::unpack));
prefix_set.extend_keys(targets.keys().map(Nibbles::unpack));
let walker = TrieWalker::new(trie_cursor, prefix_set.freeze());
// Create a hash builder to rebuild the root node since it is not available in the database.
@@ -225,7 +225,7 @@ where
}
let target_nibbles = targets.into_iter().map(Nibbles::unpack).collect::<Vec<_>>();
self.prefix_set.extend(target_nibbles.clone());
self.prefix_set.extend_keys(target_nibbles.clone());
let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?;
let walker = TrieWalker::new(trie_cursor, self.prefix_set.freeze());