feat(trie): read-only root calculation (#2233)

crates/primitives/src/checkpoints.rs

@@ -1,18 +1,70 @@
-use crate::{Address, H256};
+use crate::{
+    trie::{HashBuilderState, StoredSubNode},
+    Address, H256,
+};
+use bytes::Buf;
 use reth_codecs::{main_codec, Compact};
 
-/// Saves the progress of MerkleStage
-#[main_codec]
-#[derive(Default, Debug, Copy, Clone, PartialEq)]
-pub struct ProofCheckpoint {
-    /// The next hashed account to insert into the trie.
-    pub hashed_address: Option<H256>,
-    /// The next storage entry to insert into the trie.
-    pub storage_key: Option<H256>,
-    /// Current intermediate root for `AccountsTrie`.
-    pub account_root: Option<H256>,
-    /// Current intermediate storage root from an account.
-    pub storage_root: Option<H256>,
+/// Saves the progress of Merkle stage.
+#[derive(Default, Debug, Clone, PartialEq)]
+pub struct MerkleCheckpoint {
+    // TODO: target block?
+    /// The last hashed account key processed.
+    pub last_account_key: H256,
+    /// The last walker key processed.
+    pub last_walker_key: Vec<u8>,
+    /// Previously recorded walker stack.
+    pub walker_stack: Vec<StoredSubNode>,
+    /// The hash builder state.
+    pub state: HashBuilderState,
+}
+
+impl Compact for MerkleCheckpoint {
+    fn to_compact<B>(self, buf: &mut B) -> usize
+    where
+        B: bytes::BufMut + AsMut<[u8]>,
+    {
+        let mut len = 0;
+
+        buf.put_slice(self.last_account_key.as_slice());
+        len += self.last_account_key.len();
+
+        buf.put_u16(self.last_walker_key.len() as u16);
+        buf.put_slice(&self.last_walker_key[..]);
+        len += 2 + self.last_walker_key.len();
+
+        buf.put_u16(self.walker_stack.len() as u16);
+        len += 2;
+        for item in self.walker_stack.into_iter() {
+            len += item.to_compact(buf);
+        }
+
+        len += self.state.to_compact(buf);
+        len
+    }
+
+    fn from_compact(mut buf: &[u8], _len: usize) -> (Self, &[u8])
+    where
+        Self: Sized,
+    {
+        let last_account_key = H256::from_slice(&buf[..32]);
+        buf.advance(32);
+
+        let last_walker_key_len = buf.get_u16() as usize;
+        let last_walker_key = Vec::from(&buf[..last_walker_key_len]);
+        buf.advance(last_walker_key_len);
+
+        let walker_stack_len = buf.get_u16() as usize;
+        let mut walker_stack = Vec::with_capacity(walker_stack_len);
+        for _ in 0..walker_stack_len {
+            let (item, rest) = StoredSubNode::from_compact(buf, 0);
+            walker_stack.push(item);
+            buf = rest;
+        }
+
+        let (state, buf) = HashBuilderState::from_compact(buf, 0);
+        (MerkleCheckpoint { last_account_key, last_walker_key, walker_stack, state }, buf)
+    }
 }
 
 /// Saves the progress of AccountHashing

crates/primitives/src/lib.rs

@@ -48,7 +48,7 @@ pub use chain::{
     AllGenesisFormats, Chain, ChainInfo, ChainSpec, ChainSpecBuilder, ForkCondition, GOERLI,
     MAINNET, SEPOLIA,
 };
-pub use checkpoints::{AccountHashingCheckpoint, ProofCheckpoint, StorageHashingCheckpoint};
+pub use checkpoints::{AccountHashingCheckpoint, MerkleCheckpoint, StorageHashingCheckpoint};
 pub use constants::{
     EMPTY_OMMER_ROOT, GOERLI_GENESIS, KECCAK_EMPTY, MAINNET_GENESIS, SEPOLIA_GENESIS,
 };

crates/primitives/src/trie/branch_node.rs

@@ -1,5 +1,6 @@
 use super::TrieMask;
 use crate::H256;
+use bytes::Buf;
 use reth_codecs::Compact;
 use serde::{Deserialize, Serialize};
 
@@ -88,7 +89,7 @@ impl Compact for BranchNodeCompact {
         buf_size
     }
 
-    fn from_compact(buf: &[u8], len: usize) -> (Self, &[u8])
+    fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8])
     where
         Self: Sized,
     {
@@ -98,9 +99,9 @@ impl Compact for BranchNodeCompact {
         assert_eq!(buf.len() % hash_len, 6);
 
         // Consume the masks.
-        let (state_mask, buf) = TrieMask::from_compact(buf, len);
-        let (tree_mask, buf) = TrieMask::from_compact(buf, len);
-        let (hash_mask, buf) = TrieMask::from_compact(buf, len);
+        let (state_mask, buf) = TrieMask::from_compact(buf, 0);
+        let (tree_mask, buf) = TrieMask::from_compact(buf, 0);
+        let (hash_mask, buf) = TrieMask::from_compact(buf, 0);
 
         let mut buf = buf;
         let mut num_hashes = buf.len() / hash_len;
@@ -108,18 +109,16 @@ impl Compact for BranchNodeCompact {
 
         // Check if the root hash is present
         if hash_mask.count_ones() as usize + 1 == num_hashes {
-            let (hash, remaining) = H256::from_compact(buf, hash_len);
-            root_hash = Some(hash);
-            buf = remaining;
+            root_hash = Some(H256::from_slice(&buf[..hash_len]));
+            buf.advance(hash_len);
             num_hashes -= 1;
         }
 
         // Consume all remaining hashes.
         let mut hashes = Vec::<H256>::with_capacity(num_hashes);
         for _ in 0..num_hashes {
-            let (hash, remaining) = H256::from_compact(buf, hash_len);
-            hashes.push(hash);
-            buf = remaining;
+            hashes.push(H256::from_slice(&buf[..hash_len]));
+            buf.advance(hash_len);
         }
 
         (Self::new(state_mask, tree_mask, hash_mask, hashes, root_hash), buf)

crates/primitives/src/trie/hash_builder.rs

@@ -0,0 +1,219 @@
+use super::TrieMask;
+use crate::H256;
+use bytes::Buf;
+use reth_codecs::{derive_arbitrary, Compact};
+use serde::{Deserialize, Serialize};
+
+/// The hash builder state for storing in the database.
+/// Check the `reth-trie` crate for more info on hash builder.
+#[derive_arbitrary(compact)]
+#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
+pub struct HashBuilderState {
+    /// The current key.
+    pub key: Vec<u8>,
+    /// The builder stack.
+    pub stack: Vec<Vec<u8>>,
+    /// The current node value.
+    pub value: HashBuilderValue,
+
+    /// Group masks.
+    pub groups: Vec<TrieMask>,
+    /// Tree masks.
+    pub tree_masks: Vec<TrieMask>,
+    /// Hash masks.
+    pub hash_masks: Vec<TrieMask>,
+
+    /// Flag indicating if the current node is stored in the database.
+    pub stored_in_database: bool,
+}
+
+impl Compact for HashBuilderState {
+    fn to_compact<B>(self, buf: &mut B) -> usize
+    where
+        B: bytes::BufMut + AsMut<[u8]>,
+    {
+        let mut len = 0;
+
+        len += self.key.to_compact(buf);
+
+        buf.put_u16(self.stack.len() as u16);
+        len += 2;
+        for item in self.stack.iter() {
+            buf.put_u16(item.len() as u16);
+            buf.put_slice(&item[..]);
+            len += 2 + item.len();
+        }
+
+        len += self.value.to_compact(buf);
+
+        buf.put_u16(self.groups.len() as u16);
+        len += 2;
+        for item in self.groups.iter() {
+            len += item.to_compact(buf);
+        }
+
+        buf.put_u16(self.tree_masks.len() as u16);
+        len += 2;
+        for item in self.tree_masks.iter() {
+            len += item.to_compact(buf);
+        }
+
+        buf.put_u16(self.hash_masks.len() as u16);
+        len += 2;
+        for item in self.hash_masks.iter() {
+            len += item.to_compact(buf);
+        }
+
+        buf.put_u8(self.stored_in_database as u8);
+        len += 1;
+        len
+    }
+
+    fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8])
+    where
+        Self: Sized,
+    {
+        let (key, mut buf) = Vec::from_compact(buf, 0);
+
+        let stack_len = buf.get_u16() as usize;
+        let mut stack = Vec::with_capacity(stack_len);
+        for _ in 0..stack_len {
+            let item_len = buf.get_u16() as usize;
+            stack.push(Vec::from(&buf[..item_len]));
+            buf.advance(item_len);
+        }
+
+        let (value, mut buf) = HashBuilderValue::from_compact(buf, 0);
+
+        let groups_len = buf.get_u16() as usize;
+        let mut groups = Vec::with_capacity(groups_len);
+        for _ in 0..groups_len {
+            let (item, rest) = TrieMask::from_compact(buf, 0);
+            groups.push(item);
+            buf = rest;
+        }
+
+        let tree_masks_len = buf.get_u16() as usize;
+        let mut tree_masks = Vec::with_capacity(tree_masks_len);
+        for _ in 0..tree_masks_len {
+            let (item, rest) = TrieMask::from_compact(buf, 0);
+            tree_masks.push(item);
+            buf = rest;
+        }
+
+        let hash_masks_len = buf.get_u16() as usize;
+        let mut hash_masks = Vec::with_capacity(hash_masks_len);
+        for _ in 0..hash_masks_len {
+            let (item, rest) = TrieMask::from_compact(buf, 0);
+            hash_masks.push(item);
+            buf = rest;
+        }
+
+        let stored_in_database = buf.get_u8() != 0;
+        (Self { key, stack, value, groups, tree_masks, hash_masks, stored_in_database }, buf)
+    }
+}
+
+/// The current value of the hash builder.
+#[derive_arbitrary(compact)]
+#[derive(Clone, PartialEq, Serialize, Deserialize)]
+pub enum HashBuilderValue {
+    /// Value of the leaf node.
+    Hash(H256),
+    /// Hash of adjacent nodes.
+    Bytes(Vec<u8>),
+}
+
+impl Compact for HashBuilderValue {
+    fn to_compact<B>(self, buf: &mut B) -> usize
+    where
+        B: bytes::BufMut + AsMut<[u8]>,
+    {
+        match self {
+            Self::Hash(hash) => {
+                buf.put_u8(0);
+                1 + hash.to_compact(buf)
+            }
+            Self::Bytes(bytes) => {
+                buf.put_u8(1);
+                1 + bytes.to_compact(buf)
+            }
+        }
+    }
+
+    fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8])
+    where
+        Self: Sized,
+    {
+        match buf[0] {
+            0 => {
+                let (hash, buf) = H256::from_compact(&buf[1..], 32);
+                (Self::Hash(hash), buf)
+            }
+            1 => {
+                let (bytes, buf) = Vec::from_compact(&buf[1..], 0);
+                (Self::Bytes(bytes), buf)
+            }
+            _ => panic!("Invalid hash builder value"),
+        }
+    }
+}
+
+impl std::fmt::Debug for HashBuilderValue {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Bytes(bytes) => write!(f, "Bytes({:?})", hex::encode(bytes)),
+            Self::Hash(hash) => write!(f, "Hash({:?})", hash),
+        }
+    }
+}
+
+impl From<Vec<u8>> for HashBuilderValue {
+    fn from(value: Vec<u8>) -> Self {
+        Self::Bytes(value)
+    }
+}
+
+impl From<&[u8]> for HashBuilderValue {
+    fn from(value: &[u8]) -> Self {
+        Self::Bytes(value.to_vec())
+    }
+}
+
+impl From<H256> for HashBuilderValue {
+    fn from(value: H256) -> Self {
+        Self::Hash(value)
+    }
+}
+
+impl Default for HashBuilderValue {
+    fn default() -> Self {
+        Self::Bytes(vec![])
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use proptest::prelude::*;
+
+    #[test]
+    fn hash_builder_state_regression() {
+        let mut state = HashBuilderState::default();
+        state.stack.push(vec![]);
+        let mut buf = vec![];
+        let len = state.clone().to_compact(&mut buf);
+        let (decoded, _) = HashBuilderState::from_compact(&buf, len);
+        assert_eq!(state, decoded);
+    }
+
+    proptest! {
+        #[test]
+        fn hash_builder_state_roundtrip(state: HashBuilderState) {
+            let mut buf = vec![];
+            let len = state.clone().to_compact(&mut buf);
+            let (decoded, _) = HashBuilderState::from_compact(&buf, len);
+            assert_eq!(state, decoded);
+        }
+    }
+}

crates/primitives/src/trie/mask.rs

@@ -1,5 +1,6 @@
+use bytes::Buf;
 use derive_more::{BitAnd, BitAndAssign, BitOr, BitOrAssign, Deref, From, Not};
-use reth_codecs::Compact;
+use reth_codecs::{derive_arbitrary, Compact};
 use serde::{Deserialize, Serialize};
 
 /// A struct representing a mask of 16 bits, used for Ethereum trie operations.
@@ -26,6 +27,7 @@ use serde::{Deserialize, Serialize};
     BitOrAssign,
     Not,
 )]
+#[derive_arbitrary(compact)]
 pub struct TrieMask(u16);
 
 impl TrieMask {
@@ -66,14 +68,15 @@ impl Compact for TrieMask {
     where
         B: bytes::BufMut + AsMut<[u8]>,
     {
-        buf.put_slice(self.to_be_bytes().as_slice());
+        buf.put_u16(self.0);
         2
     }
 
-    fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8])
+    fn from_compact(mut buf: &[u8], _len: usize) -> (Self, &[u8])
     where
         Self: Sized,
     {
-        (Self(u16::from_be_bytes(buf[..2].try_into().unwrap())), &buf[2..])
+        let mask = buf.get_u16();
+        (Self(mask), buf)
     }
 }

crates/primitives/src/trie/mod.rs

@@ -1,13 +1,17 @@
 //! Collection of trie related types.
 
 mod branch_node;
+mod hash_builder;
 mod mask;
 mod nibbles;
 mod storage;
+mod subnode;
 
 pub use self::{
     branch_node::BranchNodeCompact,
+    hash_builder::{HashBuilderState, HashBuilderValue},
     mask::TrieMask,
     nibbles::{StoredNibbles, StoredNibblesSubKey},
     storage::StorageTrieEntry,
+    subnode::StoredSubNode,
 };

crates/primitives/src/trie/subnode.rs

@@ -0,0 +1,98 @@
+use super::BranchNodeCompact;
+use bytes::Buf;
+use reth_codecs::Compact;
+
+/// Walker sub node for storing intermediate state root calculation state in the database.
+/// See [crate::MerkleCheckpoint].
+#[derive(Debug, Clone, PartialEq, Default)]
+pub struct StoredSubNode {
+    /// The key of the current node.
+    pub key: Vec<u8>,
+    /// The index of the next child to visit.
+    pub nibble: Option<u8>,
+    /// The node itself.
+    pub node: Option<BranchNodeCompact>,
+}
+
+impl Compact for StoredSubNode {
+    fn to_compact<B>(self, buf: &mut B) -> usize
+    where
+        B: bytes::BufMut + AsMut<[u8]>,
+    {
+        let mut len = 0;
+
+        buf.put_u16(self.key.len() as u16);
+        buf.put_slice(&self.key[..]);
+        len += 2 + self.key.len();
+
+        if let Some(nibble) = self.nibble {
+            buf.put_u8(1);
+            buf.put_u8(nibble);
+            len += 2;
+        } else {
+            buf.put_u8(0);
+            len += 1;
+        }
+
+        if let Some(node) = self.node {
+            buf.put_u8(1);
+            len += 1;
+            len += node.to_compact(buf);
+        } else {
+            len += 1;
+            buf.put_u8(0);
+        }
+
+        len
+    }
+
+    fn from_compact(mut buf: &[u8], _len: usize) -> (Self, &[u8])
+    where
+        Self: Sized,
+    {
+        let key_len = buf.get_u16() as usize;
+        let key = Vec::from(&buf[..key_len]);
+        buf.advance(key_len);
+
+        let nibbles_exists = buf.get_u8() != 0;
+        let nibble = if nibbles_exists { Some(buf.get_u8()) } else { None };
+
+        let node_exsists = buf.get_u8() != 0;
+        let node = if node_exsists {
+            let (node, rest) = BranchNodeCompact::from_compact(buf, 0);
+            buf = rest;
+            Some(node)
+        } else {
+            None
+        };
+
+        (StoredSubNode { key, nibble, node }, buf)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{trie::TrieMask, H256};
+
+    #[test]
+    fn subnode_roundtrip() {
+        let subnode = StoredSubNode {
+            key: vec![],
+            nibble: None,
+            node: Some(BranchNodeCompact {
+                state_mask: TrieMask::new(1),
+                tree_mask: TrieMask::new(0),
+                hash_mask: TrieMask::new(1),
+                hashes: vec![H256::zero()],
+                root_hash: None,
+            }),
+        };
+
+        let mut encoded = vec![];
+        subnode.clone().to_compact(&mut encoded);
+        let (decoded, _) = StoredSubNode::from_compact(&encoded[..], 0);
+
+        assert_eq!(subnode, decoded);
+    }
+}