From c1f6e8038b7f9e657622f854dc00980d009bb58a Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 26 Sep 2024 17:46:21 +0200 Subject: [PATCH 1/5] perf: change proof internal repr to `HashMap` --- src/hash_builder/mod.rs | 11 +++--- src/proof/mod.rs | 3 ++ src/proof/proof_nodes.rs | 66 ++++++++++++++++++++++++++++++++ src/proof/retainer.rs | 17 ++++---- src/proof/verify.rs | 83 +++++++++++++++++++++++++++++++--------- 5 files changed, 148 insertions(+), 32 deletions(-) create mode 100644 src/proof/proof_nodes.rs diff --git a/src/hash_builder/mod.rs b/src/hash_builder/mod.rs index ad57cd6..cb2427f 100644 --- a/src/hash_builder/mod.rs +++ b/src/hash_builder/mod.rs @@ -5,14 +5,14 @@ use super::{ proof::ProofRetainer, BranchNodeCompact, Nibbles, TrieMask, EMPTY_ROOT_HASH, }; -use crate::HashMap; -use alloy_primitives::{hex, keccak256, Bytes, B256}; +use crate::{proof::ProofNodes, HashMap}; +use alloy_primitives::{hex, keccak256, B256}; use alloy_rlp::EMPTY_STRING_CODE; use core::cmp; use tracing::trace; #[allow(unused_imports)] -use alloc::{collections::BTreeMap, vec::Vec}; +use alloc::vec::Vec; mod value; pub use value::HashBuilderValue; @@ -91,8 +91,8 @@ impl HashBuilder { } /// Take and return the proofs retained. - pub fn take_proofs(&mut self) -> BTreeMap { - self.proof_retainer.take().map(ProofRetainer::into_proofs).unwrap_or_default() + pub fn take_proofs(&mut self) -> ProofNodes { + self.proof_retainer.take().map(ProofRetainer::into_proof_nodes).unwrap_or_default() } /// The number of total updates accrued. @@ -420,6 +420,7 @@ impl HashBuilder { mod tests { use super::*; use crate::{nodes::LeafNode, triehash_trie_root}; + use alloc::collections::BTreeMap; use alloy_primitives::{b256, hex, U256}; use alloy_rlp::Encodable; diff --git a/src/proof/mod.rs b/src/proof/mod.rs index c930664..4157ace 100644 --- a/src/proof/mod.rs +++ b/src/proof/mod.rs @@ -9,5 +9,8 @@ pub use verify::verify_proof; mod error; pub use error::ProofVerificationError; +mod proof_nodes; +pub use proof_nodes::ProofNodes; + mod retainer; pub use retainer::ProofRetainer; diff --git a/src/proof/proof_nodes.rs b/src/proof/proof_nodes.rs new file mode 100644 index 0000000..a5164a1 --- /dev/null +++ b/src/proof/proof_nodes.rs @@ -0,0 +1,66 @@ +use crate::{HashMap, Nibbles}; +use alloy_primitives::Bytes; +use core::ops::Deref; + +/// A wrapper struct for trie node key to RLP encoded trie node. +#[derive(PartialEq, Eq, Clone, Default, Debug)] +pub struct ProofNodes(HashMap); + +impl Deref for ProofNodes { + type Target = HashMap; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl ProofNodes { + /// Construct struct from iterator over keys and proof nodes. + pub fn from_iter(nodes: impl IntoIterator) -> Self { + Self(HashMap::from_iter(nodes)) + } + + /// Return iterator over proof nodes that match the target. + pub fn matching_nodes_iter<'a>( + &'a self, + target: &'a Nibbles, + ) -> impl Iterator { + 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, Bytes)> { + 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, Bytes)> { + let mut nodes = self.matching_nodes(target); + nodes.sort_unstable_by(|a, b| a.0.cmp(&b.0)); + nodes + } + + /// Insert the RLP encoded trie node at key. + pub fn insert(&mut self, key: Nibbles, node: Bytes) -> Option { + self.0.insert(key, node) + } + + /// Return the sorted vec of all proof nodes. + pub fn nodes_sorted(&self) -> Vec<(Nibbles, Bytes)> { + 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, Bytes)> { + 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 { + self.0 + } +} diff --git a/src/proof/retainer.rs b/src/proof/retainer.rs index 1853628..b5dade5 100644 --- a/src/proof/retainer.rs +++ b/src/proof/retainer.rs @@ -1,8 +1,8 @@ -use crate::Nibbles; +use crate::{proof::ProofNodes, Nibbles}; use alloy_primitives::Bytes; #[allow(unused_imports)] -use alloc::{collections::BTreeMap, vec::Vec}; +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). @@ -10,9 +10,8 @@ use alloc::{collections::BTreeMap, vec::Vec}; pub struct ProofRetainer { /// The nibbles of the target trie keys to retain proofs for. targets: Vec, - /// The map of retained proofs (RLP serialized trie nodes) - /// with their corresponding key in the trie. - proofs: BTreeMap, + /// The map retained trie node keys to RLP serialized trie nodes. + proof_nodes: ProofNodes, } impl core::iter::FromIterator for ProofRetainer { @@ -24,7 +23,7 @@ impl core::iter::FromIterator for ProofRetainer { impl ProofRetainer { /// Create new retainer with target nibbles. pub fn new(targets: Vec) -> Self { - Self { targets, proofs: Default::default() } + Self { targets, proof_nodes: Default::default() } } /// Returns `true` if the given prefix matches the retainer target. @@ -33,14 +32,14 @@ impl ProofRetainer { } /// Returns all collected proofs. - pub fn into_proofs(self) -> BTreeMap { - self.proofs + pub fn into_proof_nodes(self) -> ProofNodes { + self.proof_nodes } /// Retain the proof if the key matches any of the targets. pub fn retain(&mut self, prefix: &Nibbles, proof: &[u8]) { if prefix.is_empty() || self.matches(prefix) { - self.proofs.insert(prefix.clone(), Bytes::from(proof.to_vec())); + self.proof_nodes.insert(prefix.clone(), Bytes::from(proof.to_vec())); } } } diff --git a/src/proof/verify.rs b/src/proof/verify.rs index ac8fa58..cacdc8e 100644 --- a/src/proof/verify.rs +++ b/src/proof/verify.rs @@ -150,10 +150,9 @@ mod tests { use super::*; use crate::{ nodes::{BranchNode, ExtensionNode, LeafNode}, - proof::ProofRetainer, + proof::{ProofNodes, ProofRetainer}, triehash_trie_root, HashBuilder, TrieMask, }; - use alloc::collections::BTreeMap; use alloy_primitives::hex; use alloy_rlp::{Encodable, EMPTY_STRING_CODE}; use core::str::FromStr; @@ -164,8 +163,19 @@ mod tests { let mut hash_builder = HashBuilder::default().with_proof_retainer(ProofRetainer::default()); let root = hash_builder.root(); let proof = hash_builder.take_proofs(); - assert_eq!(proof, BTreeMap::from([(Nibbles::default(), Bytes::from([EMPTY_STRING_CODE]))])); - assert_eq!(verify_proof(root, key.clone(), None, proof.values()), Ok(())); + assert_eq!( + proof, + ProofNodes::from_iter([(Nibbles::default(), Bytes::from([EMPTY_STRING_CODE]))]) + ); + assert_eq!( + verify_proof( + root, + key.clone(), + None, + proof.into_nodes_sorted().iter().map(|(_, node)| node) + ), + Ok(()) + ); let mut dummy_proof = vec![]; BranchNode::default().encode(&mut dummy_proof); @@ -191,8 +201,16 @@ mod tests { let root = hash_builder.root(); assert_eq!(root, triehash_trie_root([(target.pack(), target.pack())])); - let proof = hash_builder.take_proofs(); - assert_eq!(verify_proof(root, target, Some(target_value.to_vec()), proof.values()), Ok(())); + let proof = hash_builder.take_proofs().into_nodes_sorted(); + assert_eq!( + verify_proof( + root, + target, + Some(target_value.to_vec()), + proof.iter().map(|(_, node)| node) + ), + Ok(()) + ); } #[test] @@ -212,8 +230,8 @@ mod tests { triehash_trie_root(range.map(|b| (B256::with_last_byte(b), B256::with_last_byte(b)))) ); - let proof = hash_builder.take_proofs(); - assert_eq!(verify_proof(root, target, None, proof.values()), Ok(())); + let proof = hash_builder.take_proofs().into_nodes_sorted(); + assert_eq!(verify_proof(root, target, None, proof.iter().map(|(_, node)| node)), Ok(())); } #[test] @@ -241,12 +259,20 @@ mod tests { triehash_trie_root(existing_keys.map(|key| (B256::from_slice(&key), value))) ); let proof = hash_builder.take_proofs(); - assert_eq!(proof, BTreeMap::from([ + assert_eq!(proof, ProofNodes::from_iter([ (Nibbles::default(), Bytes::from_str("f851a0c530c099d779362b6bd0be05039b51ccd0a8ed39e0b2abacab8fe0e3441251878080a07d4ee4f073ae7ce32a6cbcdb015eb73dd2616f33ed2e9fb6ba51c1f9ad5b697b80808080808080808080808080").unwrap()), (Nibbles::from_vec(vec![0x3]), Bytes::from_str("f85180808080808080808080a057fcbd3f97b1093cd39d0f58dafd5058e2d9f79a419e88c2498ff3952cb11a8480a07520d69a83a2bdad373a68b2c9c8c0e1e1c99b6ec80b4b933084da76d644081980808080").unwrap()), (Nibbles::from_vec(vec![0x3, 0xc]), Bytes::from_str("f842a02015000000000000000000000000000000000000000000000000000000000000a00000000000000000000000000000000000000000000000000000000000000001").unwrap()) ])); - assert_eq!(verify_proof(root, target.clone(), None, proof.values()), Ok(())); + assert_eq!( + verify_proof( + root, + target.clone(), + None, + proof.into_nodes_sorted().iter().map(|(_, node)| node) + ), + Ok(()) + ); let retainer = ProofRetainer::from_iter([target.clone()]); let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer); @@ -265,7 +291,7 @@ mod tests { ) ); let proof = hash_builder.take_proofs(); - assert_eq!(proof, BTreeMap::from([ + assert_eq!(proof, ProofNodes::from_iter([ (Nibbles::default(), Bytes::from_str("f851a0c530c099d779362b6bd0be05039b51ccd0a8ed39e0b2abacab8fe0e3441251878080a0abd80d939392f6d222f8becc15f8c6f0dbbc6833dd7e54bfbbee0c589b7fd40380808080808080808080808080").unwrap()), (Nibbles::from_vec(vec![0x3]), Bytes::from_str("f85180808080808080808080a057fcbd3f97b1093cd39d0f58dafd5058e2d9f79a419e88c2498ff3952cb11a8480a09e7b3788773773f15e26ad07b72a2c25a6374bce256d9aab6cea48fbc77d698180808080").unwrap()), (Nibbles::from_vec(vec![0x3, 0xc]), Bytes::from_str("e211a0338ac0a453edb0e40a23a70aee59e02a6c11597c34d79a5ba94da8eb20dd4d52").unwrap()), @@ -273,7 +299,12 @@ mod tests { (Nibbles::from_vec(vec![0x3, 0xc, 0x1, 0x9]), Bytes::from_str("f8419f20000000000000000000000000000000000000000000000000000000000000a00000000000000000000000000000000000000000000000000000000000000001").unwrap()), ])); assert_eq!( - verify_proof(root, target.clone(), Some(value.to_vec()), proof.values()), + verify_proof( + root, + target.clone(), + Some(value.to_vec()), + proof.into_nodes_sorted().iter().map(|(_, node)| node) + ), Ok(()) ); } @@ -296,8 +327,16 @@ mod tests { triehash_trie_root(range.map(|b| (B256::with_last_byte(b), B256::with_last_byte(b)))) ); - let proof = hash_builder.take_proofs(); - assert_eq!(verify_proof(root, target, Some(target_value.to_vec()), proof.values()), Ok(())); + let proof = hash_builder.take_proofs().into_nodes_sorted(); + assert_eq!( + verify_proof( + root, + target, + Some(target_value.to_vec()), + proof.iter().map(|(_, node)| node) + ), + Ok(()) + ); } #[test] @@ -322,15 +361,23 @@ mod tests { let proof = hash_builder.take_proofs(); - let proof1 = proof.iter().filter_map(|(k, v)| target1.starts_with(k).then_some(v)); assert_eq!( - verify_proof(root, target1.clone(), Some(target1_value.to_vec()), proof1), + verify_proof( + root, + target1.clone(), + Some(target1_value.to_vec()), + proof.matching_nodes_sorted(&target1).iter().map(|(_, node)| node) + ), Ok(()) ); - let proof2 = proof.iter().filter_map(|(k, v)| target2.starts_with(k).then_some(v)); assert_eq!( - verify_proof(root, target2.clone(), Some(target2_value.to_vec()), proof2), + verify_proof( + root, + target2.clone(), + Some(target2_value.to_vec()), + proof.matching_nodes_sorted(&target2).iter().map(|(_, node)| node) + ), Ok(()) ); } From 922eb369335b48797c9932a40e9dd1c8fa340dfa Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 26 Sep 2024 17:48:33 +0200 Subject: [PATCH 2/5] alloc::vec --- src/proof/proof_nodes.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/proof/proof_nodes.rs b/src/proof/proof_nodes.rs index a5164a1..bbf822f 100644 --- a/src/proof/proof_nodes.rs +++ b/src/proof/proof_nodes.rs @@ -2,6 +2,9 @@ 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. #[derive(PartialEq, Eq, Clone, Default, Debug)] pub struct ProofNodes(HashMap); From 31d3f50a5745e7d431fb6ce8a4618f3ebe39acc1 Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 26 Sep 2024 17:49:27 +0200 Subject: [PATCH 3/5] take_proofs -> take_proof_nodes --- src/hash_builder/mod.rs | 4 ++-- src/proof/verify.rs | 14 +++++++------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/hash_builder/mod.rs b/src/hash_builder/mod.rs index cb2427f..484ee5d 100644 --- a/src/hash_builder/mod.rs +++ b/src/hash_builder/mod.rs @@ -90,8 +90,8 @@ impl HashBuilder { (self, updates.unwrap_or_default()) } - /// Take and return the proofs retained. - pub fn take_proofs(&mut self) -> ProofNodes { + /// Take and return retained proof nodes. + pub fn take_proof_nodes(&mut self) -> ProofNodes { self.proof_retainer.take().map(ProofRetainer::into_proof_nodes).unwrap_or_default() } diff --git a/src/proof/verify.rs b/src/proof/verify.rs index cacdc8e..ca055de 100644 --- a/src/proof/verify.rs +++ b/src/proof/verify.rs @@ -162,7 +162,7 @@ mod tests { let key = Nibbles::unpack(B256::repeat_byte(42)); let mut hash_builder = HashBuilder::default().with_proof_retainer(ProofRetainer::default()); let root = hash_builder.root(); - let proof = hash_builder.take_proofs(); + let proof = hash_builder.take_proof_nodes(); assert_eq!( proof, ProofNodes::from_iter([(Nibbles::default(), Bytes::from([EMPTY_STRING_CODE]))]) @@ -201,7 +201,7 @@ mod tests { let root = hash_builder.root(); assert_eq!(root, triehash_trie_root([(target.pack(), target.pack())])); - let proof = hash_builder.take_proofs().into_nodes_sorted(); + let proof = hash_builder.take_proof_nodes().into_nodes_sorted(); assert_eq!( verify_proof( root, @@ -230,7 +230,7 @@ mod tests { triehash_trie_root(range.map(|b| (B256::with_last_byte(b), B256::with_last_byte(b)))) ); - let proof = hash_builder.take_proofs().into_nodes_sorted(); + let proof = hash_builder.take_proof_nodes().into_nodes_sorted(); assert_eq!(verify_proof(root, target, None, proof.iter().map(|(_, node)| node)), Ok(())); } @@ -258,7 +258,7 @@ mod tests { root, triehash_trie_root(existing_keys.map(|key| (B256::from_slice(&key), value))) ); - let proof = hash_builder.take_proofs(); + let proof = hash_builder.take_proof_nodes(); assert_eq!(proof, ProofNodes::from_iter([ (Nibbles::default(), Bytes::from_str("f851a0c530c099d779362b6bd0be05039b51ccd0a8ed39e0b2abacab8fe0e3441251878080a07d4ee4f073ae7ce32a6cbcdb015eb73dd2616f33ed2e9fb6ba51c1f9ad5b697b80808080808080808080808080").unwrap()), (Nibbles::from_vec(vec![0x3]), Bytes::from_str("f85180808080808080808080a057fcbd3f97b1093cd39d0f58dafd5058e2d9f79a419e88c2498ff3952cb11a8480a07520d69a83a2bdad373a68b2c9c8c0e1e1c99b6ec80b4b933084da76d644081980808080").unwrap()), @@ -290,7 +290,7 @@ mod tests { .chain([(B256::from_slice(&target.pack()), value)]) ) ); - let proof = hash_builder.take_proofs(); + let proof = hash_builder.take_proof_nodes(); assert_eq!(proof, ProofNodes::from_iter([ (Nibbles::default(), Bytes::from_str("f851a0c530c099d779362b6bd0be05039b51ccd0a8ed39e0b2abacab8fe0e3441251878080a0abd80d939392f6d222f8becc15f8c6f0dbbc6833dd7e54bfbbee0c589b7fd40380808080808080808080808080").unwrap()), (Nibbles::from_vec(vec![0x3]), Bytes::from_str("f85180808080808080808080a057fcbd3f97b1093cd39d0f58dafd5058e2d9f79a419e88c2498ff3952cb11a8480a09e7b3788773773f15e26ad07b72a2c25a6374bce256d9aab6cea48fbc77d698180808080").unwrap()), @@ -327,7 +327,7 @@ mod tests { triehash_trie_root(range.map(|b| (B256::with_last_byte(b), B256::with_last_byte(b)))) ); - let proof = hash_builder.take_proofs().into_nodes_sorted(); + let proof = hash_builder.take_proof_nodes().into_nodes_sorted(); assert_eq!( verify_proof( root, @@ -359,7 +359,7 @@ mod tests { triehash_trie_root(range.map(|b| (B256::repeat_byte(b), B256::repeat_byte(b)))) ); - let proof = hash_builder.take_proofs(); + let proof = hash_builder.take_proof_nodes(); assert_eq!( verify_proof( From 810586196ea4ecaf071d9f38c5b145e03145737b Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 26 Sep 2024 17:51:06 +0200 Subject: [PATCH 4/5] fix arbitrary proof verification test --- src/proof/verify.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/proof/verify.rs b/src/proof/verify.rs index ca055de..0de830d 100644 --- a/src/proof/verify.rs +++ b/src/proof/verify.rs @@ -588,11 +588,10 @@ mod tests { let root = hash_builder.root(); assert_eq!(root, triehash_trie_root(&hashed)); - let proofs = hash_builder.take_proofs(); + let proofs = hash_builder.take_proof_nodes(); for (key, value) in hashed { let nibbles = Nibbles::unpack(key); - let proof = proofs.iter().filter_map(|(k, v)| nibbles.starts_with(k).then_some(v)); - assert_eq!(verify_proof(root, nibbles.clone(), Some(value), proof), Ok(())); + assert_eq!(verify_proof(root, nibbles.clone(), Some(value), proofs.matching_nodes_sorted(&nibbles).iter().map(|(_, node)| node)), Ok(())); } }); } From 637d0e1a50e89b06c62472e048e633a4673d5080 Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 26 Sep 2024 17:53:00 +0200 Subject: [PATCH 5/5] change from_iter into trait impl --- src/proof/proof_nodes.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/proof/proof_nodes.rs b/src/proof/proof_nodes.rs index bbf822f..81a306c 100644 --- a/src/proof/proof_nodes.rs +++ b/src/proof/proof_nodes.rs @@ -17,12 +17,13 @@ impl Deref for ProofNodes { } } -impl ProofNodes { - /// Construct struct from iterator over keys and proof nodes. - pub fn from_iter(nodes: impl IntoIterator) -> Self { - Self(HashMap::from_iter(nodes)) +impl FromIterator<(Nibbles, Bytes)> for ProofNodes { + fn from_iter>(iter: T) -> Self { + Self(HashMap::from_iter(iter)) } +} +impl ProofNodes { /// Return iterator over proof nodes that match the target. pub fn matching_nodes_iter<'a>( &'a self,