diff --git a/src/stored/merkle.rs b/src/stored/merkle.rs index 40e380a..4b33d65 100644 --- a/src/stored/merkle.rs +++ b/src/stored/merkle.rs @@ -3,6 +3,7 @@ use core::{cell::RefCell, ops::Deref}; use alloc::{boxed::Box, format, vec::Vec}; use bumpalo::Bump; use ouroboros::self_referencing; +use sha2::{Digest, Sha256}; use crate::{ transaction::nodes::{NodeRef, TrieRoot}, @@ -85,14 +86,15 @@ impl> Store for Snapshot { let idx = node as usize; let leaf_offset = self.branches.len(); let unvisited_offset = leaf_offset + self.leaves.len(); + let mut hasher = Sha256::new(); if let Some(branch) = self.branches.get(idx) { let left = self.calc_subtree_hash(branch.left)?; let right = self.calc_subtree_hash(branch.right)?; - Ok(branch.hash_branch(&left, &right)) + Ok(branch.hash_branch(&mut hasher, &left, &right)) } else if let Some(leaf) = self.leaves.get(idx - leaf_offset) { - Ok(leaf.hash_leaf()) + Ok(leaf.hash_leaf(&mut hasher)) } else if let Some(hash) = self.unvisited_nodes.get(idx - unvisited_offset) { Ok(*hash) } else { diff --git a/src/transaction.rs b/src/transaction.rs index 566dace..cb80f14 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -2,6 +2,7 @@ pub(crate) mod nodes; use alloc::{boxed::Box, format}; use core::mem; +use sha2::{Digest, Sha256}; use crate::{stored, KeyHash, NodeHash}; use crate::{ @@ -67,9 +68,12 @@ impl, V: AsRef<[u8]>> Transaction { ) -> Result<(), TrieError>, on_modified_leaf: &mut impl FnMut(&NodeHash, &Leaf) -> Result<(), TrieError>, ) -> Result, TrieError> { + let mut hasher = Sha256::new(); + let root_hash = match &self.current_root { TrieRoot::Empty => return Ok(TrieRoot::Empty), TrieRoot::Node(node_ref) => Self::calc_root_hash_node( + &mut hasher, &self.data_store, node_ref, on_modified_leaf, @@ -87,6 +91,7 @@ impl, V: AsRef<[u8]>> Transaction { #[inline] fn calc_root_hash_node( + hasher: &mut Sha256, data_store: &S, node_ref: &NodeRef, on_modified_leaf: &mut impl FnMut(&NodeHash, &Leaf) -> Result<(), TrieError>, @@ -101,24 +106,26 @@ impl, V: AsRef<[u8]>> Transaction { match node_ref { NodeRef::ModBranch(branch) => { let left = Self::calc_root_hash_node( + hasher, data_store, &branch.left, on_modified_leaf, on_modified_branch, )?; let right = Self::calc_root_hash_node( + hasher, data_store, &branch.right, on_modified_leaf, on_modified_branch, )?; - let hash = branch.hash_branch(&left, &right); + let hash = branch.hash_branch(hasher, &left, &right); on_modified_branch(&hash, branch, left, right)?; Ok(hash) } NodeRef::ModLeaf(leaf) => { - let hash = leaf.hash_leaf(); + let hash = leaf.hash_leaf(hasher); on_modified_leaf(&hash, leaf)?; Ok(hash) diff --git a/src/transaction/nodes.rs b/src/transaction/nodes.rs index 598528f..b9704f2 100644 --- a/src/transaction/nodes.rs +++ b/src/transaction/nodes.rs @@ -1,7 +1,7 @@ use alloc::{boxed::Box, vec::Vec}; use core::{iter, mem}; -use sha2::{Digest, Sha256}; +use sha2::{digest::FixedOutputReset, Digest, Sha256}; use crate::{stored, KeyHash, NodeHash}; @@ -282,9 +282,8 @@ impl Branch { } #[inline] - pub fn hash_branch(&self, left: &NodeHash, right: &NodeHash) -> NodeHash { - let mut hasher = Sha256::new(); - + pub fn hash_branch(&self, hasher: &mut Sha256, left: &NodeHash, right: &NodeHash) -> NodeHash { + hasher.reset(); hasher.update(left); hasher.update(right); hasher.update(self.mask.bit_idx.to_le_bytes()); @@ -295,7 +294,7 @@ impl Branch { .iter() .for_each(|word| hasher.update(word.to_le_bytes())); - NodeHash::new(hasher.finalize().into()) + NodeHash::new(hasher.finalize_fixed_reset().into()) } } @@ -489,10 +488,10 @@ pub struct Leaf { impl> Leaf { #[inline] - pub fn hash_leaf(&self) -> NodeHash { - let mut hasher = Sha256::new(); + pub fn hash_leaf(&self, hasher: &mut Sha256) -> NodeHash { + hasher.reset(); hasher.update(self.key_hash.to_bytes()); hasher.update(self.value.as_ref()); - NodeHash::new(hasher.finalize().into()) + NodeHash::new(hasher.finalize_fixed_reset().into()) } }