perf: change proof internal repr to HashMap

alloy-rs · Sep 26, 2024 · c1f6e80 · c1f6e80
1 parent b286769
commit c1f6e80
Show file tree

Hide file tree

Showing 5 changed files with 148 additions and 32 deletions.
diff --git 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<Nibbles, Bytes> {
-        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
@@ -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
@@ -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<Nibbles, Bytes>);
+
+impl Deref for ProofNodes {
+    type Target = HashMap<Nibbles, Bytes>;
+
+    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<Item = (Nibbles, Bytes)>) -> 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<Item = (&'a Nibbles, &'a Bytes)> {
+        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<Bytes> {
+        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<Nibbles, Bytes> {
+        self.0
+    }
+}
diff --git a/src/proof/retainer.rs b/src/proof/retainer.rs
@@ -1,18 +1,17 @@
-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).
 #[derive(Default, Debug)]
 pub struct ProofRetainer {
     /// The nibbles of the target trie keys to retain proofs for.
     targets: Vec<Nibbles>,
-    /// The map of retained proofs (RLP serialized trie nodes)
-    /// with their corresponding key in the trie.
-    proofs: BTreeMap<Nibbles, Bytes>,
+    /// The map retained trie node keys to RLP serialized trie nodes.
+    proof_nodes: ProofNodes,
 }
 
 impl core::iter::FromIterator<Nibbles> for ProofRetainer {
@@ -24,7 +23,7 @@ impl core::iter::FromIterator<Nibbles> for ProofRetainer {
 impl ProofRetainer {
     /// Create new retainer with target nibbles.
     pub fn new(targets: Vec<Nibbles>) -> 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<Nibbles, Bytes> {
-        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
@@ -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,15 +291,20 @@ 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()),
             (Nibbles::from_vec(vec![0x3, 0xc, 0x1]), Bytes::from_str("f8518080808080a020dc5b33292bfad9013bf123f7faf1efcc5c8e00c894177fc0bfb447daef522f808080a020dc5b33292bfad9013bf123f7faf1efcc5c8e00c894177fc0bfb447daef522f80808080808080").unwrap()),
             (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(())
         );
     }