reth_trie/proof/

mod.rs

use crate::{
    hashed_cursor::{
        HashedCursorFactory, HashedCursorMetricsCache, HashedStorageCursor,
        InstrumentedHashedCursor,
    },
    node_iter::{TrieElement, TrieNodeIter},
    prefix_set::{PrefixSetMut, TriePrefixSetsMut},
    proof_v2::{self, SyncAccountValueEncoder},
    trie_cursor::{InstrumentedTrieCursor, TrieCursorFactory, TrieCursorMetricsCache},
    walker::TrieWalker,
    HashBuilder, Nibbles, TRIE_ACCOUNT_RLP_MAX_SIZE,
};
use alloy_primitives::{
    keccak256,
    map::{B256Map, B256Set, HashSet},
    Address, B256,
};
use alloy_rlp::{BufMut, Encodable};
use alloy_trie::proof::AddedRemovedKeys;
use reth_execution_errors::trie::StateProofError;
use reth_trie_common::{
    proof::ProofRetainer, AccountProof, BranchNodeMasks, BranchNodeMasksMap, DecodedMultiProofV2,
    MultiProof, MultiProofTargets, MultiProofTargetsV2, StorageMultiProof,
};

/// A struct for generating merkle proofs.
///
/// Proof generator adds the target address and slots to the prefix set, enables the proof retainer
/// on the hash builder and follows the same algorithm as the state root calculator.
/// See `StateRoot::root` for more info.
#[derive(Debug)]
pub struct Proof<T, H, K = AddedRemovedKeys> {
    /// The factory for traversing trie nodes.
    trie_cursor_factory: T,
    /// The factory for hashed cursors.
    hashed_cursor_factory: H,
    /// A set of prefix sets that have changes.
    prefix_sets: TriePrefixSetsMut,
    /// Flag indicating whether to include branch node masks in the proof.
    collect_branch_node_masks: bool,
    /// Added and removed keys for proof retention.
    added_removed_keys: Option<K>,
}

impl<T, H> Proof<T, H> {
    /// Create a new [`Proof`] instance.
    pub fn new(t: T, h: H) -> Self {
        Self {
            trie_cursor_factory: t,
            hashed_cursor_factory: h,
            prefix_sets: TriePrefixSetsMut::default(),
            collect_branch_node_masks: false,
            added_removed_keys: None,
        }
    }
}

impl<T, H, K> Proof<T, H, K> {
    /// Set the trie cursor factory.
    pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> Proof<TF, H, K> {
        Proof {
            trie_cursor_factory,
            hashed_cursor_factory: self.hashed_cursor_factory,
            prefix_sets: self.prefix_sets,
            collect_branch_node_masks: self.collect_branch_node_masks,
            added_removed_keys: self.added_removed_keys,
        }
    }

    /// Set the hashed cursor factory.
    pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> Proof<T, HF, K> {
        Proof {
            trie_cursor_factory: self.trie_cursor_factory,
            hashed_cursor_factory,
            prefix_sets: self.prefix_sets,
            collect_branch_node_masks: self.collect_branch_node_masks,
            added_removed_keys: self.added_removed_keys,
        }
    }

    /// Set the prefix sets. They have to be mutable in order to allow extension with proof target.
    pub fn with_prefix_sets_mut(mut self, prefix_sets: TriePrefixSetsMut) -> Self {
        self.prefix_sets = prefix_sets;
        self
    }

    /// Set the flag indicating whether to include branch node masks in the proof.
    pub const fn with_branch_node_masks(mut self, branch_node_masks: bool) -> Self {
        self.collect_branch_node_masks = branch_node_masks;
        self
    }

    /// Configures the proof to retain certain nodes which would otherwise fall outside the target
    /// set, when those nodes might be required to calculate the state root when keys have been
    /// added or removed to the trie.
    ///
    /// If None is given then retention of extra proofs is disabled.
    pub fn with_added_removed_keys<K2>(self, added_removed_keys: Option<K2>) -> Proof<T, H, K2> {
        Proof {
            trie_cursor_factory: self.trie_cursor_factory,
            hashed_cursor_factory: self.hashed_cursor_factory,
            prefix_sets: self.prefix_sets,
            collect_branch_node_masks: self.collect_branch_node_masks,
            added_removed_keys,
        }
    }

    /// Get a reference to the trie cursor factory.
    pub const fn trie_cursor_factory(&self) -> &T {
        &self.trie_cursor_factory
    }

    /// Get a reference to the hashed cursor factory.
    pub const fn hashed_cursor_factory(&self) -> &H {
        &self.hashed_cursor_factory
    }
}

impl<T, H, K> Proof<T, H, K>
where
    T: TrieCursorFactory + Clone,
    H: HashedCursorFactory + Clone,
    K: AsRef<AddedRemovedKeys>,
{
    /// Generate an account proof from intermediate nodes.
    pub fn account_proof(
        self,
        address: Address,
        slots: &[B256],
    ) -> Result<AccountProof, StateProofError> {
        Ok(self
            .multiproof(MultiProofTargets::from_iter([(
                keccak256(address),
                slots.iter().map(keccak256).collect(),
            )]))?
            .account_proof(address, slots)?)
    }

    /// Generate a state multiproof using the V2 proof calculator.
    ///
    /// This method uses `ProofCalculator` with `SyncAccountValueEncoder` for account proofs
    /// and `StorageProofCalculator` for storage proofs.
    pub fn multiproof_v2(
        self,
        targets: MultiProofTargetsV2,
    ) -> Result<DecodedMultiProofV2, StateProofError> {
        let MultiProofTargetsV2 { mut account_targets, storage_targets } = targets;

        let storage_prefix_sets: B256Map<_> = self
            .prefix_sets
            .storage_prefix_sets
            .into_iter()
            .map(|(addr, ps)| (addr, ps.freeze()))
            .collect();

        // Compute account proofs using the V2 proof calculator with sync account encoding.
        let account_trie_cursor = self.trie_cursor_factory.account_trie_cursor()?;
        let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?;
        let mut account_value_encoder = SyncAccountValueEncoder::new(
            self.trie_cursor_factory.clone(),
            self.hashed_cursor_factory.clone(),
        )
        .with_storage_prefix_sets(storage_prefix_sets.clone());
        let mut account_calculator =
            proof_v2::ProofCalculator::new(account_trie_cursor, hashed_account_cursor)
                .with_prefix_set(self.prefix_sets.account_prefix_set.freeze());
        let account_proofs =
            account_calculator.proof(&mut account_value_encoder, &mut account_targets)?;

        // Compute storage proofs for each targeted account.
        let mut storage_proofs =
            B256Map::with_capacity_and_hasher(storage_targets.len(), Default::default());
        for (hashed_address, mut targets) in storage_targets {
            let storage_trie_cursor =
                self.trie_cursor_factory.storage_trie_cursor(hashed_address)?;
            let hashed_storage_cursor =
                self.hashed_cursor_factory.hashed_storage_cursor(hashed_address)?;
            let mut storage_calculator = proof_v2::StorageProofCalculator::new_storage(
                storage_trie_cursor,
                hashed_storage_cursor,
            );
            if let Some(prefix_set) = storage_prefix_sets.get(&hashed_address) {
                storage_calculator = storage_calculator.with_prefix_set(prefix_set.clone());
            }
            let proofs = storage_calculator.storage_proof(hashed_address, &mut targets)?;
            storage_proofs.insert(hashed_address, proofs);
        }

        Ok(DecodedMultiProofV2 { account_proofs, storage_proofs })
    }

    /// Generate a state multiproof according to specified targets.
    pub fn multiproof(
        mut self,
        mut targets: MultiProofTargets,
    ) -> Result<MultiProof, StateProofError> {
        let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?;
        let trie_cursor = self.trie_cursor_factory.account_trie_cursor()?;

        // Create the walker.
        let mut prefix_set = self.prefix_sets.account_prefix_set.clone();
        prefix_set.extend_keys(targets.keys().map(Nibbles::unpack));
        let walker =
            TrieWalker::<_, AddedRemovedKeys>::state_trie(trie_cursor, prefix_set.freeze())
                .with_added_removed_keys(self.added_removed_keys.as_ref());

        // Create a hash builder to rebuild the root node since it is not available in the database.
        let retainer: ProofRetainer = targets.keys().map(Nibbles::unpack).collect();
        let retainer = retainer.with_added_removed_keys(self.added_removed_keys.as_ref());
        let mut hash_builder = HashBuilder::default()
            .with_proof_retainer(retainer)
            .with_updates(self.collect_branch_node_masks);

        // Initialize all storage multiproofs as empty.
        // Storage multiproofs for non-empty tries will be overwritten if necessary.
        let mut storages: B256Map<_> =
            targets.keys().map(|key| (*key, StorageMultiProof::empty())).collect();
        let mut account_rlp = Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE);
        let mut account_node_iter = TrieNodeIter::state_trie(walker, hashed_account_cursor);
        while let Some(account_node) = account_node_iter.try_next()? {
            match account_node {
                TrieElement::Branch(node) => {
                    hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
                }
                TrieElement::Leaf(hashed_address, account) => {
                    let proof_targets = targets.remove(&hashed_address);
                    let leaf_is_proof_target = proof_targets.is_some();
                    let collect_storage_masks =
                        self.collect_branch_node_masks && leaf_is_proof_target;
                    let storage_prefix_set = self
                        .prefix_sets
                        .storage_prefix_sets
                        .remove(&hashed_address)
                        .unwrap_or_default();
                    let storage_multiproof = StorageProof::new_hashed(
                        self.trie_cursor_factory.clone(),
                        self.hashed_cursor_factory.clone(),
                        hashed_address,
                    )
                    .with_prefix_set_mut(storage_prefix_set)
                    .with_branch_node_masks(collect_storage_masks)
                    .storage_multiproof(proof_targets.unwrap_or_default())?;

                    // Encode account
                    account_rlp.clear();
                    let account = account.into_trie_account(storage_multiproof.root);
                    account.encode(&mut account_rlp as &mut dyn BufMut);

                    hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp);

                    // We might be adding leaves that are not necessarily our proof targets.
                    if leaf_is_proof_target {
                        // Overwrite storage multiproof.
                        storages.insert(hashed_address, storage_multiproof);
                    }
                }
            }
        }
        let _ = hash_builder.root();
        let account_subtree = hash_builder.take_proof_nodes();
        let branch_node_masks = if self.collect_branch_node_masks {
            let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default();
            updated_branch_nodes
                .into_iter()
                .map(|(path, node)| {
                    (path, BranchNodeMasks { hash_mask: node.hash_mask, tree_mask: node.tree_mask })
                })
                .collect()
        } else {
            BranchNodeMasksMap::default()
        };

        Ok(MultiProof { account_subtree, branch_node_masks, storages })
    }
}

/// Generates storage merkle proofs.
#[derive(Debug)]
pub struct StorageProof<'a, T, H, K = AddedRemovedKeys> {
    /// The factory for traversing trie nodes.
    trie_cursor_factory: T,
    /// The factory for hashed cursors.
    hashed_cursor_factory: H,
    /// The hashed address of an account.
    hashed_address: B256,
    /// The set of storage slot prefixes that have changed.
    prefix_set: PrefixSetMut,
    /// Flag indicating whether to include branch node masks in the proof.
    collect_branch_node_masks: bool,
    /// Provided by the user to give the necessary context to retain extra proofs.
    added_removed_keys: Option<K>,
    /// Optional reference to accumulate trie cursor metrics.
    trie_cursor_metrics: Option<&'a mut TrieCursorMetricsCache>,
    /// Optional reference to accumulate hashed cursor metrics.
    hashed_cursor_metrics: Option<&'a mut HashedCursorMetricsCache>,
}

impl<T, H> StorageProof<'static, T, H> {
    /// Create a new [`StorageProof`] instance.
    pub fn new(t: T, h: H, address: Address) -> Self {
        Self::new_hashed(t, h, keccak256(address))
    }

    /// Create a new [`StorageProof`] instance with hashed address.
    pub fn new_hashed(t: T, h: H, hashed_address: B256) -> Self {
        Self {
            trie_cursor_factory: t,
            hashed_cursor_factory: h,
            hashed_address,
            prefix_set: PrefixSetMut::default(),
            collect_branch_node_masks: false,
            added_removed_keys: None,
            trie_cursor_metrics: None,
            hashed_cursor_metrics: None,
        }
    }
}

impl<'a, T, H, K> StorageProof<'a, T, H, K> {
    /// Set the trie cursor factory.
    pub fn with_trie_cursor_factory<TF>(
        self,
        trie_cursor_factory: TF,
    ) -> StorageProof<'a, TF, H, K> {
        StorageProof {
            trie_cursor_factory,
            hashed_cursor_factory: self.hashed_cursor_factory,
            hashed_address: self.hashed_address,
            prefix_set: self.prefix_set,
            collect_branch_node_masks: self.collect_branch_node_masks,
            added_removed_keys: self.added_removed_keys,
            trie_cursor_metrics: self.trie_cursor_metrics,
            hashed_cursor_metrics: self.hashed_cursor_metrics,
        }
    }

    /// Set the hashed cursor factory.
    pub fn with_hashed_cursor_factory<HF>(
        self,
        hashed_cursor_factory: HF,
    ) -> StorageProof<'a, T, HF, K> {
        StorageProof {
            trie_cursor_factory: self.trie_cursor_factory,
            hashed_cursor_factory,
            hashed_address: self.hashed_address,
            prefix_set: self.prefix_set,
            collect_branch_node_masks: self.collect_branch_node_masks,
            added_removed_keys: self.added_removed_keys,
            trie_cursor_metrics: self.trie_cursor_metrics,
            hashed_cursor_metrics: self.hashed_cursor_metrics,
        }
    }

    /// Set the changed prefixes.
    pub fn with_prefix_set_mut(mut self, prefix_set: PrefixSetMut) -> Self {
        self.prefix_set = prefix_set;
        self
    }

    /// Set the flag indicating whether to include branch node masks in the proof.
    pub const fn with_branch_node_masks(mut self, branch_node_masks: bool) -> Self {
        self.collect_branch_node_masks = branch_node_masks;
        self
    }

    /// Set the trie cursor metrics cache to accumulate metrics into.
    pub const fn with_trie_cursor_metrics(
        mut self,
        metrics: &'a mut TrieCursorMetricsCache,
    ) -> Self {
        self.trie_cursor_metrics = Some(metrics);
        self
    }

    /// Set the hashed cursor metrics cache to accumulate metrics into.
    pub const fn with_hashed_cursor_metrics(
        mut self,
        metrics: &'a mut HashedCursorMetricsCache,
    ) -> Self {
        self.hashed_cursor_metrics = Some(metrics);
        self
    }

    /// Configures the retainer to retain proofs for certain nodes which would otherwise fall
    /// outside the target set, when those nodes might be required to calculate the state root when
    /// keys have been added or removed to the trie.
    ///
    /// If None is given then retention of extra proofs is disabled.
    pub fn with_added_removed_keys<K2>(
        self,
        added_removed_keys: Option<K2>,
    ) -> StorageProof<'a, T, H, K2> {
        StorageProof {
            trie_cursor_factory: self.trie_cursor_factory,
            hashed_cursor_factory: self.hashed_cursor_factory,
            hashed_address: self.hashed_address,
            prefix_set: self.prefix_set,
            collect_branch_node_masks: self.collect_branch_node_masks,
            added_removed_keys,
            trie_cursor_metrics: self.trie_cursor_metrics,
            hashed_cursor_metrics: self.hashed_cursor_metrics,
        }
    }
}

impl<'a, T, H, K> StorageProof<'a, T, H, K>
where
    T: TrieCursorFactory,
    H: HashedCursorFactory,
    K: AsRef<AddedRemovedKeys>,
{
    /// Generate an account proof from intermediate nodes.
    pub fn storage_proof(
        self,
        slot: B256,
    ) -> Result<reth_trie_common::StorageProof, StateProofError> {
        let targets = HashSet::from_iter([keccak256(slot)]);
        Ok(self.storage_multiproof(targets)?.storage_proof(slot)?)
    }

    /// Generate storage proof.
    pub fn storage_multiproof(
        self,
        targets: B256Set,
    ) -> Result<StorageMultiProof, StateProofError> {
        let mut discard_hashed_cursor_metrics = HashedCursorMetricsCache::default();
        let hashed_cursor_metrics =
            self.hashed_cursor_metrics.unwrap_or(&mut discard_hashed_cursor_metrics);

        let hashed_storage_cursor =
            self.hashed_cursor_factory.hashed_storage_cursor(self.hashed_address)?;

        let mut hashed_storage_cursor =
            InstrumentedHashedCursor::new(hashed_storage_cursor, hashed_cursor_metrics);

        // short circuit on empty storage
        if hashed_storage_cursor.is_storage_empty()? {
            return Ok(StorageMultiProof::empty())
        }

        let mut discard_trie_cursor_metrics = TrieCursorMetricsCache::default();
        let trie_cursor_metrics =
            self.trie_cursor_metrics.unwrap_or(&mut discard_trie_cursor_metrics);

        let target_nibbles = targets.into_iter().map(Nibbles::unpack).collect::<Vec<_>>();
        let mut prefix_set = self.prefix_set;
        prefix_set.extend_keys(target_nibbles.iter().copied());

        let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?;

        let trie_cursor = InstrumentedTrieCursor::new(trie_cursor, trie_cursor_metrics);

        let walker = TrieWalker::<_>::storage_trie(trie_cursor, prefix_set.freeze())
            .with_added_removed_keys(self.added_removed_keys.as_ref());

        let retainer = ProofRetainer::from_iter(target_nibbles)
            .with_added_removed_keys(self.added_removed_keys.as_ref());
        let mut hash_builder = HashBuilder::default()
            .with_proof_retainer(retainer)
            .with_updates(self.collect_branch_node_masks);
        let mut storage_node_iter = TrieNodeIter::storage_trie(walker, hashed_storage_cursor);
        while let Some(node) = storage_node_iter.try_next()? {
            match node {
                TrieElement::Branch(node) => {
                    hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
                }
                TrieElement::Leaf(hashed_slot, value) => {
                    hash_builder.add_leaf(
                        Nibbles::unpack(hashed_slot),
                        alloy_rlp::encode_fixed_size(&value).as_ref(),
                    );
                }
            }
        }

        let root = hash_builder.root();
        let subtree = hash_builder.take_proof_nodes();
        let branch_node_masks = if self.collect_branch_node_masks {
            let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default();
            updated_branch_nodes
                .into_iter()
                .map(|(path, node)| {
                    (path, BranchNodeMasks { hash_mask: node.hash_mask, tree_mask: node.tree_mask })
                })
                .collect()
        } else {
            BranchNodeMasksMap::default()
        };

        Ok(StorageMultiProof { root, subtree, branch_node_masks })
    }
}