From 1ff76a618b27b66599e76c652617f398a54f0cec Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Tue, 11 Feb 2025 09:40:58 -0500
Subject: [PATCH] feat: add DecodedProofRetainer (#84)

---
 src/proof/decoded_retainer.rs | 48 +++++++++++++++++++++++++++++++++++
 src/proof/mod.rs              |  3 +++
 src/proof/retainer.rs         |  1 -
 3 files changed, 51 insertions(+), 1 deletion(-)
 create mode 100644 src/proof/decoded_retainer.rs

diff --git a/src/proof/decoded_retainer.rs b/src/proof/decoded_retainer.rs
new file mode 100644
index 0000000..794d1b2
--- /dev/null
+++ b/src/proof/decoded_retainer.rs
@@ -0,0 +1,48 @@
+use crate::{proof::DecodedProofNodes, Nibbles};
+use alloy_primitives::Bytes;
+
+use alloc::vec::Vec;
+
+/// Proof retainer is used to store proofs during merkle trie construction.
+/// It is intended to be used within the [`HashBuilder`](crate::HashBuilder).
+#[derive(Default, Clone, Debug)]
+pub struct DecodedProofRetainer {
+    /// The nibbles of the target trie keys to retain proofs for.
+    targets: Vec<Nibbles>,
+    /// The map retained trie node keys to RLP serialized trie nodes.
+    proof_nodes: DecodedProofNodes,
+}
+
+impl FromIterator<Nibbles> for DecodedProofRetainer {
+    fn from_iter<T: IntoIterator<Item = Nibbles>>(iter: T) -> Self {
+        Self::new(FromIterator::from_iter(iter))
+    }
+}
+
+impl DecodedProofRetainer {
+    /// Create new retainer with target nibbles.
+    pub fn new(targets: Vec<Nibbles>) -> Self {
+        Self { targets, proof_nodes: Default::default() }
+    }
+
+    /// Returns `true` if the given prefix matches the retainer target.
+    pub fn matches(&self, prefix: &Nibbles) -> bool {
+        self.targets.iter().any(|target| target.starts_with(prefix))
+    }
+
+    /// Returns all collected proofs.
+    pub fn into_proof_nodes(self) -> DecodedProofNodes {
+        self.proof_nodes
+    }
+
+    /// Retain the proof if the key matches any of the targets.
+    ///
+    /// Returns an error if the proof could not be decoded from the given proof bytes.
+    pub fn retain(&mut self, prefix: &Nibbles, proof: &[u8]) -> Result<(), alloy_rlp::Error> {
+        if prefix.is_empty() || self.matches(prefix) {
+            self.proof_nodes.insert_encoded(prefix.clone(), Bytes::from(proof.to_vec()))?;
+        }
+
+        Ok(())
+    }
+}
diff --git a/src/proof/mod.rs b/src/proof/mod.rs
index 13c1dde..642dda6 100644
--- a/src/proof/mod.rs
+++ b/src/proof/mod.rs
@@ -12,6 +12,9 @@ pub use error::ProofVerificationError;
 mod decoded_proof_nodes;
 pub use decoded_proof_nodes::DecodedProofNodes;
 
+mod decoded_retainer;
+pub use decoded_retainer::DecodedProofRetainer;
+
 mod proof_nodes;
 pub use proof_nodes::ProofNodes;
 
diff --git a/src/proof/retainer.rs b/src/proof/retainer.rs
index 0c5d74f..e8d0cc6 100644
--- a/src/proof/retainer.rs
+++ b/src/proof/retainer.rs
@@ -1,7 +1,6 @@
 use crate::{proof::ProofNodes, Nibbles};
 use alloy_primitives::Bytes;
 
-#[allow(unused_imports)]
 use alloc::vec::Vec;
 
 /// Proof retainer is used to store proofs during merkle trie construction.