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 d0d83a490c..3fdf556e9e 100644 --- a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs +++ b/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs @@ -90,8 +90,8 @@ pub(super) struct SparseTrieCacheTask, /// Whether the last state update has been received. finished_state_updates: bool, - /// Pending targets to be dispatched to the proof workers. - pending_targets: MultiProofTargetsV2, + /// Pending proof targets queued for dispatch to proof workers. + pending_targets: PendingTargets, /// Number of pending execution/prewarming updates received but not yet passed to /// `update_leaves`. pending_updates: usize, @@ -283,7 +283,7 @@ where // them to the trie, self.process_new_updates()?; self.dispatch_pending_targets(); - } else if self.pending_targets.chunking_length() > self.chunk_size.unwrap_or_default() { + } else if self.pending_targets.len() > self.chunk_size.unwrap_or_default() { // Make sure to dispatch targets if we've accumulated a lot of them. self.dispatch_pending_targets(); } @@ -505,7 +505,7 @@ where self.trie.insert_storage_trie(*address, trie); if !targets.is_empty() { - self.pending_targets.storage_targets.entry(*address).or_default().extend(targets); + self.pending_targets.extend_storage_targets(address, targets); } } @@ -535,15 +535,13 @@ where if min_len < *entry.get() { entry.insert(min_len); self.pending_targets - .account_targets - .push(Target::new(target).with_min_len(min_len)); + .push_account_target(Target::new(target).with_min_len(min_len)); } } Entry::Vacant(entry) => { entry.insert(min_len); self.pending_targets - .account_targets - .push(Target::new(target).with_min_len(min_len)); + .push_account_target(Target::new(target).with_min_len(min_len)); } } })?; @@ -653,9 +651,9 @@ where )] fn dispatch_pending_targets(&mut self) { if !self.pending_targets.is_empty() { - let chunking_length = self.pending_targets.chunking_length(); + let (targets, chunking_length) = self.pending_targets.take(); dispatch_with_chunking( - std::mem::take(&mut self.pending_targets), + targets, chunking_length, self.chunk_size, self.max_targets_for_chunking, @@ -697,6 +695,44 @@ fn encode_account_leaf_value( account_rlp_buf.clone() } +/// Pending proof targets queued for dispatch to proof workers, along with their count. +#[derive(Default)] +struct PendingTargets { + /// The proof targets. + targets: MultiProofTargetsV2, + /// Number of account + storage proof targets currently queued. + len: usize, +} + +impl PendingTargets { + /// Returns the number of pending targets. + const fn len(&self) -> usize { + self.len + } + + /// Returns `true` if there are no pending targets. + const fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Takes the pending targets, replacing with empty defaults. + fn take(&mut self) -> (MultiProofTargetsV2, usize) { + (std::mem::take(&mut self.targets), std::mem::take(&mut self.len)) + } + + /// Adds a target to the account targets. + fn push_account_target(&mut self, target: Target) { + self.targets.account_targets.push(target); + self.len += 1; + } + + /// Extends storage targets for the given address. + fn extend_storage_targets(&mut self, address: &B256, targets: Vec) { + self.len += targets.len(); + self.targets.storage_targets.entry(*address).or_default().extend(targets); + } +} + /// Message type for the sparse trie task. enum SparseTrieTaskMessage { /// A hashed state update ready to be processed.