Skip to main content

reth_trie/proof_v2/
value.rs

1//! Generic value encoder types for proof calculation with lazy evaluation.
2
3use crate::{
4    hashed_cursor::HashedCursorFactory, proof_v2::ProofCalculator, trie_cursor::TrieCursorFactory,
5};
6use alloy_primitives::{B256, U256};
7use alloy_rlp::Encodable;
8use reth_execution_errors::trie::StateProofError;
9use reth_primitives_traits::Account;
10use std::rc::Rc;
11
12/// A trait for deferred RLP-encoding of leaf values.
13pub trait DeferredValueEncoder {
14    /// RLP encodes the value into the provided buffer.
15    ///
16    /// # Arguments
17    ///
18    /// * `buf` - A mutable buffer to encode the data into
19    fn encode(self, buf: &mut Vec<u8>) -> Result<(), StateProofError>;
20}
21
22/// A trait for RLP-encoding values for proof calculation. This trait is designed to allow the lazy
23/// computation of leaf values in a generic way.
24///
25/// When calculating a leaf value in a storage trie the [`DeferredValueEncoder`] simply holds onto
26/// the slot value, and the `encode` method synchronously RLP-encodes it.
27///
28/// When calculating a leaf value in the accounts trie we create a [`DeferredValueEncoder`] to
29/// initiate any asynchronous computation of the account's storage root we want to do. Later we call
30/// [`DeferredValueEncoder::encode`] to obtain the result of that computation and RLP-encode it.
31pub trait LeafValueEncoder {
32    /// The type of value being encoded (e.g., U256 for storage, Account for accounts).
33    type Value;
34
35    /// The type that will compute and encode the value when needed.
36    type DeferredEncoder: DeferredValueEncoder;
37
38    /// Returns an encoder that will RLP-encode the value when its `encode` method is called.
39    ///
40    /// # Arguments
41    ///
42    /// * `key` - The key the value was stored at in the DB
43    /// * `value` - The value to encode
44    ///
45    /// The returned deferred encoder will be called as late as possible in the algorithm to
46    /// maximize the time available for parallel computation (e.g., storage root calculation).
47    fn deferred_encoder(&mut self, key: B256, value: Self::Value) -> Self::DeferredEncoder;
48}
49
50/// An encoder for storage slot values.
51///
52/// This encoder simply RLP-encodes U256 storage values directly.
53#[derive(Debug, Clone, Copy, Default)]
54pub struct StorageValueEncoder;
55
56/// The deferred encoder for a storage slot value.
57#[derive(Debug, Clone, Copy)]
58pub struct StorageDeferredValueEncoder(U256);
59
60impl DeferredValueEncoder for StorageDeferredValueEncoder {
61    fn encode(self, buf: &mut Vec<u8>) -> Result<(), StateProofError> {
62        self.0.encode(buf);
63        Ok(())
64    }
65}
66
67impl LeafValueEncoder for StorageValueEncoder {
68    type Value = U256;
69    type DeferredEncoder = StorageDeferredValueEncoder;
70
71    fn deferred_encoder(&mut self, _key: B256, value: Self::Value) -> Self::DeferredEncoder {
72        StorageDeferredValueEncoder(value)
73    }
74}
75
76/// An account value encoder that synchronously computes storage roots.
77///
78/// This encoder contains factories for creating trie and hashed cursors. Storage roots are
79/// computed synchronously within the deferred encoder using a `StorageProofCalculator`.
80#[derive(Debug, Clone)]
81pub struct SyncAccountValueEncoder<T, H> {
82    /// Factory for creating trie cursors.
83    trie_cursor_factory: Rc<T>,
84    /// Factory for creating hashed cursors.
85    hashed_cursor_factory: Rc<H>,
86}
87
88impl<T, H> SyncAccountValueEncoder<T, H> {
89    /// Create a new account value encoder with the given factories.
90    pub fn new(trie_cursor_factory: T, hashed_cursor_factory: H) -> Self {
91        Self {
92            trie_cursor_factory: Rc::new(trie_cursor_factory),
93            hashed_cursor_factory: Rc::new(hashed_cursor_factory),
94        }
95    }
96}
97
98/// The deferred encoder for an account value with synchronous storage root calculation.
99#[derive(Debug, Clone)]
100pub struct SyncAccountDeferredValueEncoder<T, H> {
101    trie_cursor_factory: Rc<T>,
102    hashed_cursor_factory: Rc<H>,
103    hashed_address: B256,
104    account: Account,
105}
106
107impl<T, H> DeferredValueEncoder for SyncAccountDeferredValueEncoder<T, H>
108where
109    T: TrieCursorFactory,
110    H: HashedCursorFactory,
111{
112    fn encode(self, buf: &mut Vec<u8>) -> Result<(), StateProofError> {
113        let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?;
114        let hashed_cursor =
115            self.hashed_cursor_factory.hashed_storage_cursor(self.hashed_address)?;
116
117        let mut storage_proof_calculator = ProofCalculator::new_storage(trie_cursor, hashed_cursor);
118        let root_node = storage_proof_calculator.storage_root_node(self.hashed_address)?;
119        let storage_root = storage_proof_calculator
120            .compute_root_hash(&[root_node])?
121            .expect("storage_root_node returns a node at empty path");
122
123        let trie_account = self.account.into_trie_account(storage_root);
124        trie_account.encode(buf);
125
126        Ok(())
127    }
128}
129
130impl<T, H> LeafValueEncoder for SyncAccountValueEncoder<T, H>
131where
132    T: TrieCursorFactory,
133    H: HashedCursorFactory,
134{
135    type Value = Account;
136    type DeferredEncoder = SyncAccountDeferredValueEncoder<T, H>;
137
138    fn deferred_encoder(
139        &mut self,
140        hashed_address: B256,
141        account: Self::Value,
142    ) -> Self::DeferredEncoder {
143        // Return a deferred encoder that will synchronously compute the storage root when encode()
144        // is called.
145        SyncAccountDeferredValueEncoder {
146            trie_cursor_factory: self.trie_cursor_factory.clone(),
147            hashed_cursor_factory: self.hashed_cursor_factory.clone(),
148            hashed_address,
149            account,
150        }
151    }
152}