Skip to main content

reth_stages_types/
checkpoints.rs

1use super::StageId;
2#[cfg(test)]
3use alloc::vec;
4use alloc::{format, string::String, vec::Vec};
5use alloy_primitives::{Address, BlockNumber, B256, U256};
6use core::ops::RangeInclusive;
7use reth_trie_common::{hash_builder::HashBuilderState, StoredSubNode};
8
9/// Saves the progress of Merkle stage.
10#[derive(Default, Debug, Clone, PartialEq, Eq)]
11pub struct MerkleCheckpoint {
12    /// The target block number.
13    pub target_block: BlockNumber,
14    /// The last hashed account key processed.
15    pub last_account_key: B256,
16    /// Previously recorded walker stack.
17    pub walker_stack: Vec<StoredSubNode>,
18    /// The hash builder state.
19    pub state: HashBuilderState,
20    /// Optional storage root checkpoint for the last processed account.
21    pub storage_root_checkpoint: Option<StorageRootMerkleCheckpoint>,
22}
23
24impl MerkleCheckpoint {
25    /// Creates a new Merkle checkpoint.
26    pub const fn new(
27        target_block: BlockNumber,
28        last_account_key: B256,
29        walker_stack: Vec<StoredSubNode>,
30        state: HashBuilderState,
31    ) -> Self {
32        Self { target_block, last_account_key, walker_stack, state, storage_root_checkpoint: None }
33    }
34}
35
36#[cfg(any(test, feature = "reth-codec"))]
37impl reth_codecs::Compact for MerkleCheckpoint {
38    fn to_compact<B>(&self, buf: &mut B) -> usize
39    where
40        B: bytes::BufMut + AsMut<[u8]>,
41    {
42        let mut len = 0;
43
44        buf.put_u64(self.target_block);
45        len += 8;
46
47        buf.put_slice(self.last_account_key.as_slice());
48        len += self.last_account_key.len();
49
50        buf.put_u16(self.walker_stack.len() as u16);
51        len += 2;
52        for item in &self.walker_stack {
53            len += item.to_compact(buf);
54        }
55
56        len += self.state.to_compact(buf);
57
58        // Encode the optional storage root checkpoint
59        match &self.storage_root_checkpoint {
60            Some(checkpoint) => {
61                // one means Some
62                buf.put_u8(1);
63                len += 1;
64                len += checkpoint.to_compact(buf);
65            }
66            None => {
67                // zero means None
68                buf.put_u8(0);
69                len += 1;
70            }
71        }
72
73        len
74    }
75
76    fn from_compact(mut buf: &[u8], _len: usize) -> (Self, &[u8]) {
77        use bytes::Buf;
78        let target_block = buf.get_u64();
79
80        let last_account_key = B256::from_slice(&buf[..32]);
81        buf.advance(32);
82
83        let walker_stack_len = buf.get_u16() as usize;
84        let mut walker_stack = Vec::with_capacity(walker_stack_len);
85        for _ in 0..walker_stack_len {
86            let (item, rest) = StoredSubNode::from_compact(buf, 0);
87            walker_stack.push(item);
88            buf = rest;
89        }
90
91        let (state, mut buf) = HashBuilderState::from_compact(buf, 0);
92
93        // Decode the storage root checkpoint if it exists
94        let (storage_root_checkpoint, buf) = if buf.is_empty() {
95            (None, buf)
96        } else {
97            match buf.get_u8() {
98                1 => {
99                    let (checkpoint, rest) = StorageRootMerkleCheckpoint::from_compact(buf, 0);
100                    (Some(checkpoint), rest)
101                }
102                _ => (None, buf),
103            }
104        };
105
106        (Self { target_block, last_account_key, walker_stack, state, storage_root_checkpoint }, buf)
107    }
108}
109
110/// Saves the progress of a storage root computation.
111///
112/// This contains the walker stack, hash builder state, and the last storage key processed.
113#[derive(Debug, Clone, PartialEq, Eq)]
114pub struct StorageRootMerkleCheckpoint {
115    /// The last storage key processed.
116    pub last_storage_key: B256,
117    /// Previously recorded walker stack.
118    pub walker_stack: Vec<StoredSubNode>,
119    /// The hash builder state.
120    pub state: HashBuilderState,
121    /// The account nonce.
122    pub account_nonce: u64,
123    /// The account balance.
124    pub account_balance: U256,
125    /// The account bytecode hash.
126    pub account_bytecode_hash: B256,
127}
128
129impl StorageRootMerkleCheckpoint {
130    /// Creates a new storage root merkle checkpoint.
131    pub const fn new(
132        last_storage_key: B256,
133        walker_stack: Vec<StoredSubNode>,
134        state: HashBuilderState,
135        account_nonce: u64,
136        account_balance: U256,
137        account_bytecode_hash: B256,
138    ) -> Self {
139        Self {
140            last_storage_key,
141            walker_stack,
142            state,
143            account_nonce,
144            account_balance,
145            account_bytecode_hash,
146        }
147    }
148}
149
150#[cfg(any(test, feature = "reth-codec"))]
151impl reth_codecs::Compact for StorageRootMerkleCheckpoint {
152    fn to_compact<B>(&self, buf: &mut B) -> usize
153    where
154        B: bytes::BufMut + AsMut<[u8]>,
155    {
156        let mut len = 0;
157
158        buf.put_slice(self.last_storage_key.as_slice());
159        len += self.last_storage_key.len();
160
161        buf.put_u16(self.walker_stack.len() as u16);
162        len += 2;
163        for item in &self.walker_stack {
164            len += item.to_compact(buf);
165        }
166
167        len += self.state.to_compact(buf);
168
169        // Encode account fields
170        buf.put_u64(self.account_nonce);
171        len += 8;
172
173        let balance_len = self.account_balance.byte_len() as u8;
174        buf.put_u8(balance_len);
175        len += 1;
176        len += self.account_balance.to_compact(buf);
177
178        buf.put_slice(self.account_bytecode_hash.as_slice());
179        len += 32;
180
181        len
182    }
183
184    fn from_compact(mut buf: &[u8], _len: usize) -> (Self, &[u8]) {
185        use bytes::Buf;
186
187        let last_storage_key = B256::from_slice(&buf[..32]);
188        buf.advance(32);
189
190        let walker_stack_len = buf.get_u16() as usize;
191        let mut walker_stack = Vec::with_capacity(walker_stack_len);
192        for _ in 0..walker_stack_len {
193            let (item, rest) = StoredSubNode::from_compact(buf, 0);
194            walker_stack.push(item);
195            buf = rest;
196        }
197
198        let (state, mut buf) = HashBuilderState::from_compact(buf, 0);
199
200        // Decode account fields
201        let account_nonce = buf.get_u64();
202        let balance_len = buf.get_u8() as usize;
203        let (account_balance, mut buf) = U256::from_compact(buf, balance_len);
204        let account_bytecode_hash = B256::from_slice(&buf[..32]);
205        buf.advance(32);
206
207        (
208            Self {
209                last_storage_key,
210                walker_stack,
211                state,
212                account_nonce,
213                account_balance,
214                account_bytecode_hash,
215            },
216            buf,
217        )
218    }
219}
220
221/// Saves the progress of `AccountHashing` stage.
222#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
223#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
224#[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
225#[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
226#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
227pub struct AccountHashingCheckpoint {
228    /// The next account to start hashing from.
229    pub address: Option<Address>,
230    /// Block range which this checkpoint is valid for.
231    pub block_range: CheckpointBlockRange,
232    /// Progress measured in accounts.
233    pub progress: EntitiesCheckpoint,
234}
235
236/// Saves the progress of `StorageHashing` stage.
237#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
238#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
239#[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
240#[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
241#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
242pub struct StorageHashingCheckpoint {
243    /// The next account to start hashing from.
244    pub address: Option<Address>,
245    /// The next storage slot to start hashing from.
246    pub storage: Option<B256>,
247    /// Block range which this checkpoint is valid for.
248    pub block_range: CheckpointBlockRange,
249    /// Progress measured in storage slots.
250    pub progress: EntitiesCheckpoint,
251}
252
253/// Saves the progress of Execution stage.
254#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
255#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
256#[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
257#[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
258#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
259pub struct ExecutionCheckpoint {
260    /// Block range which this checkpoint is valid for.
261    pub block_range: CheckpointBlockRange,
262    /// Progress measured in gas.
263    pub progress: EntitiesCheckpoint,
264}
265
266/// Saves the progress of Headers stage.
267#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
268#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
269#[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
270#[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
271#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
272pub struct HeadersCheckpoint {
273    /// Block range which this checkpoint is valid for.
274    pub block_range: CheckpointBlockRange,
275    /// Progress measured in gas.
276    pub progress: EntitiesCheckpoint,
277}
278
279/// Saves the progress of Index History stages.
280#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
281#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
282#[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
283#[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
284#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
285pub struct IndexHistoryCheckpoint {
286    /// Block range which this checkpoint is valid for.
287    pub block_range: CheckpointBlockRange,
288    /// Progress measured in changesets.
289    pub progress: EntitiesCheckpoint,
290}
291
292/// Saves the progress of `MerkleChangeSets` stage.
293///
294/// Note: This type is only kept for backward compatibility with the Compact codec.
295/// The `MerkleChangeSets` stage has been removed.
296#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
297#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
298#[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
299#[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
300#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
301pub struct MerkleChangeSetsCheckpoint {
302    /// Block range which this checkpoint is valid for.
303    pub block_range: CheckpointBlockRange,
304}
305
306/// Saves the progress of abstract stage iterating over or downloading entities.
307#[derive(Debug, Default, PartialEq, Eq, Clone, Copy)]
308#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
309#[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
310#[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
311#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
312pub struct EntitiesCheckpoint {
313    /// Number of entities already processed.
314    pub processed: u64,
315    /// Total entities to be processed.
316    pub total: u64,
317}
318
319impl EntitiesCheckpoint {
320    /// Formats entities checkpoint as percentage, i.e. `processed / total`.
321    ///
322    /// Return [None] if `total == 0`.
323    pub fn fmt_percentage(&self) -> Option<String> {
324        if self.total == 0 {
325            return None
326        }
327
328        // Calculate percentage with 2 decimal places.
329        let percentage = 100.0 * self.processed as f64 / self.total as f64;
330
331        // Truncate to 2 decimal places, rounding down so that 99.999% becomes 99.99% and not 100%.
332        #[cfg(not(feature = "std"))]
333        {
334            // Manual floor implementation using integer arithmetic for no_std
335            let scaled = (percentage * 100.0) as u64;
336            Some(format!("{:.2}%", scaled as f64 / 100.0))
337        }
338        #[cfg(feature = "std")]
339        Some(format!("{:.2}%", (percentage * 100.0).floor() / 100.0))
340    }
341}
342
343/// Saves the block range. Usually, it's used to check the validity of some stage checkpoint across
344/// multiple executions.
345#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
346#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
347#[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
348#[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
349#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
350pub struct CheckpointBlockRange {
351    /// The first block of the range, inclusive.
352    pub from: BlockNumber,
353    /// The last block of the range, inclusive.
354    pub to: BlockNumber,
355}
356
357impl From<RangeInclusive<BlockNumber>> for CheckpointBlockRange {
358    fn from(range: RangeInclusive<BlockNumber>) -> Self {
359        Self { from: *range.start(), to: *range.end() }
360    }
361}
362
363impl From<&RangeInclusive<BlockNumber>> for CheckpointBlockRange {
364    fn from(range: &RangeInclusive<BlockNumber>) -> Self {
365        Self { from: *range.start(), to: *range.end() }
366    }
367}
368
369/// Saves the progress of a stage.
370#[derive(Debug, Default, PartialEq, Eq, Clone, Copy)]
371#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
372#[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
373#[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
374#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
375pub struct StageCheckpoint {
376    /// The maximum block processed by the stage.
377    pub block_number: BlockNumber,
378    /// Stage-specific checkpoint. None if stage uses only block-based checkpoints.
379    pub stage_checkpoint: Option<StageUnitCheckpoint>,
380}
381
382impl StageCheckpoint {
383    /// Creates a new [`StageCheckpoint`] with only `block_number` set.
384    pub fn new(block_number: BlockNumber) -> Self {
385        Self { block_number, ..Default::default() }
386    }
387
388    /// Sets the block number.
389    pub const fn with_block_number(mut self, block_number: BlockNumber) -> Self {
390        self.block_number = block_number;
391        self
392    }
393
394    /// Sets the block range, if checkpoint uses block range.
395    pub fn with_block_range(mut self, stage_id: &StageId, from: u64, to: u64) -> Self {
396        self.stage_checkpoint = Some(match stage_id {
397            StageId::Execution => StageUnitCheckpoint::Execution(ExecutionCheckpoint::default()),
398            StageId::AccountHashing => {
399                StageUnitCheckpoint::Account(AccountHashingCheckpoint::default())
400            }
401            StageId::StorageHashing => {
402                StageUnitCheckpoint::Storage(StorageHashingCheckpoint::default())
403            }
404            StageId::IndexStorageHistory | StageId::IndexAccountHistory => {
405                StageUnitCheckpoint::IndexHistory(IndexHistoryCheckpoint::default())
406            }
407            _ => return self,
408        });
409        if let Some(ref mut checkpoint) = self.stage_checkpoint {
410            checkpoint.set_block_range(from, to);
411        }
412        self
413    }
414
415    /// Get the underlying [`EntitiesCheckpoint`], if any, to determine the number of entities
416    /// processed, and the number of total entities to process.
417    pub fn entities(&self) -> Option<EntitiesCheckpoint> {
418        let stage_checkpoint = self.stage_checkpoint?;
419
420        match stage_checkpoint {
421            StageUnitCheckpoint::Account(AccountHashingCheckpoint {
422                progress: entities, ..
423            }) |
424            StageUnitCheckpoint::Storage(StorageHashingCheckpoint {
425                progress: entities, ..
426            }) |
427            StageUnitCheckpoint::Entities(entities) |
428            StageUnitCheckpoint::Execution(ExecutionCheckpoint { progress: entities, .. }) |
429            StageUnitCheckpoint::Headers(HeadersCheckpoint { progress: entities, .. }) |
430            StageUnitCheckpoint::IndexHistory(IndexHistoryCheckpoint {
431                progress: entities,
432                ..
433            }) => Some(entities),
434            StageUnitCheckpoint::MerkleChangeSets(_) | StageUnitCheckpoint::Finish(_) => None,
435        }
436    }
437}
438
439#[cfg(any(test, feature = "reth-codec"))]
440reth_codecs::impl_compression_for_compact!(StageCheckpoint);
441
442/// Saves the progress of the Finish stage.
443#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
444#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
445#[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
446#[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
447#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
448pub struct FinishCheckpoint {
449    /// The highest block with a partially persisted state and trie.
450    pub partial_state_trie: Option<BlockNumber>,
451}
452
453impl FinishCheckpoint {
454    /// Returns the highest block with a partially persisted state and trie.
455    pub const fn partial_state_trie(&self) -> Option<BlockNumber> {
456        self.partial_state_trie
457    }
458}
459
460// TODO(alexey): add a merkle checkpoint. Currently it's hard because [`MerkleCheckpoint`]
461//  is not a Copy type.
462/// Stage-specific checkpoint metrics.
463#[derive(Debug, PartialEq, Eq, Clone, Copy)]
464#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
465#[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
466#[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
467#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
468pub enum StageUnitCheckpoint {
469    /// Saves the progress of `AccountHashing` stage.
470    Account(AccountHashingCheckpoint),
471    /// Saves the progress of `StorageHashing` stage.
472    Storage(StorageHashingCheckpoint),
473    /// Saves the progress of abstract stage iterating over or downloading entities.
474    Entities(EntitiesCheckpoint),
475    /// Saves the progress of Execution stage.
476    Execution(ExecutionCheckpoint),
477    /// Saves the progress of Headers stage.
478    Headers(HeadersCheckpoint),
479    /// Saves the progress of Index History stage.
480    IndexHistory(IndexHistoryCheckpoint),
481    /// Saves the progress of `MerkleChangeSets` stage.
482    ///
483    /// Note: This variant is only kept for backward compatibility with the Compact codec.
484    /// The `MerkleChangeSets` stage has been removed.
485    MerkleChangeSets(MerkleChangeSetsCheckpoint),
486    /// Saves the progress of the Finish stage.
487    Finish(FinishCheckpoint),
488}
489
490impl StageUnitCheckpoint {
491    /// Sets the block range. Returns old block range, or `None` if checkpoint doesn't use block
492    /// range.
493    pub const fn set_block_range(&mut self, from: u64, to: u64) -> Option<CheckpointBlockRange> {
494        match self {
495            Self::Account(AccountHashingCheckpoint { block_range, .. }) |
496            Self::Storage(StorageHashingCheckpoint { block_range, .. }) |
497            Self::Execution(ExecutionCheckpoint { block_range, .. }) |
498            Self::IndexHistory(IndexHistoryCheckpoint { block_range, .. }) => {
499                let old_range = *block_range;
500                *block_range = CheckpointBlockRange { from, to };
501
502                Some(old_range)
503            }
504            _ => None,
505        }
506    }
507}
508
509#[cfg(test)]
510impl Default for StageUnitCheckpoint {
511    fn default() -> Self {
512        Self::Account(AccountHashingCheckpoint::default())
513    }
514}
515
516/// Generates [`StageCheckpoint`] getter and builder methods.
517macro_rules! stage_unit_checkpoints {
518    ($(($index:expr,$enum_variant:tt,$checkpoint_ty:ty,#[doc = $fn_get_doc:expr]$fn_get_name:ident,#[doc = $fn_build_doc:expr]$fn_build_name:ident)),+) => {
519        impl StageCheckpoint {
520            $(
521                #[doc = $fn_get_doc]
522                pub const fn $fn_get_name(&self) -> Option<$checkpoint_ty> {
523                    match self.stage_checkpoint {
524                        Some(StageUnitCheckpoint::$enum_variant(checkpoint)) => Some(checkpoint),
525                        _ => None,
526                    }
527                }
528
529                #[doc = $fn_build_doc]
530                pub const fn $fn_build_name(
531                    mut self,
532                    checkpoint: $checkpoint_ty,
533                ) -> Self {
534                    self.stage_checkpoint = Some(StageUnitCheckpoint::$enum_variant(checkpoint));
535                    self
536                }
537            )+
538        }
539    };
540}
541
542stage_unit_checkpoints!(
543    (
544        0,
545        Account,
546        AccountHashingCheckpoint,
547        /// Returns the account hashing stage checkpoint, if any.
548        account_hashing_stage_checkpoint,
549        /// Sets the stage checkpoint to account hashing.
550        with_account_hashing_stage_checkpoint
551    ),
552    (
553        1,
554        Storage,
555        StorageHashingCheckpoint,
556        /// Returns the storage hashing stage checkpoint, if any.
557        storage_hashing_stage_checkpoint,
558        /// Sets the stage checkpoint to storage hashing.
559        with_storage_hashing_stage_checkpoint
560    ),
561    (
562        2,
563        Entities,
564        EntitiesCheckpoint,
565        /// Returns the entities stage checkpoint, if any.
566        entities_stage_checkpoint,
567        /// Sets the stage checkpoint to entities.
568        with_entities_stage_checkpoint
569    ),
570    (
571        3,
572        Execution,
573        ExecutionCheckpoint,
574        /// Returns the execution stage checkpoint, if any.
575        execution_stage_checkpoint,
576        /// Sets the stage checkpoint to execution.
577        with_execution_stage_checkpoint
578    ),
579    (
580        4,
581        Headers,
582        HeadersCheckpoint,
583        /// Returns the headers stage checkpoint, if any.
584        headers_stage_checkpoint,
585        /// Sets the stage checkpoint to headers.
586        with_headers_stage_checkpoint
587    ),
588    (
589        5,
590        IndexHistory,
591        IndexHistoryCheckpoint,
592        /// Returns the index history stage checkpoint, if any.
593        index_history_stage_checkpoint,
594        /// Sets the stage checkpoint to index history.
595        with_index_history_stage_checkpoint
596    ),
597    (
598        6,
599        Finish,
600        FinishCheckpoint,
601        /// Returns the finish stage checkpoint, if any.
602        finish_stage_checkpoint,
603        /// Sets the stage checkpoint to finish.
604        with_finish_stage_checkpoint
605    )
606);
607
608#[cfg(test)]
609mod tests {
610    use super::*;
611    use alloy_primitives::b256;
612    use rand::Rng;
613    use reth_codecs::Compact;
614
615    #[test]
616    fn merkle_checkpoint_roundtrip() {
617        let mut rng = rand::rng();
618        let checkpoint = MerkleCheckpoint {
619            target_block: rng.random(),
620            last_account_key: rng.random(),
621            walker_stack: vec![StoredSubNode {
622                key: B256::random_with(&mut rng).to_vec(),
623                nibble: Some(rng.random()),
624                node: None,
625            }],
626            state: HashBuilderState::default(),
627            storage_root_checkpoint: None,
628        };
629
630        let mut buf = Vec::new();
631        let encoded = checkpoint.to_compact(&mut buf);
632        let (decoded, _) = MerkleCheckpoint::from_compact(&buf, encoded);
633        assert_eq!(decoded, checkpoint);
634    }
635
636    #[test]
637    fn storage_root_merkle_checkpoint_roundtrip() {
638        let mut rng = rand::rng();
639        let checkpoint = StorageRootMerkleCheckpoint {
640            last_storage_key: rng.random(),
641            walker_stack: vec![StoredSubNode {
642                key: B256::random_with(&mut rng).to_vec(),
643                nibble: Some(rng.random()),
644                node: None,
645            }],
646            state: HashBuilderState::default(),
647            account_nonce: 0,
648            account_balance: U256::ZERO,
649            account_bytecode_hash: B256::ZERO,
650        };
651
652        let mut buf = Vec::new();
653        let encoded = checkpoint.to_compact(&mut buf);
654        let (decoded, _) = StorageRootMerkleCheckpoint::from_compact(&buf, encoded);
655        assert_eq!(decoded, checkpoint);
656    }
657
658    #[test]
659    fn merkle_checkpoint_with_storage_root_roundtrip() {
660        let mut rng = rand::rng();
661
662        // Create a storage root checkpoint
663        let storage_checkpoint = StorageRootMerkleCheckpoint {
664            last_storage_key: rng.random(),
665            walker_stack: vec![StoredSubNode {
666                key: B256::random_with(&mut rng).to_vec(),
667                nibble: Some(rng.random()),
668                node: None,
669            }],
670            state: HashBuilderState::default(),
671            account_nonce: 1,
672            account_balance: U256::from(1),
673            account_bytecode_hash: b256!(
674                "0x0fffffffffffffffffffffffffffffff0fffffffffffffffffffffffffffffff"
675            ),
676        };
677
678        // Create a merkle checkpoint with the storage root checkpoint
679        let checkpoint = MerkleCheckpoint {
680            target_block: rng.random(),
681            last_account_key: rng.random(),
682            walker_stack: vec![StoredSubNode {
683                key: B256::random_with(&mut rng).to_vec(),
684                nibble: Some(rng.random()),
685                node: None,
686            }],
687            state: HashBuilderState::default(),
688            storage_root_checkpoint: Some(storage_checkpoint),
689        };
690
691        let mut buf = Vec::new();
692        let encoded = checkpoint.to_compact(&mut buf);
693        let (decoded, _) = MerkleCheckpoint::from_compact(&buf, encoded);
694        assert_eq!(decoded, checkpoint);
695    }
696
697    #[test]
698    fn finish_checkpoint_roundtrip() {
699        let finish_checkpoint = FinishCheckpoint { partial_state_trie: Some(21) };
700        let checkpoint = StageCheckpoint::new(42).with_finish_stage_checkpoint(finish_checkpoint);
701
702        let mut buf = Vec::new();
703        let encoded = checkpoint.to_compact(&mut buf);
704        let (decoded, _) = StageCheckpoint::from_compact(&buf, encoded);
705
706        assert_eq!(decoded, checkpoint);
707        assert_eq!(decoded.finish_stage_checkpoint().unwrap().partial_state_trie(), Some(21));
708    }
709}