reth_trie_common/hash_builder/
state.rs

1use crate::TrieMask;
2use alloy_trie::{hash_builder::HashBuilderValue, nodes::RlpNode, HashBuilder};
3use nybbles::Nibbles;
4
5/// The hash builder state for storing in the database.
6/// Check the `reth-trie` crate for more info on hash builder.
7#[derive(Debug, Clone, PartialEq, Eq, Default)]
8#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
9#[cfg_attr(
10    feature = "arbitrary",
11    derive(arbitrary::Arbitrary),
12    reth_codecs::add_arbitrary_tests(compact)
13)]
14pub struct HashBuilderState {
15    /// The current key.
16    pub key: Vec<u8>,
17    /// The current node value.
18    pub value: HashBuilderValue,
19    /// The builder stack.
20    pub stack: Vec<RlpNode>,
21
22    /// Group masks.
23    pub groups: Vec<TrieMask>,
24    /// Tree masks.
25    pub tree_masks: Vec<TrieMask>,
26    /// Hash masks.
27    pub hash_masks: Vec<TrieMask>,
28
29    /// Flag indicating if the current node is stored in the database.
30    pub stored_in_database: bool,
31}
32
33impl From<HashBuilderState> for HashBuilder {
34    fn from(state: HashBuilderState) -> Self {
35        Self {
36            key: Nibbles::from_nibbles_unchecked(state.key),
37            stack: state.stack,
38            value: state.value,
39            groups: state.groups,
40            tree_masks: state.tree_masks,
41            hash_masks: state.hash_masks,
42            stored_in_database: state.stored_in_database,
43            updated_branch_nodes: None,
44            proof_retainer: None,
45            rlp_buf: Vec::with_capacity(32),
46        }
47    }
48}
49
50impl From<HashBuilder> for HashBuilderState {
51    fn from(state: HashBuilder) -> Self {
52        Self {
53            key: state.key.into(),
54            stack: state.stack,
55            value: state.value,
56            groups: state.groups,
57            tree_masks: state.tree_masks,
58            hash_masks: state.hash_masks,
59            stored_in_database: state.stored_in_database,
60        }
61    }
62}
63
64#[cfg(any(test, feature = "reth-codec"))]
65impl reth_codecs::Compact for HashBuilderState {
66    fn to_compact<B>(&self, buf: &mut B) -> usize
67    where
68        B: bytes::BufMut + AsMut<[u8]>,
69    {
70        let mut len = 0;
71
72        len += self.key.to_compact(buf);
73
74        buf.put_u16(self.stack.len() as u16);
75        len += 2;
76        for item in &self.stack {
77            buf.put_u16(item.len() as u16);
78            buf.put_slice(&item[..]);
79            len += 2 + item.len();
80        }
81
82        len += self.value.to_compact(buf);
83
84        buf.put_u16(self.groups.len() as u16);
85        len += 2;
86        for item in &self.groups {
87            len += (*item).to_compact(buf);
88        }
89
90        buf.put_u16(self.tree_masks.len() as u16);
91        len += 2;
92        for item in &self.tree_masks {
93            len += (*item).to_compact(buf);
94        }
95
96        buf.put_u16(self.hash_masks.len() as u16);
97        len += 2;
98        for item in &self.hash_masks {
99            len += (*item).to_compact(buf);
100        }
101
102        buf.put_u8(self.stored_in_database as u8);
103        len += 1;
104        len
105    }
106
107    fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8]) {
108        use bytes::Buf;
109
110        let (key, mut buf) = Vec::from_compact(buf, 0);
111
112        let stack_len = buf.get_u16() as usize;
113        let mut stack = Vec::with_capacity(stack_len);
114        for _ in 0..stack_len {
115            let item_len = buf.get_u16() as usize;
116            stack.push(RlpNode::from_raw(&buf[..item_len]).unwrap());
117            buf.advance(item_len);
118        }
119
120        let (value, mut buf) = HashBuilderValue::from_compact(buf, 0);
121
122        let groups_len = buf.get_u16() as usize;
123        let mut groups = Vec::with_capacity(groups_len);
124        for _ in 0..groups_len {
125            let (item, rest) = TrieMask::from_compact(buf, 0);
126            groups.push(item);
127            buf = rest;
128        }
129
130        let tree_masks_len = buf.get_u16() as usize;
131        let mut tree_masks = Vec::with_capacity(tree_masks_len);
132        for _ in 0..tree_masks_len {
133            let (item, rest) = TrieMask::from_compact(buf, 0);
134            tree_masks.push(item);
135            buf = rest;
136        }
137
138        let hash_masks_len = buf.get_u16() as usize;
139        let mut hash_masks = Vec::with_capacity(hash_masks_len);
140        for _ in 0..hash_masks_len {
141            let (item, rest) = TrieMask::from_compact(buf, 0);
142            hash_masks.push(item);
143            buf = rest;
144        }
145
146        let stored_in_database = buf.get_u8() != 0;
147        (Self { key, stack, value, groups, tree_masks, hash_masks, stored_in_database }, buf)
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use reth_codecs::Compact;
155
156    #[test]
157    fn hash_builder_state_regression() {
158        let mut state = HashBuilderState::default();
159        state.stack.push(Default::default());
160        let mut buf = vec![];
161        let len = state.clone().to_compact(&mut buf);
162        let (decoded, _) = HashBuilderState::from_compact(&buf, len);
163        assert_eq!(state, decoded);
164    }
165
166    #[cfg(feature = "arbitrary")]
167    proptest::proptest! {
168        #[test]
169        fn hash_builder_state_roundtrip(state in proptest_arbitrary_interop::arb::<HashBuilderState>()) {
170            let mut buf = vec![];
171            let len = state.to_compact(&mut buf);
172            let (decoded, _) = HashBuilderState::from_compact(&buf, len);
173            assert_eq!(state, decoded);
174        }
175    }
176}