diff --git a/crates/engine/tree/Cargo.toml b/crates/engine/tree/Cargo.toml index a651843c38ee..01d7e7e20241 100644 --- a/crates/engine/tree/Cargo.toml +++ b/crates/engine/tree/Cargo.toml @@ -120,6 +120,5 @@ test-utils = [ "reth-tracing", "reth-trie/test-utils", "reth-prune-types?/test-utils", - "reth-primitives-traits/test-utils", "reth-trie-db/test-utils", ] diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs index 64e82f05a75d..32bfbf68604b 100644 --- a/crates/engine/tree/src/tree/root.rs +++ b/crates/engine/tree/src/tree/root.rs @@ -111,8 +111,10 @@ pub(crate) enum StateRootMessage { /// Handle to track proof calculation ordering #[derive(Debug, Default)] pub(crate) struct ProofSequencer { - /// The next expected proof sequence number + /// The next proof sequence number to be produced. next_sequence: u64, + /// The next sequence number expected to be delivered. + next_to_deliver: u64, /// Buffer for out-of-order proofs pending_proofs: BTreeMap, } @@ -132,19 +134,31 @@ impl ProofSequencer { /// Adds a proof and returns all sequential proofs if we have a continuous sequence pub(crate) fn add_proof(&mut self, sequence: u64, proof: MultiProof) -> Vec { - if sequence < self.next_sequence { - return vec![proof]; + if sequence >= self.next_to_deliver { + self.pending_proofs.insert(sequence, proof); } - // Insert the new proof into pending proofs - self.pending_proofs.insert(sequence, proof); + // return early if we don't have the next expected proof + if !self.pending_proofs.contains_key(&self.next_to_deliver) { + return Vec::new() + } let mut consecutive_proofs = Vec::with_capacity(self.pending_proofs.len()); + let mut current_sequence = self.next_to_deliver; - // Keep taking proofs from pending_proofs as long as they form a consecutive sequence - while let Some(proof) = self.pending_proofs.remove(&self.next_sequence) { + // keep collecting proofs as long as we have consecutive sequence numbers + while let Some(proof) = self.pending_proofs.remove(¤t_sequence) { consecutive_proofs.push(proof); - self.next_sequence += 1; + current_sequence += 1; + + // if we don't have the next number, stop collecting + if !self.pending_proofs.contains_key(¤t_sequence) { + break; + } + } + + if !consecutive_proofs.is_empty() { + self.next_to_deliver += consecutive_proofs.len() as u64; } consecutive_proofs @@ -844,6 +858,7 @@ mod tests { let mut sequencer = ProofSequencer::new(); let proof1 = MultiProof::default(); let proof2 = MultiProof::default(); + sequencer.next_sequence = 2; let ready = sequencer.add_proof(0, proof1); assert_eq!(ready.len(), 1); @@ -860,6 +875,7 @@ mod tests { let proof1 = MultiProof::default(); let proof2 = MultiProof::default(); let proof3 = MultiProof::default(); + sequencer.next_sequence = 3; let ready = sequencer.add_proof(2, proof3); assert_eq!(ready.len(), 0); @@ -879,6 +895,7 @@ mod tests { let mut sequencer = ProofSequencer::new(); let proof1 = MultiProof::default(); let proof3 = MultiProof::default(); + sequencer.next_sequence = 3; let ready = sequencer.add_proof(0, proof1); assert_eq!(ready.len(), 1); @@ -898,7 +915,7 @@ mod tests { assert_eq!(ready.len(), 1); let ready = sequencer.add_proof(0, proof2); - assert_eq!(ready.len(), 1); + assert_eq!(ready.len(), 0); assert!(!sequencer.has_pending()); } @@ -906,6 +923,7 @@ mod tests { fn test_add_proof_batch_processing() { let mut sequencer = ProofSequencer::new(); let proofs: Vec<_> = (0..5).map(|_| MultiProof::default()).collect(); + sequencer.next_sequence = 5; sequencer.add_proof(4, proofs[4].clone()); sequencer.add_proof(2, proofs[2].clone());