1use super::{
4 nodes::{BranchNodeRef, ExtensionNodeRef, LeafNodeRef},
5 proof::ProofRetainer,
6 BranchNodeCompact, Nibbles, TrieMask, EMPTY_ROOT_HASH,
7};
8use crate::{nodes::RlpNode, proof::ProofNodes, HashMap};
9use alloc::vec::Vec;
10use alloy_primitives::{keccak256, B256};
11use alloy_rlp::EMPTY_STRING_CODE;
12use core::cmp;
13use tracing::trace;
14
15mod value;
16pub use value::{HashBuilderValue, HashBuilderValueRef};
17
18#[derive(Debug, Clone, Default)]
43#[allow(missing_docs)]
44pub struct HashBuilder {
45 pub key: Nibbles,
46 pub value: HashBuilderValue,
47 pub stack: Vec<RlpNode>,
48
49 pub state_masks: Vec<TrieMask>,
50 pub tree_masks: Vec<TrieMask>,
51 pub hash_masks: Vec<TrieMask>,
52
53 pub stored_in_database: bool,
54
55 pub updated_branch_nodes: Option<HashMap<Nibbles, BranchNodeCompact>>,
56 pub proof_retainer: Option<ProofRetainer>,
57
58 pub rlp_buf: Vec<u8>,
59}
60
61impl HashBuilder {
62 pub fn with_updates(mut self, retain_updates: bool) -> Self {
66 self.set_updates(retain_updates);
67 self
68 }
69
70 pub fn with_proof_retainer(mut self, retainer: ProofRetainer) -> Self {
72 self.proof_retainer = Some(retainer);
73 self
74 }
75
76 pub fn set_updates(&mut self, retain_updates: bool) {
80 if retain_updates {
81 self.updated_branch_nodes = Some(HashMap::default());
82 }
83 }
84
85 pub fn split(mut self) -> (Self, HashMap<Nibbles, BranchNodeCompact>) {
87 let updates = self.updated_branch_nodes.take();
88 (self, updates.unwrap_or_default())
89 }
90
91 pub fn take_proof_nodes(&mut self) -> ProofNodes {
93 self.proof_retainer.take().map(ProofRetainer::into_proof_nodes).unwrap_or_default()
94 }
95
96 pub fn updates_len(&self) -> usize {
99 self.updated_branch_nodes.as_ref().map(|u| u.len()).unwrap_or(0)
100 }
101
102 #[cfg(feature = "std")]
104 pub fn print_stack(&self) {
105 println!("============ STACK ===============");
106 for item in &self.stack {
107 println!("{}", alloy_primitives::hex::encode(item));
108 }
109 println!("============ END STACK ===============");
110 }
111
112 pub fn add_leaf(&mut self, key: Nibbles, value: &[u8]) {
118 assert!(key > self.key, "add_leaf key {:?} self.key {:?}", key, self.key);
119 self.add_leaf_unchecked(key, value);
120 }
121
122 pub fn add_leaf_unchecked(&mut self, key: Nibbles, value: &[u8]) {
127 debug_assert!(key > self.key, "add_leaf_unchecked key {:?} self.key {:?}", key, self.key);
128 if !self.key.is_empty() {
129 self.update(&key);
130 }
131 self.set_key_value(key, HashBuilderValueRef::Bytes(value));
132 }
133
134 pub fn add_branch(&mut self, key: Nibbles, value: B256, stored_in_database: bool) {
136 assert!(
137 key > self.key || (self.key.is_empty() && key.is_empty()),
138 "add_branch key {:?} self.key {:?}",
139 key,
140 self.key
141 );
142 if !self.key.is_empty() {
143 self.update(&key);
144 } else if key.is_empty() {
145 self.stack.push(RlpNode::word_rlp(&value));
146 }
147 self.set_key_value(key, HashBuilderValueRef::Hash(&value));
148 self.stored_in_database = stored_in_database;
149 }
150
151 pub fn root(&mut self) -> B256 {
153 if !self.key.is_empty() {
155 self.update(&Nibbles::default());
156 self.key.clear();
157 self.value.clear();
158 }
159 let root = self.current_root();
160 if root == EMPTY_ROOT_HASH {
161 if let Some(proof_retainer) = self.proof_retainer.as_mut() {
162 proof_retainer.retain(&Nibbles::default(), &[EMPTY_STRING_CODE])
163 }
164 }
165 root
166 }
167
168 #[inline]
169 fn set_key_value(&mut self, key: Nibbles, value: HashBuilderValueRef<'_>) {
170 self.log_key_value("old value");
171 self.key = key;
172 self.value.set_from_ref(value);
173 self.log_key_value("new value");
174 }
175
176 fn log_key_value(&self, msg: &str) {
177 trace!(target: "trie::hash_builder",
178 key = ?self.key,
179 value = ?self.value,
180 "{msg}",
181 );
182 }
183
184 fn current_root(&self) -> B256 {
185 if let Some(node_ref) = self.stack.last() {
186 if let Some(hash) = node_ref.as_hash() {
187 hash
188 } else {
189 keccak256(node_ref)
190 }
191 } else {
192 EMPTY_ROOT_HASH
193 }
194 }
195
196 fn update(&mut self, succeeding: &Nibbles) {
201 let mut build_extensions = false;
202 let mut current = self.key.clone();
204 debug_assert!(!current.is_empty());
205
206 trace!(target: "trie::hash_builder", ?current, ?succeeding, "updating merkle tree");
207
208 let mut i = 0usize;
209 loop {
210 let _span = tracing::trace_span!(target: "trie::hash_builder", "loop", i, ?current, build_extensions).entered();
211
212 let preceding_exists = !self.state_masks.is_empty();
213 let preceding_len = self.state_masks.len().saturating_sub(1);
214
215 let common_prefix_len = succeeding.common_prefix_length(current.as_slice());
216 let len = cmp::max(preceding_len, common_prefix_len);
217 assert!(len < current.len(), "len {} current.len {}", len, current.len());
218
219 trace!(
220 target: "trie::hash_builder",
221 ?len,
222 ?common_prefix_len,
223 ?preceding_len,
224 preceding_exists,
225 "prefix lengths after comparing keys"
226 );
227
228 let extra_digit = current[len];
230 if self.state_masks.len() <= len {
231 let new_len = len + 1;
232 trace!(target: "trie::hash_builder", new_len, old_len = self.state_masks.len(), "scaling state masks to fit");
233 self.state_masks.resize(new_len, TrieMask::default());
234 }
235 self.state_masks[len] |= TrieMask::from_nibble(extra_digit);
236 trace!(
237 target: "trie::hash_builder",
238 ?extra_digit,
239 state_masks = ?self.state_masks,
240 );
241
242 if self.tree_masks.len() < current.len() {
244 self.resize_masks(current.len());
245 }
246
247 let mut len_from = len;
248 if !succeeding.is_empty() || preceding_exists {
249 len_from += 1;
250 }
251 trace!(target: "trie::hash_builder", "skipping {len_from} nibbles");
252
253 let short_node_key = current.slice(len_from..);
255 trace!(target: "trie::hash_builder", ?short_node_key);
256
257 if !build_extensions {
259 match self.value.as_ref() {
260 HashBuilderValueRef::Bytes(leaf_value) => {
261 let leaf_node = LeafNodeRef::new(&short_node_key, leaf_value);
262 self.rlp_buf.clear();
263 let rlp = leaf_node.rlp(&mut self.rlp_buf);
264 trace!(
265 target: "trie::hash_builder",
266 ?leaf_node,
267 ?rlp,
268 "pushing leaf node",
269 );
270 self.stack.push(rlp);
271 self.retain_proof_from_buf(¤t.slice(..len_from));
272 }
273 HashBuilderValueRef::Hash(hash) => {
274 trace!(target: "trie::hash_builder", ?hash, "pushing branch node hash");
275 self.stack.push(RlpNode::word_rlp(hash));
276
277 if self.stored_in_database {
278 self.tree_masks[current.len() - 1] |=
279 TrieMask::from_nibble(current.last().unwrap());
280 }
281 self.hash_masks[current.len() - 1] |=
282 TrieMask::from_nibble(current.last().unwrap());
283
284 build_extensions = true;
285 }
286 }
287 }
288
289 if build_extensions && !short_node_key.is_empty() {
290 self.update_masks(¤t, len_from);
291 let stack_last = self.stack.pop().expect("there should be at least one stack item");
292 let extension_node = ExtensionNodeRef::new(&short_node_key, &stack_last);
293
294 self.rlp_buf.clear();
295 let rlp = extension_node.rlp(&mut self.rlp_buf);
296 trace!(
297 target: "trie::hash_builder",
298 ?extension_node,
299 ?rlp,
300 "pushing extension node",
301 );
302 self.stack.push(rlp);
303 self.retain_proof_from_buf(¤t.slice(..len_from));
304 self.resize_masks(len_from);
305 }
306
307 if preceding_len <= common_prefix_len && !succeeding.is_empty() {
308 trace!(target: "trie::hash_builder", "no common prefix to create branch nodes from, returning");
309 return;
310 }
311
312 if !succeeding.is_empty() || preceding_exists {
314 let children = self.push_branch_node(¤t, len);
316 self.store_branch_node(¤t, len, children);
318 }
319
320 self.state_masks.resize(len, TrieMask::default());
321 self.resize_masks(len);
322
323 if preceding_len == 0 {
324 trace!(target: "trie::hash_builder", "0 or 1 state masks means we have no more elements to process");
325 return;
326 }
327
328 current.truncate(preceding_len);
329 trace!(target: "trie::hash_builder", ?current, "truncated nibbles to {} bytes", preceding_len);
330
331 trace!(target: "trie::hash_builder", state_masks = ?self.state_masks, "popping empty state masks");
332 while self.state_masks.last() == Some(&TrieMask::default()) {
333 self.state_masks.pop();
334 }
335
336 build_extensions = true;
337
338 i += 1;
339 }
340 }
341
342 fn push_branch_node(&mut self, current: &Nibbles, len: usize) -> Vec<B256> {
349 let state_mask = self.state_masks[len];
350 let hash_mask = self.hash_masks[len];
351 let branch_node = BranchNodeRef::new(&self.stack, state_mask);
352 let children = if self.updated_branch_nodes.is_some() {
354 branch_node.child_hashes(hash_mask).collect()
355 } else {
356 vec![]
357 };
358
359 self.rlp_buf.clear();
360 let rlp = branch_node.rlp(&mut self.rlp_buf);
361 self.retain_proof_from_buf(¤t.slice(..len));
362
363 let first_child_idx = self.stack.len() - state_mask.count_ones() as usize;
365 trace!(
366 target: "trie::hash_builder",
367 new_len = first_child_idx,
368 old_len = self.stack.len(),
369 "resizing stack to prepare branch node"
370 );
371 self.stack.resize_with(first_child_idx, Default::default);
372
373 trace!(target: "trie::hash_builder", ?rlp, "pushing branch node with {state_mask:?} mask from stack");
374 self.stack.push(rlp);
375 children
376 }
377
378 fn store_branch_node(&mut self, current: &Nibbles, len: usize, children: Vec<B256>) {
383 if len > 0 {
384 let parent_index = len - 1;
385 self.hash_masks[parent_index] |= TrieMask::from_nibble(current[parent_index]);
386 }
387
388 let store_in_db_trie = !self.tree_masks[len].is_empty() || !self.hash_masks[len].is_empty();
389 if store_in_db_trie {
390 if len > 0 {
391 let parent_index = len - 1;
392 self.tree_masks[parent_index] |= TrieMask::from_nibble(current[parent_index]);
393 }
394
395 if self.updated_branch_nodes.is_some() {
396 let common_prefix = current.slice(..len);
397 let node = BranchNodeCompact::new(
398 self.state_masks[len],
399 self.tree_masks[len],
400 self.hash_masks[len],
401 children,
402 (len == 0).then(|| self.current_root()),
403 );
404 trace!(target: "trie::hash_builder", ?node, "intermediate node");
405 self.updated_branch_nodes.as_mut().unwrap().insert(common_prefix, node);
406 }
407 }
408 }
409
410 fn retain_proof_from_buf(&mut self, prefix: &Nibbles) {
411 if let Some(proof_retainer) = self.proof_retainer.as_mut() {
412 proof_retainer.retain(prefix, &self.rlp_buf)
413 }
414 }
415
416 fn update_masks(&mut self, current: &Nibbles, len_from: usize) {
417 if len_from > 0 {
418 let flag = TrieMask::from_nibble(current[len_from - 1]);
419
420 self.hash_masks[len_from - 1] &= !flag;
421
422 if !self.tree_masks[current.len() - 1].is_empty() {
423 self.tree_masks[len_from - 1] |= flag;
424 }
425 }
426 }
427
428 fn resize_masks(&mut self, new_len: usize) {
429 trace!(
430 target: "trie::hash_builder",
431 new_len,
432 old_tree_mask_len = self.tree_masks.len(),
433 old_hash_mask_len = self.hash_masks.len(),
434 "resizing tree/hash masks"
435 );
436 self.tree_masks.resize(new_len, TrieMask::default());
437 self.hash_masks.resize(new_len, TrieMask::default());
438 }
439}
440
441#[cfg(test)]
442mod tests {
443 use super::*;
444 use crate::{nodes::LeafNode, triehash_trie_root};
445 use alloc::collections::BTreeMap;
446 use alloy_primitives::{b256, hex, U256};
447 use alloy_rlp::Encodable;
448
449 fn assert_hashed_trie_root<'a, I, K>(iter: I)
451 where
452 I: Iterator<Item = (K, &'a U256)>,
453 K: AsRef<[u8]> + Ord,
454 {
455 let hashed = iter
456 .map(|(k, v)| (keccak256(k.as_ref()), alloy_rlp::encode(v).to_vec()))
457 .collect::<BTreeMap<_, _>>();
459
460 let mut hb = HashBuilder::default();
461
462 hashed.iter().for_each(|(key, val)| {
463 let nibbles = Nibbles::unpack(key);
464 hb.add_leaf(nibbles, val);
465 });
466
467 assert_eq!(hb.root(), triehash_trie_root(&hashed));
468 }
469
470 fn assert_trie_root<I, K, V>(iter: I)
472 where
473 I: IntoIterator<Item = (K, V)>,
474 K: AsRef<[u8]> + Ord,
475 V: AsRef<[u8]>,
476 {
477 let mut hb = HashBuilder::default();
478
479 let data = iter.into_iter().collect::<BTreeMap<_, _>>();
480 data.iter().for_each(|(key, val)| {
481 let nibbles = Nibbles::unpack(key);
482 hb.add_leaf(nibbles, val.as_ref());
483 });
484
485 assert_eq!(hb.root(), triehash_trie_root(data));
486 }
487
488 #[test]
489 fn empty() {
490 assert_eq!(HashBuilder::default().root(), EMPTY_ROOT_HASH);
491 }
492
493 #[test]
494 #[cfg(feature = "arbitrary")]
495 #[cfg_attr(miri, ignore = "no proptest")]
496 fn arbitrary_hashed_root() {
497 use proptest::prelude::*;
498 proptest!(|(state: BTreeMap<B256, U256>)| {
499 assert_hashed_trie_root(state.iter());
500 });
501 }
502
503 #[test]
504 fn test_generates_branch_node() {
505 let mut hb = HashBuilder::default().with_updates(true);
506
507 let data = BTreeMap::from([
525 (
526 hex!("1000000000000000000000000000000000000000000000000000000000000000").to_vec(),
529 Vec::new(),
530 ),
531 (
532 hex!("1100000000000000000000000000000000000000000000000000000000000000").to_vec(),
533 Vec::new(),
534 ),
535 (
536 hex!("1110000000000000000000000000000000000000000000000000000000000000").to_vec(),
537 Vec::new(),
538 ),
539 (
540 hex!("1200000000000000000000000000000000000000000000000000000000000000").to_vec(),
541 Vec::new(),
542 ),
543 (
544 hex!("1220000000000000000000000000000000000000000000000000000000000000").to_vec(),
545 Vec::new(),
546 ),
547 (
548 hex!("1320000000000000000000000000000000000000000000000000000000000000").to_vec(),
551 Vec::new(),
552 ),
553 ]);
554 data.iter().for_each(|(key, val)| {
555 let nibbles = Nibbles::unpack(key);
556 hb.add_leaf(nibbles, val.as_ref());
557 });
558 let _root = hb.root();
559
560 let (_, updates) = hb.split();
561
562 let update = updates.get(&Nibbles::from_nibbles_unchecked(hex!("01"))).unwrap();
563 assert_eq!(update.state_mask, TrieMask::new(0b1111));
565 assert_eq!(update.tree_mask, TrieMask::new(0b0000));
567 assert_eq!(update.hash_mask, TrieMask::new(0b0110));
569 assert_eq!(update.hashes.len(), 2);
571
572 assert_eq!(_root, triehash_trie_root(data));
573 }
574
575 #[test]
576 fn test_root_raw_data() {
577 let data = [
578 (hex!("646f").to_vec(), hex!("76657262").to_vec()),
579 (hex!("676f6f64").to_vec(), hex!("7075707079").to_vec()),
580 (hex!("676f6b32").to_vec(), hex!("7075707079").to_vec()),
581 (hex!("676f6b34").to_vec(), hex!("7075707079").to_vec()),
582 ];
583 assert_trie_root(data);
584 }
585
586 #[test]
587 fn test_root_rlp_hashed_data() {
588 let data: HashMap<_, _, _> = HashMap::from([
589 (B256::with_last_byte(1), U256::from(2)),
590 (B256::with_last_byte(3), U256::from(4)),
591 ]);
592 assert_hashed_trie_root(data.iter());
593 }
594
595 #[test]
596 fn test_root_known_hash() {
597 let root_hash = b256!("45596e474b536a6b4d64764e4f75514d544577646c414e684271706871446456");
598 let mut hb = HashBuilder::default();
599 hb.add_branch(Nibbles::default(), root_hash, false);
600 assert_eq!(hb.root(), root_hash);
601 }
602
603 #[test]
604 fn manual_branch_node_ok() {
605 let raw_input = vec![
606 (hex!("646f").to_vec(), hex!("76657262").to_vec()),
607 (hex!("676f6f64").to_vec(), hex!("7075707079").to_vec()),
608 ];
609 let expected = triehash_trie_root(raw_input.clone());
610
611 let mut hb = HashBuilder::default();
613 for (key, val) in &raw_input {
614 hb.add_leaf(Nibbles::unpack(key), val.as_slice());
615 }
616
617 let leaf1 = LeafNode::new(Nibbles::unpack(&raw_input[0].0[1..]), raw_input[0].1.clone());
621 let leaf2 = LeafNode::new(Nibbles::unpack(&raw_input[1].0[1..]), raw_input[1].1.clone());
622 let mut branch: [&dyn Encodable; 17] = [b""; 17];
623 branch[4] = &leaf1;
626 branch[7] = &leaf2;
627 let mut branch_node_rlp = Vec::new();
628 alloy_rlp::encode_list::<_, dyn Encodable>(&branch, &mut branch_node_rlp);
629 let branch_node_hash = keccak256(branch_node_rlp);
630
631 let mut hb2 = HashBuilder::default();
632 hb2.add_branch(Nibbles::from_nibbles_unchecked([0x6]), branch_node_hash, false);
634
635 assert_eq!(hb.root(), expected);
636 assert_eq!(hb2.root(), expected);
637 }
638}