alloy_trie/proof/
retainer.rs

1use crate::{proof::ProofNodes, Nibbles};
2use alloy_primitives::Bytes;
3
4use alloc::vec::Vec;
5
6/// Proof retainer is used to store proofs during merkle trie construction.
7/// It is intended to be used within the [`HashBuilder`](crate::HashBuilder).
8#[derive(Default, Clone, Debug)]
9pub struct ProofRetainer {
10    /// The nibbles of the target trie keys to retain proofs for.
11    targets: Vec<Nibbles>,
12    /// The map retained trie node keys to RLP serialized trie nodes.
13    proof_nodes: ProofNodes,
14}
15
16impl FromIterator<Nibbles> for ProofRetainer {
17    fn from_iter<T: IntoIterator<Item = Nibbles>>(iter: T) -> Self {
18        Self::new(FromIterator::from_iter(iter))
19    }
20}
21
22impl ProofRetainer {
23    /// Create new retainer with target nibbles.
24    pub fn new(targets: Vec<Nibbles>) -> Self {
25        Self { targets, proof_nodes: Default::default() }
26    }
27
28    /// Returns `true` if the given prefix matches the retainer target.
29    pub fn matches(&self, prefix: &Nibbles) -> bool {
30        self.targets.iter().any(|target| target.starts_with(prefix))
31    }
32
33    /// Returns all collected proofs.
34    pub fn into_proof_nodes(self) -> ProofNodes {
35        self.proof_nodes
36    }
37
38    /// Retain the proof if the key matches any of the targets.
39    pub fn retain(&mut self, prefix: &Nibbles, proof: &[u8]) {
40        if prefix.is_empty() || self.matches(prefix) {
41            self.proof_nodes.insert(prefix.clone(), Bytes::from(proof.to_vec()));
42        }
43    }
44}