Skip to content

Commit

Permalink
feat: add DecodedProofNodes struct (#81)
Browse files Browse the repository at this point in the history
  • Loading branch information
Rjected authored Feb 7, 2025
1 parent f7f6feb commit 77ab71a
Show file tree
Hide file tree
Showing 3 changed files with 106 additions and 1 deletion.
103 changes: 103 additions & 0 deletions src/proof/decoded_proof_nodes.rs
Original file line number Diff line number Diff line change
@@ -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);
}
}
3 changes: 3 additions & 0 deletions src/proof/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
1 change: 0 additions & 1 deletion src/proof/proof_nodes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit 77ab71a

Please sign in to comment.