Skip to main content

reth_stages/stages/
merkle.rs

1use alloy_consensus::BlockHeader;
2use alloy_primitives::{BlockNumber, Sealable, B256};
3use reth_codecs::Compact;
4use reth_consensus::ConsensusError;
5use reth_db_api::{
6    tables,
7    transaction::{DbTx, DbTxMut},
8};
9use reth_primitives_traits::{GotExpected, SealedHeader};
10use reth_provider::{
11    ChangeSetReader, DBProvider, HeaderProvider, ProviderError, StageCheckpointReader,
12    StageCheckpointWriter, StatsReader, StorageChangeSetReader, StorageSettingsCache, TrieWriter,
13};
14use reth_stages_api::{
15    BlockErrorKind, EntitiesCheckpoint, ExecInput, ExecOutput, MerkleCheckpoint, Stage,
16    StageCheckpoint, StageError, StageId, StorageRootMerkleCheckpoint, UnwindInput, UnwindOutput,
17};
18use reth_trie::{IntermediateStateRootState, StateRoot, StateRootProgress, StoredSubNode};
19use reth_trie_db::DatabaseStateRoot;
20
21use std::fmt::Debug;
22
23type DbStateRoot<'a, TX, A> = StateRoot<
24    reth_trie_db::DatabaseTrieCursorFactory<&'a TX, A>,
25    reth_trie_db::DatabaseHashedCursorFactory<&'a TX>,
26>;
27use tracing::*;
28
29// TODO: automate the process outlined below so the user can just send in a debugging package
30/// The error message that we include in invalid state root errors to tell users what information
31/// they should include in a bug report, since true state root errors can be impossible to debug
32/// with just basic logs.
33pub const INVALID_STATE_ROOT_ERROR_MESSAGE: &str = r#"
34Invalid state root error on stage verification!
35This is an error that likely requires a report to the reth team with additional information.
36Please include the following information in your report:
37 * This error message
38 * The state root of the block that was rejected
39 * The output of `reth db stats --checksum` from the database that was being used. This will take a long time to run!
40 * 50-100 lines of logs before and after the first occurrence of the log message with the state root of the block that was rejected.
41 * The debug logs from __the same time period__. To find the default location for these logs, run:
42   `reth --help | grep -A 4 'log.file.directory'`
43
44Once you have this information, please submit a github issue at https://github.com/paradigmxyz/reth/issues/new
45"#;
46
47/// The default threshold (in number of blocks) for switching from incremental trie building
48/// of changes to whole rebuild.
49pub const MERKLE_STAGE_DEFAULT_REBUILD_THRESHOLD: u64 = 100_000;
50
51/// The default threshold (in number of blocks) to run the stage in incremental mode. The
52/// incremental mode will calculate the state root for a large range of blocks by calculating the
53/// new state root for this many blocks, in batches, repeating until we reach the desired block
54/// number.
55pub const MERKLE_STAGE_DEFAULT_INCREMENTAL_THRESHOLD: u64 = 7_000;
56
57/// The merkle hashing stage uses input from
58/// [`AccountHashingStage`][crate::stages::AccountHashingStage] and
59/// [`StorageHashingStage`][crate::stages::StorageHashingStage] to calculate intermediate hashes
60/// and state roots.
61///
62/// This stage should be run with the above two stages, otherwise it is a no-op.
63///
64/// This stage is split in two: one for calculating hashes and one for unwinding.
65///
66/// When run in execution, it's going to be executed AFTER the hashing stages, to generate
67/// the state root. When run in unwind mode, it's going to be executed BEFORE the hashing stages,
68/// so that it unwinds the intermediate hashes based on the unwound hashed state from the hashing
69/// stages. The order of these two variants is important. The unwind variant should be added to the
70/// pipeline before the execution variant.
71///
72/// An example pipeline to only hash state would be:
73///
74/// - [`MerkleStage::Unwind`]
75/// - [`AccountHashingStage`][crate::stages::AccountHashingStage]
76/// - [`StorageHashingStage`][crate::stages::StorageHashingStage]
77/// - [`MerkleStage::Execution`]
78#[derive(Debug, Clone)]
79pub enum MerkleStage {
80    /// The execution portion of the merkle stage.
81    Execution {
82        // TODO: make struct for holding incremental settings, for code reuse between `Execution`
83        // variant and `Both`
84        /// The threshold (in number of blocks) for switching from incremental trie building
85        /// of changes to whole rebuild.
86        rebuild_threshold: u64,
87        /// The threshold (in number of blocks) to run the stage in incremental mode. The
88        /// incremental mode will calculate the state root by calculating the new state root for
89        /// some number of blocks, repeating until we reach the desired block number.
90        incremental_threshold: u64,
91    },
92    /// The unwind portion of the merkle stage.
93    Unwind {
94        /// Whether every child of a changed branch path should be walked.
95        walk_all_changed_branch_children: bool,
96    },
97    /// Able to execute and unwind. Used for tests
98    #[cfg(any(test, feature = "test-utils"))]
99    Both {
100        /// The threshold (in number of blocks) for switching from incremental trie building
101        /// of changes to whole rebuild.
102        rebuild_threshold: u64,
103        /// The threshold (in number of blocks) to run the stage in incremental mode. The
104        /// incremental mode will calculate the state root by calculating the new state root for
105        /// some number of blocks, repeating until we reach the desired block number.
106        incremental_threshold: u64,
107    },
108}
109
110impl MerkleStage {
111    /// Stage default for the [`MerkleStage::Execution`].
112    pub const fn default_execution() -> Self {
113        Self::Execution {
114            rebuild_threshold: MERKLE_STAGE_DEFAULT_REBUILD_THRESHOLD,
115            incremental_threshold: MERKLE_STAGE_DEFAULT_INCREMENTAL_THRESHOLD,
116        }
117    }
118
119    /// Stage default for the [`MerkleStage::Unwind`].
120    pub const fn default_unwind() -> Self {
121        Self::new_unwind(false)
122    }
123
124    /// Create a new instance of [`MerkleStage::Unwind`].
125    pub const fn new_unwind(walk_all_changed_branch_children: bool) -> Self {
126        Self::Unwind { walk_all_changed_branch_children }
127    }
128
129    /// Create new instance of [`MerkleStage::Execution`].
130    pub const fn new_execution(rebuild_threshold: u64, incremental_threshold: u64) -> Self {
131        Self::Execution { rebuild_threshold, incremental_threshold }
132    }
133
134    /// Gets the hashing progress
135    pub fn get_execution_checkpoint(
136        &self,
137        provider: &impl StageCheckpointReader,
138    ) -> Result<Option<MerkleCheckpoint>, StageError> {
139        let buf =
140            provider.get_stage_checkpoint_progress(StageId::MerkleExecute)?.unwrap_or_default();
141
142        if buf.is_empty() {
143            return Ok(None)
144        }
145
146        let (checkpoint, _) = MerkleCheckpoint::from_compact(&buf, buf.len());
147        Ok(Some(checkpoint))
148    }
149
150    /// Saves the hashing progress
151    pub fn save_execution_checkpoint(
152        &self,
153        provider: &impl StageCheckpointWriter,
154        checkpoint: Option<MerkleCheckpoint>,
155    ) -> Result<(), StageError> {
156        let mut buf = vec![];
157        if let Some(checkpoint) = checkpoint {
158            debug!(
159                target: "sync::stages::merkle::exec",
160                last_account_key = ?checkpoint.last_account_key,
161                "Saving inner merkle checkpoint"
162            );
163            checkpoint.to_compact(&mut buf);
164        }
165        Ok(provider.save_stage_checkpoint_progress(StageId::MerkleExecute, buf)?)
166    }
167}
168
169impl<Provider> Stage<Provider> for MerkleStage
170where
171    Provider: DBProvider<Tx: DbTxMut>
172        + TrieWriter
173        + StatsReader
174        + HeaderProvider
175        + ChangeSetReader
176        + StorageChangeSetReader
177        + StorageSettingsCache
178        + StageCheckpointReader
179        + StageCheckpointWriter,
180{
181    /// Return the id of the stage
182    fn id(&self) -> StageId {
183        match self {
184            Self::Execution { .. } => StageId::MerkleExecute,
185            Self::Unwind { .. } => StageId::MerkleUnwind,
186            #[cfg(any(test, feature = "test-utils"))]
187            Self::Both { .. } => StageId::Other("MerkleBoth"),
188        }
189    }
190
191    /// Execute the stage.
192    fn execute(&mut self, provider: &Provider, input: ExecInput) -> Result<ExecOutput, StageError> {
193        let (threshold, incremental_threshold) = match self {
194            Self::Unwind { .. } => {
195                info!(target: "sync::stages::merkle::unwind", "Stage is always skipped");
196                return Ok(ExecOutput::done(StageCheckpoint::new(input.target())))
197            }
198            Self::Execution { rebuild_threshold, incremental_threshold } => {
199                (*rebuild_threshold, *incremental_threshold)
200            }
201            #[cfg(any(test, feature = "test-utils"))]
202            Self::Both { rebuild_threshold, incremental_threshold } => {
203                (*rebuild_threshold, *incremental_threshold)
204            }
205        };
206
207        let range = input.next_block_range();
208        let (from_block, to_block) = range.clone().into_inner();
209        let current_block_number = input.checkpoint().block_number;
210
211        let target_block = provider
212            .header_by_number(to_block)?
213            .ok_or_else(|| ProviderError::HeaderNotFound(to_block.into()))?;
214        let target_block_root = target_block.state_root();
215
216        let (trie_root, entities_checkpoint) = if range.is_empty() {
217            (target_block_root, input.checkpoint().entities_stage_checkpoint().unwrap_or_default())
218        } else if to_block - from_block > threshold || from_block == 1 {
219            let mut checkpoint = self.get_execution_checkpoint(provider)?;
220
221            // if there are more blocks than threshold it is faster to rebuild the trie
222            let mut entities_checkpoint = if let Some(checkpoint) =
223                checkpoint.as_ref().filter(|c| c.target_block == to_block)
224            {
225                debug!(
226                    target: "sync::stages::merkle::exec",
227                    current = ?current_block_number,
228                    target = ?to_block,
229                    last_account_key = ?checkpoint.last_account_key,
230                    "Continuing inner merkle checkpoint"
231                );
232
233                input.checkpoint().entities_stage_checkpoint()
234            } else {
235                debug!(
236                    target: "sync::stages::merkle::exec",
237                    current = ?current_block_number,
238                    target = ?to_block,
239                    previous_checkpoint = ?checkpoint,
240                    "Rebuilding trie"
241                );
242                // Reset the checkpoint and clear trie tables
243                checkpoint = None;
244                self.save_execution_checkpoint(provider, None)?;
245                provider.tx_ref().clear::<tables::AccountsTrie>()?;
246                provider.tx_ref().clear::<tables::StoragesTrie>()?;
247
248                None
249            }
250            .unwrap_or(EntitiesCheckpoint {
251                processed: 0,
252                total: (provider.count_entries::<tables::HashedAccounts>()? +
253                    provider.count_entries::<tables::HashedStorages>()?)
254                    as u64,
255            });
256
257            let tx = provider.tx_ref();
258            let progress = reth_trie_db::with_adapter!(provider, |A| {
259                DbStateRoot::<_, A>::from_tx(tx)
260                    .with_intermediate_state(checkpoint.map(IntermediateStateRootState::from))
261                    .root_with_progress()
262            })
263            .map_err(|e| {
264                error!(target: "sync::stages::merkle", %e, ?current_block_number, ?to_block, "State root with progress failed! {INVALID_STATE_ROOT_ERROR_MESSAGE}");
265                StageError::Fatal(Box::new(e))
266            })?;
267            match progress {
268                StateRootProgress::Progress(state, hashed_entries_walked, updates) => {
269                    provider.write_trie_updates(updates)?;
270
271                    let mut checkpoint = MerkleCheckpoint::new(
272                        to_block,
273                        state.account_root_state.last_hashed_key,
274                        state
275                            .account_root_state
276                            .walker_stack
277                            .into_iter()
278                            .map(StoredSubNode::from)
279                            .collect(),
280                        state.account_root_state.hash_builder.into(),
281                    );
282
283                    // Save storage root state if present
284                    if let Some(storage_state) = state.storage_root_state {
285                        checkpoint.storage_root_checkpoint =
286                            Some(StorageRootMerkleCheckpoint::new(
287                                storage_state.state.last_hashed_key,
288                                storage_state
289                                    .state
290                                    .walker_stack
291                                    .into_iter()
292                                    .map(StoredSubNode::from)
293                                    .collect(),
294                                storage_state.state.hash_builder.into(),
295                                storage_state.account,
296                            ));
297                    }
298                    self.save_execution_checkpoint(provider, Some(checkpoint))?;
299
300                    entities_checkpoint.processed += hashed_entries_walked as u64;
301
302                    return Ok(ExecOutput::in_progress(
303                        input.checkpoint().with_entities_stage_checkpoint(entities_checkpoint),
304                    ))
305                }
306                StateRootProgress::Complete(root, hashed_entries_walked, updates) => {
307                    provider.write_trie_updates(updates)?;
308
309                    entities_checkpoint.processed += hashed_entries_walked as u64;
310
311                    (root, entities_checkpoint)
312                }
313            }
314        } else {
315            debug!(target: "sync::stages::merkle::exec", current = ?current_block_number, target = ?to_block, "Updating trie in chunks");
316            let mut final_root = None;
317            for start_block in range.step_by(incremental_threshold as usize) {
318                let chunk_to = std::cmp::min(start_block + incremental_threshold - 1, to_block);
319                let chunk_range = start_block..=chunk_to;
320                debug!(
321                    target: "sync::stages::merkle::exec",
322                    current = ?current_block_number,
323                    target = ?to_block,
324                    incremental_threshold,
325                    chunk_range = ?chunk_range,
326                    "Processing chunk"
327                );
328                let (root, updates) = reth_trie_db::with_adapter!(provider, |A| {
329                    DbStateRoot::<_, A>::incremental_root_with_updates(provider, chunk_range)
330                })
331                .map_err(|e| {
332                    error!(target: "sync::stages::merkle", %e, ?current_block_number, ?to_block, "Incremental state root failed! {INVALID_STATE_ROOT_ERROR_MESSAGE}");
333                    StageError::Fatal(Box::new(e))
334                })?;
335                provider.write_trie_updates(updates)?;
336                final_root = Some(root);
337            }
338
339            // if we had no final root, we must have not looped above, which should not be possible
340            let final_root = final_root.ok_or(StageError::Fatal(
341                "Incremental merkle hashing did not produce a final root".into(),
342            ))?;
343
344            let total_hashed_entries = (provider.count_entries::<tables::HashedAccounts>()? +
345                provider.count_entries::<tables::HashedStorages>()?)
346                as u64;
347
348            let entities_checkpoint = EntitiesCheckpoint {
349                // This is fine because `range` doesn't have an upper bound, so in this `else`
350                // branch we're just hashing all remaining accounts and storage slots we have in the
351                // database.
352                processed: total_hashed_entries,
353                total: total_hashed_entries,
354            };
355            // Save the checkpoint
356            (final_root, entities_checkpoint)
357        };
358
359        // Reset the checkpoint
360        self.save_execution_checkpoint(provider, None)?;
361
362        validate_state_root(trie_root, SealedHeader::seal_slow(target_block), to_block)?;
363
364        Ok(ExecOutput::done(
365            StageCheckpoint::new(to_block).with_entities_stage_checkpoint(entities_checkpoint),
366        ))
367    }
368
369    /// Unwind the stage.
370    fn unwind(
371        &mut self,
372        provider: &Provider,
373        input: UnwindInput,
374    ) -> Result<UnwindOutput, StageError> {
375        let tx = provider.tx_ref();
376        let range = input.unwind_block_range();
377        if matches!(self, Self::Execution { .. }) {
378            info!(target: "sync::stages::merkle::unwind", "Stage is always skipped");
379            return Ok(UnwindOutput { checkpoint: StageCheckpoint::new(input.unwind_to) })
380        }
381        let walk_all_changed_branch_children = match self {
382            Self::Unwind { walk_all_changed_branch_children } => *walk_all_changed_branch_children,
383            #[cfg(any(test, feature = "test-utils"))]
384            Self::Both { .. } => false,
385            Self::Execution { .. } => unreachable!(),
386        };
387
388        let mut entities_checkpoint =
389            input.checkpoint.entities_stage_checkpoint().unwrap_or(EntitiesCheckpoint {
390                processed: 0,
391                total: (tx.entries::<tables::HashedAccounts>()? +
392                    tx.entries::<tables::HashedStorages>()?) as u64,
393            });
394
395        if input.unwind_to == 0 {
396            tx.clear::<tables::AccountsTrie>()?;
397            tx.clear::<tables::StoragesTrie>()?;
398
399            entities_checkpoint.processed = 0;
400
401            return Ok(UnwindOutput {
402                checkpoint: StageCheckpoint::new(input.unwind_to)
403                    .with_entities_stage_checkpoint(entities_checkpoint),
404            })
405        }
406
407        // Unwind trie only if there are transitions
408        if range.is_empty() {
409            info!(target: "sync::stages::merkle::unwind", "Nothing to unwind");
410        } else {
411            let (block_root, updates) = reth_trie_db::with_adapter!(provider, |A| {
412                DbStateRoot::<_, A>::incremental_root_calculator(provider, range).and_then(
413                    |calculator| {
414                        calculator
415                            .with_walk_all_changed_branch_children(walk_all_changed_branch_children)
416                            .root_with_updates()
417                    },
418                )
419            })
420            .map_err(|e| StageError::Fatal(Box::new(e)))?;
421
422            // Validate the calculated state root
423            let target = provider
424                .header_by_number(input.unwind_to)?
425                .ok_or_else(|| ProviderError::HeaderNotFound(input.unwind_to.into()))?;
426
427            validate_state_root(block_root, SealedHeader::seal_slow(target), input.unwind_to)?;
428
429            // Validation passed, apply unwind changes to the database.
430            provider.write_trie_updates(updates)?;
431
432            // Update entities checkpoint to reflect the unwind operation
433            // Since we're unwinding, we need to recalculate the total entities at the target block
434            let accounts = tx.entries::<tables::HashedAccounts>()?;
435            let storages = tx.entries::<tables::HashedStorages>()?;
436            let total = (accounts + storages) as u64;
437            entities_checkpoint.total = total;
438            entities_checkpoint.processed = total;
439        }
440
441        Ok(UnwindOutput {
442            checkpoint: StageCheckpoint::new(input.unwind_to)
443                .with_entities_stage_checkpoint(entities_checkpoint),
444        })
445    }
446}
447
448/// Check that the computed state root matches the root in the expected header.
449#[inline]
450fn validate_state_root<H: BlockHeader + Sealable + Debug>(
451    got: B256,
452    expected: SealedHeader<H>,
453    target_block: BlockNumber,
454) -> Result<(), StageError> {
455    if got == expected.state_root() {
456        Ok(())
457    } else {
458        error!(target: "sync::stages::merkle", ?target_block, ?got, ?expected, "Failed to verify block state root! {INVALID_STATE_ROOT_ERROR_MESSAGE}");
459        Err(StageError::Block {
460            error: BlockErrorKind::Validation(ConsensusError::BodyStateRootDiff(
461                GotExpected { got, expected: expected.state_root() }.into(),
462            )),
463            block: Box::new(expected.block_with_parent()),
464        })
465    }
466}
467
468#[cfg(test)]
469mod tests {
470    use super::*;
471    use crate::test_utils::{
472        stage_test_suite_ext, ExecuteStageTestRunner, StageTestRunner, StorageKind,
473        TestRunnerError, TestStageDB, UnwindStageTestRunner,
474    };
475    use alloy_primitives::{keccak256, U256};
476    use assert_matches::assert_matches;
477    use reth_db_api::cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO};
478    use reth_primitives_traits::{SealedBlock, StorageEntry};
479    use reth_provider::{providers::StaticFileWriter, StaticFileProviderFactory};
480    use reth_stages_api::StageUnitCheckpoint;
481    use reth_static_file_types::StaticFileSegment;
482    use reth_testing_utils::generators::{
483        self, random_block, random_block_range, random_changeset_range,
484        random_contract_account_range, BlockParams, BlockRangeParams,
485    };
486    use reth_trie::test_utils::{state_root, state_root_prehashed};
487    use std::collections::BTreeMap;
488
489    stage_test_suite_ext!(MerkleTestRunner, merkle);
490
491    /// Execute from genesis so as to merkelize whole state
492    #[tokio::test]
493    async fn execute_clean_merkle() {
494        let (previous_stage, stage_progress) = (500, 0);
495
496        // Set up the runner
497        let mut runner = MerkleTestRunner::default();
498        // set low threshold so we hash the whole storage
499        let input = ExecInput {
500            target: Some(previous_stage),
501            checkpoint: Some(StageCheckpoint::new(stage_progress)),
502        };
503
504        runner.seed_execution(input).expect("failed to seed execution");
505
506        let rx = runner.execute(input);
507
508        // Assert the successful result
509        let result = rx.await.unwrap();
510        assert_matches!(
511            result,
512            Ok(ExecOutput {
513                checkpoint: StageCheckpoint {
514                    block_number,
515                    stage_checkpoint: Some(StageUnitCheckpoint::Entities(EntitiesCheckpoint {
516                        processed,
517                        total
518                    }))
519                },
520                done: true
521            }) if block_number == previous_stage && processed == total &&
522                total == (
523                    runner.db.count_entries::<tables::HashedAccounts>().unwrap() +
524                    runner.db.count_entries::<tables::HashedStorages>().unwrap()
525                ) as u64
526        );
527
528        // Validate the stage execution
529        assert!(runner.validate_execution(input, result.ok()).is_ok(), "execution validation");
530    }
531
532    /// Update small trie
533    #[tokio::test]
534    async fn execute_small_merkle() {
535        let (previous_stage, stage_progress) = (2, 1);
536
537        // Set up the runner
538        let mut runner = MerkleTestRunner::default();
539        let input = ExecInput {
540            target: Some(previous_stage),
541            checkpoint: Some(StageCheckpoint::new(stage_progress)),
542        };
543
544        runner.seed_execution(input).expect("failed to seed execution");
545
546        let rx = runner.execute(input);
547
548        // Assert the successful result
549        let result = rx.await.unwrap();
550        assert_matches!(
551            result,
552            Ok(ExecOutput {
553                checkpoint: StageCheckpoint {
554                    block_number,
555                    stage_checkpoint: Some(StageUnitCheckpoint::Entities(EntitiesCheckpoint {
556                        processed,
557                        total
558                    }))
559                },
560                done: true
561            }) if block_number == previous_stage && processed == total &&
562                total == (
563                    runner.db.count_entries::<tables::HashedAccounts>().unwrap() +
564                    runner.db.count_entries::<tables::HashedStorages>().unwrap()
565                ) as u64
566        );
567
568        // Validate the stage execution
569        assert!(runner.validate_execution(input, result.ok()).is_ok(), "execution validation");
570    }
571
572    #[tokio::test]
573    async fn execute_chunked_merkle() {
574        let (previous_stage, stage_progress) = (200, 100);
575        let clean_threshold = 100;
576        let incremental_threshold = 10;
577
578        // Set up the runner
579        let mut runner =
580            MerkleTestRunner { db: TestStageDB::default(), clean_threshold, incremental_threshold };
581
582        let input = ExecInput {
583            target: Some(previous_stage),
584            checkpoint: Some(StageCheckpoint::new(stage_progress)),
585        };
586
587        runner.seed_execution(input).expect("failed to seed execution");
588        let rx = runner.execute(input);
589
590        // Assert the successful result
591        let result = rx.await.unwrap();
592        assert_matches!(
593            result,
594            Ok(ExecOutput {
595                checkpoint: StageCheckpoint {
596                    block_number,
597                    stage_checkpoint: Some(StageUnitCheckpoint::Entities(EntitiesCheckpoint {
598                        processed,
599                        total
600                    }))
601                },
602                done: true
603            }) if block_number == previous_stage && processed == total &&
604                total == (
605                    runner.db.count_entries::<tables::HashedAccounts>().unwrap() +
606                    runner.db.count_entries::<tables::HashedStorages>().unwrap()
607                ) as u64
608        );
609
610        // Validate the stage execution
611        let provider = runner.db.factory.provider().unwrap();
612        let header = provider.header_by_number(previous_stage).unwrap().unwrap();
613        let expected_root = header.state_root;
614
615        let actual_root = runner
616            .db
617            .query_with_provider(|provider| {
618                Ok(reth_trie_db::with_adapter!(provider, |A| {
619                    DbStateRoot::<_, A>::incremental_root_with_updates(
620                        &provider,
621                        stage_progress + 1..=previous_stage,
622                    )
623                }))
624            })
625            .unwrap();
626
627        assert_eq!(
628            actual_root.unwrap().0,
629            expected_root,
630            "State root mismatch after chunked processing"
631        );
632    }
633
634    struct MerkleTestRunner {
635        db: TestStageDB,
636        clean_threshold: u64,
637        incremental_threshold: u64,
638    }
639
640    impl Default for MerkleTestRunner {
641        fn default() -> Self {
642            Self {
643                db: TestStageDB::default(),
644                clean_threshold: 10000,
645                incremental_threshold: 10000,
646            }
647        }
648    }
649
650    impl StageTestRunner for MerkleTestRunner {
651        type S = MerkleStage;
652
653        fn db(&self) -> &TestStageDB {
654            &self.db
655        }
656
657        fn stage(&self) -> Self::S {
658            Self::S::Both {
659                rebuild_threshold: self.clean_threshold,
660                incremental_threshold: self.incremental_threshold,
661            }
662        }
663    }
664
665    impl ExecuteStageTestRunner for MerkleTestRunner {
666        type Seed = Vec<SealedBlock<reth_ethereum_primitives::Block>>;
667
668        #[allow(clippy::clone_on_copy)]
669        fn seed_execution(&mut self, input: ExecInput) -> Result<Self::Seed, TestRunnerError> {
670            let stage_progress = input.checkpoint().block_number;
671            let start = stage_progress + 1;
672            let end = input.target();
673            let mut rng = generators::rng();
674
675            let mut preblocks = vec![];
676            if stage_progress > 0 {
677                preblocks.append(&mut random_block_range(
678                    &mut rng,
679                    0..=stage_progress - 1,
680                    BlockRangeParams {
681                        parent: Some(B256::ZERO),
682                        tx_count: 0..1,
683                        ..Default::default()
684                    },
685                ));
686                self.db.insert_blocks(preblocks.iter(), StorageKind::Static)?;
687            }
688
689            let num_of_accounts = 31;
690            let accounts = random_contract_account_range(&mut rng, &mut (0..num_of_accounts))
691                .into_iter()
692                .collect::<BTreeMap<_, _>>();
693
694            self.db.insert_accounts_and_storages(
695                accounts.iter().map(|(addr, acc)| (*addr, (acc.clone(), std::iter::empty()))),
696            )?;
697
698            let (header, body) = random_block(
699                &mut rng,
700                stage_progress,
701                BlockParams { parent: preblocks.last().map(|b| b.hash()), ..Default::default() },
702            )
703            .split_sealed_header_body();
704            let mut header = header.unseal();
705
706            header.state_root = state_root(
707                accounts
708                    .clone()
709                    .into_iter()
710                    .map(|(address, account)| (address, (account, std::iter::empty()))),
711            );
712            let sealed_head = SealedBlock::<reth_ethereum_primitives::Block>::from_sealed_parts(
713                SealedHeader::seal_slow(header),
714                body,
715            );
716
717            let head_hash = sealed_head.hash();
718            let mut blocks = vec![sealed_head];
719            blocks.extend(random_block_range(
720                &mut rng,
721                start..=end,
722                BlockRangeParams { parent: Some(head_hash), tx_count: 0..3, ..Default::default() },
723            ));
724            let last_block = blocks.last().cloned().unwrap();
725            self.db.insert_blocks(blocks.iter(), StorageKind::Static)?;
726
727            let (transitions, final_state) = random_changeset_range(
728                &mut rng,
729                blocks.iter(),
730                accounts.into_iter().map(|(addr, acc)| (addr, (acc, Vec::new()))),
731                0..3,
732                0..256,
733            );
734            // add block changeset from block 1.
735            self.db.insert_changesets(transitions, Some(start))?;
736            self.db.insert_accounts_and_storages(final_state)?;
737
738            // Calculate state root
739            let root = self.db.query(|tx| {
740                let mut accounts = BTreeMap::default();
741                let mut accounts_cursor = tx.cursor_read::<tables::HashedAccounts>()?;
742                let mut storage_cursor = tx.cursor_dup_read::<tables::HashedStorages>()?;
743                for entry in accounts_cursor.walk_range(..)? {
744                    let (key, account) = entry?;
745                    let mut storage_entries = Vec::new();
746                    let mut entry = storage_cursor.seek_exact(key)?;
747                    while let Some((_, storage)) = entry {
748                        storage_entries.push(storage);
749                        entry = storage_cursor.next_dup()?;
750                    }
751                    let storage = storage_entries
752                        .into_iter()
753                        .filter(|v| !v.value.is_zero())
754                        .map(|v| (v.key, v.value))
755                        .collect::<Vec<_>>();
756                    accounts.insert(key, (account, storage));
757                }
758
759                Ok(state_root_prehashed(accounts))
760            })?;
761
762            let static_file_provider = self.db.factory.static_file_provider();
763            let mut writer =
764                static_file_provider.latest_writer(StaticFileSegment::Headers).unwrap();
765            let mut last_header = last_block.clone_sealed_header();
766            last_header.set_state_root(root);
767
768            let hash = last_header.hash_slow();
769            writer.prune_headers(1).unwrap();
770            writer.commit().unwrap();
771            writer.append_header(&last_header, &hash).unwrap();
772            writer.commit().unwrap();
773
774            Ok(blocks)
775        }
776
777        fn validate_execution(
778            &self,
779            _input: ExecInput,
780            _output: Option<ExecOutput>,
781        ) -> Result<(), TestRunnerError> {
782            // The execution is validated within the stage
783            Ok(())
784        }
785    }
786
787    impl UnwindStageTestRunner for MerkleTestRunner {
788        fn validate_unwind(&self, _input: UnwindInput) -> Result<(), TestRunnerError> {
789            // The unwind is validated within the stage
790            Ok(())
791        }
792
793        fn before_unwind(&self, input: UnwindInput) -> Result<(), TestRunnerError> {
794            let target_block = input.unwind_to + 1;
795
796            self.db
797                .commit(|tx| {
798                    let mut storage_changesets_cursor =
799                        tx.cursor_dup_read::<tables::StorageChangeSets>().unwrap();
800                    let mut storage_cursor =
801                        tx.cursor_dup_write::<tables::HashedStorages>().unwrap();
802
803                    let mut tree: BTreeMap<B256, BTreeMap<B256, U256>> = BTreeMap::new();
804
805                    let mut rev_changeset_walker =
806                        storage_changesets_cursor.walk_back(None).unwrap();
807                    while let Some((bn_address, entry)) =
808                        rev_changeset_walker.next().transpose().unwrap()
809                    {
810                        if bn_address.block_number() < target_block {
811                            break
812                        }
813
814                        tree.entry(keccak256(bn_address.address()))
815                            .or_default()
816                            .insert(keccak256(entry.key), entry.value);
817                    }
818                    for (hashed_address, storage) in tree {
819                        for (hashed_slot, value) in storage {
820                            let storage_entry = storage_cursor
821                                .seek_by_key_subkey(hashed_address, hashed_slot)
822                                .unwrap();
823                            if storage_entry.is_some_and(|v| v.key == hashed_slot) {
824                                storage_cursor.delete_current().unwrap();
825                            }
826
827                            if !value.is_zero() {
828                                let storage_entry = StorageEntry { key: hashed_slot, value };
829                                storage_cursor.upsert(hashed_address, &storage_entry).unwrap();
830                            }
831                        }
832                    }
833
834                    let mut changeset_cursor =
835                        tx.cursor_dup_write::<tables::AccountChangeSets>().unwrap();
836                    let mut rev_changeset_walker = changeset_cursor.walk_back(None).unwrap();
837
838                    while let Some((block_number, account_before_tx)) =
839                        rev_changeset_walker.next().transpose().unwrap()
840                    {
841                        if block_number < target_block {
842                            break
843                        }
844
845                        if let Some(acc) = account_before_tx.info {
846                            tx.put::<tables::HashedAccounts>(
847                                keccak256(account_before_tx.address),
848                                acc,
849                            )
850                            .unwrap();
851                        } else {
852                            tx.delete::<tables::HashedAccounts>(
853                                keccak256(account_before_tx.address),
854                                None,
855                            )
856                            .unwrap();
857                        }
858                    }
859                    Ok(())
860                })
861                .unwrap();
862            Ok(())
863        }
864    }
865}