reth_trie_common/hash_builder/
state.rs

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