diff --git a/crates/common/trie/node.rs b/crates/common/trie/node.rs index 8818c54f8a..183a6bd463 100644 --- a/crates/common/trie/node.rs +++ b/crates/common/trie/node.rs @@ -4,6 +4,7 @@ mod leaf; use std::{ array, + collections::btree_map::Range, sync::{Arc, OnceLock}, }; @@ -51,23 +52,30 @@ impl NodeRef { } } - pub fn commit(&mut self, acc: &mut Vec<(NodeHash, Vec)>) -> NodeHash { + pub fn commit( + &mut self, + buffer: &mut Vec, + acc: &mut Vec<(NodeHash, std::ops::Range)>, + ) -> NodeHash { match *self { NodeRef::Node(ref mut node, ref mut hash) => { match Arc::make_mut(node) { Node::Branch(node) => { for node in &mut node.choices { - node.commit(acc); + node.commit(buffer, acc); } } Node::Extension(node) => { - node.child.commit(acc); + node.child.commit(buffer, acc); } Node::Leaf(_) => {} } let hash = hash.get_or_init(|| node.compute_hash()); - acc.push((*hash, node.encode_to_vec())); + let first = buffer.len(); + node.encode(buffer); + let last = buffer.len(); + acc.push((*hash, (first..last))); let hash = *hash; *self = hash.into(); diff --git a/crates/common/trie/trie.rs b/crates/common/trie/trie.rs index 7011a3d144..27e3e28f3d 100644 --- a/crates/common/trie/trie.rs +++ b/crates/common/trie/trie.rs @@ -176,8 +176,13 @@ impl Trie { pub fn commit(&mut self) -> Result<(), TrieError> { if self.root.is_valid() { let mut acc = Vec::new(); - self.root.commit(&mut acc); - self.db.put_batch(acc)?; // we'll try to avoid calling this for every commit + let mut buffer = Vec::new(); + self.root.commit(&mut buffer, &mut acc); + let updates: Vec<(NodeHash, Vec)> = acc + .iter() + .map(|(hash, range)| (*hash, buffer[range.clone()].to_vec())) + .collect(); + self.db.put_batch(updates)?; // we'll try to avoid calling this for every commit } Ok(()) @@ -187,11 +192,16 @@ impl Trie { /// Nodes are given with their hash pre-calculated. pub fn commit_without_storing(&mut self) -> Vec { let mut acc = Vec::new(); + let mut buffer = Vec::new(); if self.root.is_valid() { - self.root.commit(&mut acc); + self.root.commit(&mut buffer, &mut acc); } - acc + let updates: Vec<(NodeHash, Vec)> = acc + .iter() + .map(|(hash, range)| (*hash, buffer[range.clone()].to_vec())) + .collect(); + updates } /// Obtain a merkle proof for the given path.