reth_codecs/alloy/
trie.rs

1//! Native Compact codec impl for alloy-trie types.
2
3use crate::Compact;
4use alloc::vec::Vec;
5use alloy_primitives::{FixedBytesSliceExt, 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        let hashes = self.hashes.as_slice().as_flattened();
71        buf_size += hashes.len();
72        buf.put_slice(hashes);
73
74        buf_size
75    }
76
77    fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8]) {
78        let hash_len = B256::len_bytes();
79
80        // Assert the buffer is long enough to contain the masks and the hashes.
81        assert_eq!(buf.len() % hash_len, 6);
82
83        // Consume the masks.
84        let (state_mask, buf) = TrieMask::from_compact(buf, 0);
85        let (tree_mask, buf) = TrieMask::from_compact(buf, 0);
86        let (hash_mask, buf) = TrieMask::from_compact(buf, 0);
87
88        let mut buf = buf;
89        let mut num_hashes = buf.len() / hash_len;
90        let mut root_hash = None;
91
92        // Check if the root hash is present
93        if hash_mask.count_ones() as usize + 1 == num_hashes {
94            root_hash = Some(B256::from_slice(&buf[..hash_len]));
95            buf.advance(hash_len);
96            num_hashes -= 1;
97        }
98
99        // Consume all remaining hashes.
100        let hashes_bytes = hash_len * num_hashes;
101        let mut hashes = Vec::<B256>::with_capacity(num_hashes);
102        // SAFETY: We're copying all the bytes at once into the spare capacity, which is enough.
103        unsafe {
104            hashes.as_mut_ptr().cast::<u8>().copy_from_nonoverlapping(buf.as_ptr(), hashes_bytes);
105            hashes.set_len(num_hashes);
106            buf.advance(hashes_bytes);
107        }
108
109        (Self::new(state_mask, tree_mask, hash_mask, hashes, root_hash), buf)
110    }
111}
112
113impl Compact for TrieMask {
114    fn to_compact<B>(&self, buf: &mut B) -> usize
115    where
116        B: bytes::BufMut + AsMut<[u8]>,
117    {
118        buf.put_u16(self.get());
119        2
120    }
121
122    fn from_compact(mut buf: &[u8], _len: usize) -> (Self, &[u8]) {
123        let mask = buf.get_u16();
124        (Self::new(mask), buf)
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131    use alloy_primitives::hex;
132
133    #[test]
134    fn node_encoding() {
135        let n = BranchNodeCompact::new(
136            0xf607,
137            0x0005,
138            0x4004,
139            vec![
140                hex!("90d53cd810cc5d4243766cd4451e7b9d14b736a1148b26b3baac7617f617d321").into(),
141                hex!("cc35c964dda53ba6c0b87798073a9628dbc9cd26b5cce88eb69655a9c609caf1").into(),
142            ],
143            Some(hex!("aaaabbbb0006767767776fffffeee44444000005567645600000000eeddddddd").into()),
144        );
145
146        let mut out = Vec::new();
147        let compact_len = n.to_compact(&mut out);
148        assert_eq!(BranchNodeCompact::from_compact(&out, compact_len).0, n);
149    }
150}
reth_codecs/alloy/trie.rs

reth_codecs/alloy/
trie.rs
