diff --git a/crates/engine/tree/src/tree/payload_processor/multiproof.rs b/crates/engine/tree/src/tree/payload_processor/multiproof.rs index 321de725bec..aa850fefe4c 100644 --- a/crates/engine/tree/src/tree/payload_processor/multiproof.rs +++ b/crates/engine/tree/src/tree/payload_processor/multiproof.rs @@ -732,7 +732,12 @@ impl MultiProofTask { /// Handles request for proof prefetch. /// /// Returns a number of proofs that were spawned. - #[instrument(level = "debug", target = "engine::tree::payload_processor::multiproof", skip_all, fields(accounts = targets.len()))] + #[instrument( + level = "debug", + target = "engine::tree::payload_processor::multiproof", + skip_all, + fields(accounts = targets.len(), chunks = 0) + )] fn on_prefetch_proof(&mut self, targets: MultiProofTargets) -> u64 { let proof_targets = self.get_prefetch_proof_targets(targets); self.fetched_proof_targets.extend_ref(&proof_targets); @@ -774,10 +779,16 @@ impl MultiProofTask { chunks += 1; }; - if should_chunk && let Some(chunk_size) = self.chunk_size { + if should_chunk && + let Some(chunk_size) = self.chunk_size && + proof_targets.chunking_length() > chunk_size + { + let mut chunks = 0usize; for proof_targets_chunk in proof_targets.chunks(chunk_size) { dispatch(proof_targets_chunk); + chunks += 1; } + tracing::Span::current().record("chunks", chunks); } else { dispatch(proof_targets); } @@ -863,7 +874,12 @@ impl MultiProofTask { /// Handles state updates. /// /// Returns a number of proofs that were spawned. - #[instrument(level = "debug", target = "engine::tree::payload_processor::multiproof", skip(self, update), fields(accounts = update.len()))] + #[instrument( + level = "debug", + target = "engine::tree::payload_processor::multiproof", + skip(self, update), + fields(accounts = update.len(), chunks = 0) + )] fn on_state_update(&mut self, source: StateChangeSource, update: EvmState) -> u64 { let hashed_state_update = evm_state_to_hashed_post_state(update); @@ -923,10 +939,16 @@ impl MultiProofTask { chunks += 1; }; - if should_chunk && let Some(chunk_size) = self.chunk_size { + if should_chunk && + let Some(chunk_size) = self.chunk_size && + not_fetched_state_update.chunking_length() > chunk_size + { + let mut chunks = 0usize; for chunk in not_fetched_state_update.chunks(chunk_size) { dispatch(chunk); + chunks += 1; } + tracing::Span::current().record("chunks", chunks); } else { dispatch(not_fetched_state_update); } diff --git a/crates/trie/common/src/hashed_state.rs b/crates/trie/common/src/hashed_state.rs index 8fb994daddd..4bb59bb00b8 100644 --- a/crates/trie/common/src/hashed_state.rs +++ b/crates/trie/common/src/hashed_state.rs @@ -278,6 +278,15 @@ impl HashedPostState { ChunkedHashedPostState::new(self, size) } + /// Returns the number of items that will be considered during chunking in `[Self::chunks]`. + pub fn chunking_length(&self) -> usize { + self.accounts.len() + + self.storages + .values() + .map(|storage| if storage.wiped { 1 } else { 0 } + storage.storage.len()) + .sum::() + } + /// Extend this hashed post state with contents of another. /// Entries in the second hashed post state take precedence. pub fn extend(&mut self, other: Self) { @@ -1239,4 +1248,65 @@ mod tests { assert_eq!(storage3.zero_valued_slots.len(), 1); assert!(storage3.zero_valued_slots.contains(&B256::from([4; 32]))); } + + #[test] + fn test_hashed_post_state_chunking_length() { + let addr1 = B256::from([1; 32]); + let addr2 = B256::from([2; 32]); + let addr3 = B256::from([3; 32]); + let addr4 = B256::from([4; 32]); + let slot1 = B256::from([1; 32]); + let slot2 = B256::from([2; 32]); + let slot3 = B256::from([3; 32]); + + let state = HashedPostState { + accounts: B256Map::from_iter([(addr1, None), (addr2, None), (addr4, None)]), + storages: B256Map::from_iter([ + ( + addr1, + HashedStorage { + wiped: false, + storage: B256Map::from_iter([ + (slot1, U256::ZERO), + (slot2, U256::ZERO), + (slot3, U256::ZERO), + ]), + }, + ), + ( + addr2, + HashedStorage { + wiped: true, + storage: B256Map::from_iter([ + (slot1, U256::ZERO), + (slot2, U256::ZERO), + (slot3, U256::ZERO), + ]), + }, + ), + ( + addr3, + HashedStorage { + wiped: false, + storage: B256Map::from_iter([ + (slot1, U256::ZERO), + (slot2, U256::ZERO), + (slot3, U256::ZERO), + ]), + }, + ), + ]), + }; + + let chunking_length = state.chunking_length(); + for size in 1..=state.clone().chunks(1).count() { + let chunk_count = state.clone().chunks(size).count(); + let expected_count = chunking_length.div_ceil(size); + assert_eq!( + chunk_count, expected_count, + "chunking_length: {}, size: {}", + chunking_length, size + ); + } + } } diff --git a/crates/trie/common/src/proofs.rs b/crates/trie/common/src/proofs.rs index b7961f047a4..a8e0bb59b93 100644 --- a/crates/trie/common/src/proofs.rs +++ b/crates/trie/common/src/proofs.rs @@ -89,6 +89,11 @@ impl MultiProofTargets { pub fn chunks(self, size: usize) -> ChunkedMultiProofTargets { ChunkedMultiProofTargets::new(self, size) } + + /// Returns the number of items that will be considered during chunking in `[Self::chunks]`. + pub fn chunking_length(&self) -> usize { + self.values().map(|slots| 1 + slots.len().saturating_sub(1)).sum::() + } } /// An iterator that yields chunks of the proof targets of at most `size` account and storage @@ -1067,4 +1072,33 @@ mod tests { acc.storage_root = EMPTY_ROOT_HASH; assert_eq!(acc, inverse); } + + #[test] + fn test_multiproof_targets_chunking_length() { + let mut targets = MultiProofTargets::default(); + targets.insert(B256::with_last_byte(1), B256Set::default()); + targets.insert( + B256::with_last_byte(2), + B256Set::from_iter([B256::with_last_byte(10), B256::with_last_byte(20)]), + ); + targets.insert( + B256::with_last_byte(3), + B256Set::from_iter([ + B256::with_last_byte(30), + B256::with_last_byte(31), + B256::with_last_byte(32), + ]), + ); + + let chunking_length = targets.chunking_length(); + for size in 1..=targets.clone().chunks(1).count() { + let chunk_count = targets.clone().chunks(size).count(); + let expected_count = chunking_length.div_ceil(size); + assert_eq!( + chunk_count, expected_count, + "chunking_length: {}, size: {}", + chunking_length, size + ); + } + } }