reth_codecs/alloy/
trie.rs

Help
1//! Native Compact codec impl for alloy-trie types.
2
3use crate::Compact;
4use alloc::vec::Vec;
5use alloy_primitives::B256;
6use alloy_trie::{
7    hash_builder::{HashBuilderValue, HashBuilderValueRef},
8    BranchNodeCompact, TrieMask,
9};
10use bytes::{Buf, BufMut};
11
12/// Identifier for [`HashBuilderValueRef::Hash`]
13const HASH_BUILDER_TYPE_HASH: u8 = 0;
14
15/// Identifier for [`HashBuilderValueRef::Bytes`]
16const HASH_BUILDER_TYPE_BYTES: u8 = 1;
17
18impl Compact for HashBuilderValue {
19    fn to_compact<B>(&self, buf: &mut B) -> usize
20    where
21        B: BufMut + AsMut<[u8]>,
22    {
23        match self.as_ref() {
24            HashBuilderValueRef::Hash(hash) => {
25                buf.put_u8(HASH_BUILDER_TYPE_HASH);
26                1 + hash.to_compact(buf)
27            }
28            HashBuilderValueRef::Bytes(bytes) => {
29                buf.put_u8(HASH_BUILDER_TYPE_BYTES);
30                1 + bytes.to_compact(buf)
31            }
32        }
33    }
34
35    fn from_compact(mut buf: &[u8], _: usize) -> (Self, &[u8]) {
36        let mut this = Self::default();
37        let buf = match buf.get_u8() {
38            HASH_BUILDER_TYPE_HASH => {
39                let (hash, buf) = B256::from_compact(buf, 32);
40                this.set_from_ref(HashBuilderValueRef::Hash(&hash));
41                buf
42            }
43            HASH_BUILDER_TYPE_BYTES => {
44                let (bytes, buf) = Vec::from_compact(buf, 0);
45                this.set_bytes_owned(bytes);
46                buf
47            }
48            _ => unreachable!("Junk data in database: unknown HashBuilderValue variant"),
49        };
50        (this, buf)
51    }
52}
53
54impl Compact for BranchNodeCompact {
55    fn to_compact<B>(&self, buf: &mut B) -> usize
56    where
57        B: bytes::BufMut + AsMut<[u8]>,
58    {
59        let mut buf_size = 0;
60
61        buf_size += self.state_mask.to_compact(buf);
62        buf_size += self.tree_mask.to_compact(buf);
63        buf_size += self.hash_mask.to_compact(buf);
64
65        if let Some(root_hash) = self.root_hash {
66            buf_size += B256::len_bytes();
67            buf.put_slice(root_hash.as_slice());
68        }
69
70        for hash in self.hashes.iter() {
71            buf_size += B256::len_bytes();
72            buf.put_slice(hash.as_slice());
73        }
74
75        buf_size
76    }
77
78    fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8]) {
79        let hash_len = B256::len_bytes();
80
81        // Assert the buffer is long enough to contain the masks and the hashes.
82        assert_eq!(buf.len() % hash_len, 6);
83
84        // Consume the masks.
85        let (state_mask, buf) = TrieMask::from_compact(buf, 0);
86        let (tree_mask, buf) = TrieMask::from_compact(buf, 0);
87        let (hash_mask, buf) = TrieMask::from_compact(buf, 0);
88
89        let mut buf = buf;
90        let mut num_hashes = buf.len() / hash_len;
91        let mut root_hash = None;
92
93        // Check if the root hash is present
94        if hash_mask.count_ones() as usize + 1 == num_hashes {
95            root_hash = Some(B256::from_slice(&buf[..hash_len]));
96            buf.advance(hash_len);
97            num_hashes -= 1;
98        }
99
100        // Consume all remaining hashes.
101        let mut hashes = Vec::<B256>::with_capacity(num_hashes);
102        for _ in 0..num_hashes {
103            hashes.push(B256::from_slice(&buf[..hash_len]));
104            buf.advance(hash_len);
105        }
106
107        (Self::new(state_mask, tree_mask, hash_mask, hashes, root_hash), buf)
108    }
109}
110
111impl Compact for TrieMask {
112    fn to_compact<B>(&self, buf: &mut B) -> usize
113    where
114        B: bytes::BufMut + AsMut<[u8]>,
115    {
116        buf.put_u16(self.get());
117        2
118    }
119
120    fn from_compact(mut buf: &[u8], _len: usize) -> (Self, &[u8]) {
121        let mask = buf.get_u16();
122        (Self::new(mask), buf)
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use alloy_primitives::hex;
130
131    #[test]
132    fn node_encoding() {
133        let n = BranchNodeCompact::new(
134            0xf607,
135            0x0005,
136            0x4004,
137            vec![
138                hex!("90d53cd810cc5d4243766cd4451e7b9d14b736a1148b26b3baac7617f617d321").into(),
139                hex!("cc35c964dda53ba6c0b87798073a9628dbc9cd26b5cce88eb69655a9c609caf1").into(),
140            ],
141            Some(hex!("aaaabbbb0006767767776fffffeee44444000005567645600000000eeddddddd").into()),
142        );
143
144        let mut out = Vec::new();
145        let compact_len = n.to_compact(&mut out);
146        assert_eq!(BranchNodeCompact::from_compact(&out, compact_len).0, n);
147    }
148}
reth_codecs/alloy/trie.rs

reth_codecs/alloy/
trie.rs
