feat(stages): rich checkpoint (#2701)

Co-authored-by: Roman Krasiuk <rokrassyuk@gmail.com>

crates/primitives/src/account.rs

@@ -135,7 +135,7 @@ impl Compact for Bytecode {
                     jump_map: JumpMap::from_slice(buf),
                 },
             }),
-            _ => unreachable!(),
+            _ => unreachable!("Junk data in database: unknown BytecodeState variant"),
         };
         (decoded, &[])
     }

crates/primitives/src/checkpoints.rs

@@ -1,9 +1,11 @@
 use crate::{
     trie::{hash_builder::HashBuilderState, StoredSubNode},
-    Address, BlockNumber, H256,
+    Address, BlockNumber, TxNumber, H256,
 };
-use bytes::Buf;
-use reth_codecs::{main_codec, Compact};
+use bytes::{Buf, BufMut};
+use reth_codecs::{derive_arbitrary, main_codec, Compact};
+use serde::{Deserialize, Serialize};
+use std::fmt::{Display, Formatter};
 
 /// Saves the progress of Merkle stage.
 #[derive(Default, Debug, Clone, PartialEq)]
@@ -97,7 +99,7 @@ impl Compact for MerkleCheckpoint {
 
 /// Saves the progress of AccountHashing
 #[main_codec]
-#[derive(Default, Debug, Copy, Clone, PartialEq)]
+#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
 pub struct AccountHashingCheckpoint {
     /// The next account to start hashing from
     pub address: Option<Address>,
@@ -109,7 +111,7 @@ pub struct AccountHashingCheckpoint {
 
 /// Saves the progress of StorageHashing
 #[main_codec]
-#[derive(Default, Debug, Copy, Clone, PartialEq)]
+#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
 pub struct StorageHashingCheckpoint {
     /// The next account to start hashing from
     pub address: Option<Address>,
@@ -121,6 +123,88 @@ pub struct StorageHashingCheckpoint {
     pub to: u64,
 }
 
+/// Saves the progress of a stage.
+#[main_codec]
+#[derive(Debug, Default, PartialEq, Eq, Clone, Copy)]
+pub struct StageCheckpoint {
+    /// The maximum block processed by the stage.
+    pub block_number: BlockNumber,
+    /// Stage-specific checkpoint. None if stage uses only block-based checkpoints.
+    pub stage_checkpoint: Option<StageUnitCheckpoint>,
+}
+
+impl StageCheckpoint {
+    /// Creates a new [`StageCheckpoint`] with only `block_number` set.
+    pub fn new(block_number: BlockNumber) -> Self {
+        Self { block_number, ..Default::default() }
+    }
+}
+
+// TODO(alexey): ideally, we'd want to display block number + stage-specific metric (if available)
+//  in places like logs or traces
+impl Display for StageCheckpoint {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        Display::fmt(&self.block_number, f)
+    }
+}
+
+// TODO(alexey): add a merkle checkpoint. Currently it's hard because [`MerkleCheckpoint`]
+//  is not a Copy type.
+/// Stage-specific checkpoint metrics.
+#[derive_arbitrary(compact)]
+#[derive(Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize)]
+pub enum StageUnitCheckpoint {
+    /// Saves the progress of transaction-indexed stages.
+    Transaction(TxNumber),
+    /// Saves the progress of AccountHashing stage.
+    Account(AccountHashingCheckpoint),
+    /// Saves the progress of StorageHashing stage.
+    Storage(StorageHashingCheckpoint),
+}
+
+impl Compact for StageUnitCheckpoint {
+    fn to_compact<B>(self, buf: &mut B) -> usize
+    where
+        B: BufMut + AsMut<[u8]>,
+    {
+        match self {
+            StageUnitCheckpoint::Transaction(data) => {
+                buf.put_u8(0);
+                1 + data.to_compact(buf)
+            }
+            StageUnitCheckpoint::Account(data) => {
+                buf.put_u8(1);
+                1 + data.to_compact(buf)
+            }
+            StageUnitCheckpoint::Storage(data) => {
+                buf.put_u8(2);
+                1 + data.to_compact(buf)
+            }
+        }
+    }
+
+    fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8])
+    where
+        Self: Sized,
+    {
+        match buf[0] {
+            0 => {
+                let (data, buf) = TxNumber::from_compact(&buf[1..], buf.len() - 1);
+                (Self::Transaction(data), buf)
+            }
+            1 => {
+                let (data, buf) = AccountHashingCheckpoint::from_compact(&buf[1..], buf.len() - 1);
+                (Self::Account(data), buf)
+            }
+            2 => {
+                let (data, buf) = StorageHashingCheckpoint::from_compact(&buf[1..], buf.len() - 1);
+                (Self::Storage(data), buf)
+            }
+            _ => unreachable!("Junk data in database: unknown StageUnitCheckpoint variant"),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -146,4 +230,30 @@ mod tests {
         let (decoded, _) = MerkleCheckpoint::from_compact(&buf, encoded);
         assert_eq!(decoded, checkpoint);
     }
+
+    #[test]
+    fn stage_unit_checkpoint_roundtrip() {
+        let mut rng = rand::thread_rng();
+        let checkpoints = vec![
+            StageUnitCheckpoint::Transaction(rng.gen()),
+            StageUnitCheckpoint::Account(AccountHashingCheckpoint {
+                address: Some(Address::from_low_u64_be(rng.gen())),
+                from: rng.gen(),
+                to: rng.gen(),
+            }),
+            StageUnitCheckpoint::Storage(StorageHashingCheckpoint {
+                address: Some(Address::from_low_u64_be(rng.gen())),
+                storage: Some(H256::from_low_u64_be(rng.gen())),
+                from: rng.gen(),
+                to: rng.gen(),
+            }),
+        ];
+
+        for checkpoint in checkpoints {
+            let mut buf = Vec::new();
+            let encoded = checkpoint.clone().to_compact(&mut buf);
+            let (decoded, _) = StageUnitCheckpoint::from_compact(&buf, encoded);
+            assert_eq!(decoded, checkpoint);
+        }
+    }
 }

crates/primitives/src/lib.rs

@@ -50,7 +50,10 @@ pub use chain::{
     AllGenesisFormats, Chain, ChainInfo, ChainSpec, ChainSpecBuilder, ForkCondition, GOERLI,
     MAINNET, SEPOLIA,
 };
-pub use checkpoints::{AccountHashingCheckpoint, MerkleCheckpoint, StorageHashingCheckpoint};
+pub use checkpoints::{
+    AccountHashingCheckpoint, MerkleCheckpoint, StageCheckpoint, StageUnitCheckpoint,
+    StorageHashingCheckpoint,
+};
 pub use compression::*;
 pub use constants::{
     EMPTY_OMMER_ROOT, GOERLI_GENESIS, KECCAK_EMPTY, MAINNET_GENESIS, SEPOLIA_GENESIS,

crates/primitives/src/transaction/mod.rs

@@ -261,7 +261,7 @@ impl Compact for Transaction {
                 let (tx, buf) = TxEip1559::from_compact(buf, buf.len());
                 (Transaction::Eip1559(tx), buf)
             }
-            _ => unreachable!(),
+            _ => unreachable!("Junk data in database: unknown Transaction variant"),
         }
     }
 }
@@ -717,7 +717,7 @@ impl Compact for TransactionKind {
                 let (addr, buf) = Address::from_compact(buf, buf.len());
                 (TransactionKind::Call(addr), buf)
             }
-            _ => unreachable!(),
+            _ => unreachable!("Junk data in database: unknown TransactionKind variant"),
         }
     }
 }

crates/primitives/src/trie/hash_builder/value.rs

@@ -42,7 +42,7 @@ impl Compact for HashBuilderValue {
                 let (bytes, buf) = Vec::from_compact(&buf[1..], 0);
                 (Self::Bytes(bytes), buf)
             }
-            _ => panic!("Invalid hash builder value"),
+            _ => unreachable!("Junk data in database: unknown HashBuilderValue variant"),
         }
     }
 }