diff --git a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs b/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs index a3bd37a356..905c2da435 100644 --- a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs +++ b/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs @@ -236,6 +236,11 @@ pub(super) struct SparseTrieCacheTask, /// Storage trie updates. hashed address -> slot -> update. storage_updates: B256Map>, + + /// Account updates that are buffered but were not yet applied to the trie. + new_account_updates: B256Map, + /// Storage updates that are buffered but were not yet applied to the trie. + new_storage_updates: B256Map>, /// Account updates that are blocked by storage root calculation or account reveal. /// /// Those are being moved into `account_updates` once storage roots @@ -294,6 +299,8 @@ where max_targets_for_chunking: DEFAULT_MAX_TARGETS_FOR_CHUNKING, account_updates: Default::default(), storage_updates: Default::default(), + new_account_updates: Default::default(), + new_storage_updates: Default::default(), pending_account_updates: Default::default(), fetched_account_targets: Default::default(), fetched_storage_targets: Default::default(), @@ -389,18 +396,16 @@ where // If we don't have any pending messages, we can spend some time on computing // storage roots and promoting account updates. self.dispatch_pending_targets(); + self.process_new_updates()?; self.promote_pending_account_updates()?; self.dispatch_pending_targets(); } else if self.updates.is_empty() || self.pending_updates > MAX_PENDING_UPDATES { // If we don't have any pending updates OR we've accumulated a lot already, apply // them to the trie, - self.process_leaf_updates()?; + self.process_new_updates()?; self.dispatch_pending_targets(); - } else if self.updates.is_empty() || - self.pending_targets.chunking_length() > self.chunk_size.unwrap_or_default() - { - // Make sure to dispatch targets if we don't have any updates or if we've - // accumulated a lot of them. + } else if self.pending_targets.chunking_length() > self.chunk_size.unwrap_or_default() { + // Make sure to dispatch targets if we've accumulated a lot of them. self.dispatch_pending_targets(); } @@ -450,13 +455,13 @@ where for target in targets.account_targets { // Only touch accounts that are not yet present in the updates set. - self.account_updates.entry(target.key()).or_insert(LeafUpdate::Touched); + self.new_account_updates.entry(target.key()).or_insert(LeafUpdate::Touched); } for (address, slots) in targets.storage_targets { for slot in slots { // Only touch storages that are not yet present in the updates set. - self.storage_updates + self.new_storage_updates .entry(address) .or_default() .entry(slot.key()) @@ -465,7 +470,7 @@ where // Touch corresponding account leaf to make sure its revealed in accounts trie for // storage root update. - self.account_updates.entry(address).or_insert(LeafUpdate::Touched); + self.new_account_updates.entry(address).or_insert(LeafUpdate::Touched); } } @@ -486,15 +491,18 @@ where } else { alloy_rlp::encode_fixed_size(&value).to_vec() }; - self.storage_updates + self.new_storage_updates .entry(address) .or_default() .insert(slot, LeafUpdate::Changed(encoded)); + + // Remove an existing storage update if it exists. + self.storage_updates.get_mut(&address).and_then(|updates| updates.remove(&slot)); } // Make sure account is tracked in `account_updates` so that it is revealed in accounts // trie for storage root update. - self.account_updates.entry(address).or_insert(LeafUpdate::Touched); + self.new_account_updates.entry(address).or_insert(LeafUpdate::Touched); // Make sure account is tracked in `pending_account_updates` so that once storage root // is computed, it will be updated in the accounts trie. @@ -506,7 +514,7 @@ where // // This might overwrite an existing update, which is fine, because storage root from it // is already tracked in the trie and can be easily fetched again. - self.account_updates.insert(address, LeafUpdate::Touched); + self.new_account_updates.insert(address, LeafUpdate::Touched); // Track account in `pending_account_updates` so that once storage root is computed, // it will be updated in the accounts trie. @@ -523,6 +531,50 @@ where }) } + #[instrument( + level = "debug", + target = "engine::tree::payload_processor::sparse_trie", + skip_all + )] + fn process_new_updates(&mut self) -> SparseTrieResult<()> { + self.pending_updates = 0; + + // Firstly apply all new storage and account updates to the tries. + self.process_leaf_updates(true)?; + + for (address, mut new) in self.new_storage_updates.drain() { + let updates = self.storage_updates.entry(address).or_default(); + for (slot, new) in new.drain() { + match updates.entry(slot) { + Entry::Occupied(mut entry) => { + // Only overwrite existing entries with new values + if new.is_changed() { + entry.insert(new); + } + } + Entry::Vacant(entry) => { + entry.insert(new); + } + } + } + } + + for (address, new) in self.new_account_updates.drain() { + match self.account_updates.entry(address) { + Entry::Occupied(mut entry) => { + if new.is_changed() { + entry.insert(new); + } + } + Entry::Vacant(entry) => { + entry.insert(new); + } + } + } + + Ok(()) + } + /// Applies all account and storage leaf updates to corresponding tries and collects any new /// multiproof targets. #[instrument( @@ -530,12 +582,12 @@ where target = "engine::tree::payload_processor::sparse_trie", skip_all )] - fn process_leaf_updates(&mut self) -> SparseTrieResult<()> { - self.pending_updates = 0; + fn process_leaf_updates(&mut self, new: bool) -> SparseTrieResult<()> { + let storage_updates = + if new { &mut self.new_storage_updates } else { &mut self.storage_updates }; // Process all storage updates in parallel, skipping tries with no pending updates. - let storage_results = self - .storage_updates + let storage_results = storage_updates .iter_mut() .filter(|(_, updates)| !updates.is_empty()) .map(|(address, updates)| { @@ -575,7 +627,7 @@ where } // Process account trie updates and fill the account targets. - self.process_account_leaf_updates()?; + self.process_account_leaf_updates(new)?; Ok(()) } @@ -583,12 +635,14 @@ where /// Invokes `update_leaves` for the accounts trie and collects any new targets. /// /// Returns whether any updates were drained (applied to the trie). - fn process_account_leaf_updates(&mut self) -> SparseTrieResult { - let updates_len_before = self.account_updates.len(); + fn process_account_leaf_updates(&mut self, new: bool) -> SparseTrieResult { + let account_updates = + if new { &mut self.new_account_updates } else { &mut self.account_updates }; - self.trie.trie_mut().update_leaves( - &mut self.account_updates, - |target, min_len| match self.fetched_account_targets.entry(target) { + let updates_len_before = account_updates.len(); + + self.trie.trie_mut().update_leaves(account_updates, |target, min_len| { + match self.fetched_account_targets.entry(target) { Entry::Occupied(mut entry) => { if min_len < *entry.get() { entry.insert(min_len); @@ -603,10 +657,10 @@ where .account_targets .push(Target::new(target).with_min_len(min_len)); } - }, - )?; + } + })?; - Ok(self.account_updates.len() < updates_len_before) + Ok(account_updates.len() < updates_len_before) } /// Iterates through all storage tries for which all updates were processed, computes their @@ -618,7 +672,7 @@ where skip_all )] fn promote_pending_account_updates(&mut self) -> SparseTrieResult<()> { - self.process_leaf_updates()?; + self.process_leaf_updates(false)?; if self.pending_account_updates.is_empty() { return Ok(()); @@ -711,7 +765,7 @@ where // // We need to keep iterating if any updates are being drained because that might // indicate that more pending account updates can be promoted. - if !self.process_account_leaf_updates()? { + if !self.process_account_leaf_updates(false)? { break } } diff --git a/crates/trie/sparse/src/traits.rs b/crates/trie/sparse/src/traits.rs index 2f0f9b3351..915f38434a 100644 --- a/crates/trie/sparse/src/traits.rs +++ b/crates/trie/sparse/src/traits.rs @@ -24,6 +24,18 @@ pub enum LeafUpdate { Touched, } +impl LeafUpdate { + /// Returns true if the leaf update is a change. + pub const fn is_changed(&self) -> bool { + matches!(self, Self::Changed(_)) + } + + /// Returns true if the leaf update is a touched update. + pub const fn is_touched(&self) -> bool { + matches!(self, Self::Touched) + } +} + /// Trait defining common operations for revealed sparse trie implementations. /// /// This trait abstracts over different sparse trie implementations (serial vs parallel)