From 77ab71aaeb41f3f607cf655c56032750b8816548 Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Fri, 7 Feb 2025 17:37:17 -0500
Subject: [PATCH] feat: add DecodedProofNodes struct (#81)

---
 src/proof/decoded_proof_nodes.rs | 103 +++++++++++++++++++++++++++++++
 src/proof/mod.rs                 |   3 +
 src/proof/proof_nodes.rs         |   1 -
 3 files changed, 106 insertions(+), 1 deletion(-)
 create mode 100644 src/proof/decoded_proof_nodes.rs

diff --git a/src/proof/decoded_proof_nodes.rs b/src/proof/decoded_proof_nodes.rs
new file mode 100644
index 0000000..5c27699
--- /dev/null
+++ b/src/proof/decoded_proof_nodes.rs
@@ -0,0 +1,103 @@
+use crate::{nodes::TrieNode, proof::ProofNodes, HashMap, Nibbles};
+use alloy_primitives::Bytes;
+use alloy_rlp::Decodable;
+use core::ops::Deref;
+
+use alloc::vec::Vec;
+
+/// A wrapper struct for trie node key to RLP encoded trie node.
+#[derive(PartialEq, Eq, Clone, Default, Debug)]
+pub struct DecodedProofNodes(HashMap<Nibbles, TrieNode>);
+
+impl Deref for DecodedProofNodes {
+    type Target = HashMap<Nibbles, TrieNode>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl FromIterator<(Nibbles, TrieNode)> for DecodedProofNodes {
+    fn from_iter<T: IntoIterator<Item = (Nibbles, TrieNode)>>(iter: T) -> Self {
+        Self(HashMap::from_iter(iter))
+    }
+}
+
+impl Extend<(Nibbles, TrieNode)> for DecodedProofNodes {
+    fn extend<T: IntoIterator<Item = (Nibbles, TrieNode)>>(&mut self, iter: T) {
+        self.0.extend(iter);
+    }
+}
+
+impl TryFrom<ProofNodes> for DecodedProofNodes {
+    type Error = alloy_rlp::Error;
+
+    fn try_from(proof_nodes: ProofNodes) -> Result<Self, Self::Error> {
+        let mut decoded_proof_nodes =
+            HashMap::with_capacity_and_hasher(proof_nodes.len(), Default::default());
+        for (key, node) in proof_nodes.into_inner() {
+            decoded_proof_nodes.insert(key, TrieNode::decode(&mut &node[..])?);
+        }
+        Ok(Self(decoded_proof_nodes))
+    }
+}
+
+impl DecodedProofNodes {
+    /// Return iterator over proof nodes that match the target.
+    pub fn matching_nodes_iter<'a>(
+        &'a self,
+        target: &'a Nibbles,
+    ) -> impl Iterator<Item = (&'a Nibbles, &'a TrieNode)> {
+        self.0.iter().filter(|(key, _)| target.starts_with(key))
+    }
+
+    /// Return the vec of proof nodes that match the target.
+    pub fn matching_nodes(&self, target: &Nibbles) -> Vec<(Nibbles, TrieNode)> {
+        self.matching_nodes_iter(target).map(|(key, node)| (key.clone(), node.clone())).collect()
+    }
+
+    /// Return the sorted vec of proof nodes that match the target.
+    pub fn matching_nodes_sorted(&self, target: &Nibbles) -> Vec<(Nibbles, TrieNode)> {
+        let mut nodes = self.matching_nodes(target);
+        nodes.sort_unstable_by(|a, b| a.0.cmp(&b.0));
+        nodes
+    }
+
+    /// Insert the trie node at key.
+    pub fn insert(&mut self, key: Nibbles, node: TrieNode) -> Option<TrieNode> {
+        self.0.insert(key, node)
+    }
+
+    /// Insert the RLP encoded trie nodoe at key
+    pub fn insert_encoded(
+        &mut self,
+        key: Nibbles,
+        node: Bytes,
+    ) -> Result<Option<TrieNode>, alloy_rlp::Error> {
+        Ok(self.0.insert(key, TrieNode::decode(&mut &node[..])?))
+    }
+
+    /// Return the sorted vec of all proof nodes.
+    pub fn nodes_sorted(&self) -> Vec<(Nibbles, TrieNode)> {
+        let mut nodes = Vec::from_iter(self.0.iter().map(|(k, v)| (k.clone(), v.clone())));
+        nodes.sort_unstable_by(|a, b| a.0.cmp(&b.0));
+        nodes
+    }
+
+    /// Convert into sorted vec of all proof nodes.
+    pub fn into_nodes_sorted(self) -> Vec<(Nibbles, TrieNode)> {
+        let mut nodes = Vec::from_iter(self.0);
+        nodes.sort_unstable_by(|a, b| a.0.cmp(&b.0));
+        nodes
+    }
+
+    /// Convert wrapper struct into inner map.
+    pub fn into_inner(self) -> HashMap<Nibbles, TrieNode> {
+        self.0
+    }
+
+    /// Extends with the elements of another `DecodedProofNodes`.
+    pub fn extend_from(&mut self, other: Self) {
+        self.extend(other.0);
+    }
+}
diff --git a/src/proof/mod.rs b/src/proof/mod.rs
index 4157ace..13c1dde 100644
--- a/src/proof/mod.rs
+++ b/src/proof/mod.rs
@@ -9,6 +9,9 @@ pub use verify::verify_proof;
 mod error;
 pub use error::ProofVerificationError;
 
+mod decoded_proof_nodes;
+pub use decoded_proof_nodes::DecodedProofNodes;
+
 mod proof_nodes;
 pub use proof_nodes::ProofNodes;
 
diff --git a/src/proof/proof_nodes.rs b/src/proof/proof_nodes.rs
index d1f5f4b..0bc1043 100644
--- a/src/proof/proof_nodes.rs
+++ b/src/proof/proof_nodes.rs
@@ -2,7 +2,6 @@ use crate::{HashMap, Nibbles};
 use alloy_primitives::Bytes;
 use core::ops::Deref;
 
-#[allow(unused_imports)]
 use alloc::vec::Vec;
 
 /// A wrapper struct for trie node key to RLP encoded trie node.