jmt/node_type.rs
1// Copyright (c) The Diem Core Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4//! Node types of [`JellyfishMerkleTree`](crate::JellyfishMerkleTree)
5//!
6//! This module defines two types of Jellyfish Merkle tree nodes: [`InternalNode`]
7//! and [`LeafNode`] as building blocks of a 256-bit
8//! [`JellyfishMerkleTree`](crate::JellyfishMerkleTree). [`InternalNode`] represents a 4-level
9//! binary tree to optimize for IOPS: it compresses a tree with 31 nodes into one node with 16
10//! chidren at the lowest level. [`LeafNode`] stores the full key and the value associated.
11use crate::storage::TreeReader;
12
13use crate::SimpleHasher;
14use alloc::format;
15use alloc::vec::Vec;
16use alloc::{boxed::Box, vec};
17use anyhow::Context;
18use borsh::{BorshDeserialize, BorshSerialize};
19use num_derive::{FromPrimitive, ToPrimitive};
20#[cfg(any(test))]
21use proptest::prelude::*;
22#[cfg(any(test))]
23use proptest_derive::Arbitrary;
24use serde::{Deserialize, Serialize};
25
26use crate::proof::SparseMerkleNode;
27use crate::{
28 types::{
29 nibble::{nibble_path::NibblePath, Nibble},
30 proof::{SparseMerkleInternalNode, SparseMerkleLeafNode},
31 Version,
32 },
33 KeyHash, ValueHash, SPARSE_MERKLE_PLACEHOLDER_HASH,
34};
35
36/// The unique key of each node.
37#[derive(
38 Clone,
39 Debug,
40 Hash,
41 Eq,
42 PartialEq,
43 Ord,
44 PartialOrd,
45 Serialize,
46 Deserialize,
47 borsh::BorshSerialize,
48 borsh::BorshDeserialize,
49)]
50#[cfg_attr(any(test), derive(Arbitrary))]
51pub struct NodeKey {
52 // The version at which the node is created.
53 version: Version,
54 // The nibble path this node represents in the tree.
55 nibble_path: NibblePath,
56}
57
58impl NodeKey {
59 /// Creates a new `NodeKey`.
60 pub fn new(version: Version, nibble_path: NibblePath) -> Self {
61 Self {
62 version,
63 nibble_path,
64 }
65 }
66
67 /// A shortcut to generate a node key consisting of a version and an empty nibble path.
68 pub(crate) fn new_empty_path(version: Version) -> Self {
69 Self::new(version, NibblePath::new(vec![]))
70 }
71
72 /// Gets the version.
73 pub fn version(&self) -> Version {
74 self.version
75 }
76
77 /// Gets the nibble path.
78 pub fn nibble_path(&self) -> &NibblePath {
79 &self.nibble_path
80 }
81
82 /// Generates a child node key based on this node key.
83 pub(crate) fn gen_child_node_key(&self, version: Version, n: Nibble) -> Self {
84 let mut node_nibble_path = self.nibble_path().clone();
85 node_nibble_path.push(n);
86 Self::new(version, node_nibble_path)
87 }
88
89 /// Generates parent node key at the same version based on this node key.
90 pub(crate) fn gen_parent_node_key(&self) -> Self {
91 let mut node_nibble_path = self.nibble_path().clone();
92 assert!(
93 node_nibble_path.pop().is_some(),
94 "Current node key is root.",
95 );
96 Self::new(self.version, node_nibble_path)
97 }
98
99 /// Sets the version to the given version.
100 pub(crate) fn set_version(&mut self, version: Version) {
101 self.version = version;
102 }
103}
104
105#[derive(
106 Clone,
107 Debug,
108 Eq,
109 PartialEq,
110 borsh::BorshSerialize,
111 borsh::BorshDeserialize,
112 Serialize,
113 Deserialize,
114)]
115pub enum NodeType {
116 Leaf,
117 Internal { leaf_count: usize },
118}
119
120#[cfg(any(test))]
121impl Arbitrary for NodeType {
122 type Parameters = ();
123 type Strategy = BoxedStrategy<Self>;
124
125 fn arbitrary_with(_args: ()) -> Self::Strategy {
126 prop_oneof![
127 Just(NodeType::Leaf),
128 (2..100usize).prop_map(|leaf_count| NodeType::Internal { leaf_count })
129 ]
130 .boxed()
131 }
132}
133
134/// Each child of [`InternalNode`] encapsulates a nibble forking at this node.
135#[derive(
136 Clone,
137 Debug,
138 Eq,
139 PartialEq,
140 borsh::BorshSerialize,
141 borsh::BorshDeserialize,
142 Serialize,
143 Deserialize,
144)]
145#[cfg_attr(any(test), derive(Arbitrary))]
146pub struct Child {
147 /// The hash value of this child node.
148 pub hash: [u8; 32],
149 /// `version`, the `nibble_path` of the ['NodeKey`] of this [`InternalNode`] the child belongs
150 /// to and the child's index constitute the [`NodeKey`] to uniquely identify this child node
151 /// from the storage. Used by `[`NodeKey::gen_child_node_key`].
152 pub version: Version,
153 /// Indicates if the child is a leaf, or if it's an internal node, the total number of leaves
154 /// under it.
155 pub node_type: NodeType,
156}
157
158impl Child {
159 pub fn new(hash: [u8; 32], version: Version, node_type: NodeType) -> Self {
160 Self {
161 hash,
162 version,
163 node_type,
164 }
165 }
166
167 pub fn is_leaf(&self) -> bool {
168 matches!(self.node_type, NodeType::Leaf)
169 }
170
171 pub fn leaf_count(&self) -> usize {
172 match self.node_type {
173 NodeType::Leaf => 1,
174 NodeType::Internal { leaf_count } => leaf_count,
175 }
176 }
177}
178
179/// [`Children`] is just a collection of children belonging to a [`InternalNode`], indexed from 0 to
180/// 15, inclusive.
181#[derive(
182 Debug,
183 Clone,
184 PartialEq,
185 Eq,
186 Default,
187 borsh::BorshSerialize,
188 borsh::BorshDeserialize,
189 Serialize,
190 Deserialize,
191)]
192pub struct Children {
193 /// The actual children. We box this array to avoid stack overflows, since the space consumed
194 /// is somewhat large
195 children: Box<[Option<Child>; 16]>,
196 num_children: usize,
197}
198
199#[cfg(any(test))]
200impl Arbitrary for Children {
201 type Parameters = ();
202 type Strategy = BoxedStrategy<Self>;
203
204 fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
205 (any::<Box<[Option<Child>; 16]>>().prop_map(|children| {
206 let num_children = children.iter().filter(|child| child.is_some()).count();
207 Self {
208 children,
209 num_children,
210 }
211 }))
212 .boxed()
213 }
214}
215
216impl Children {
217 /// Create an empty set of children.
218 pub fn new() -> Self {
219 Default::default()
220 }
221
222 /// Insert a new child. Insert is guaranteed not to allocate.
223 pub fn insert(&mut self, nibble: Nibble, child: Child) {
224 let idx = nibble.as_usize();
225 if self.children[idx].is_none() {
226 self.num_children += 1;
227 }
228 self.children[idx] = Some(child);
229 }
230
231 /// Get the child at the provided nibble.
232 pub fn get(&self, nibble: Nibble) -> &Option<Child> {
233 &self.children[nibble.as_usize()]
234 }
235
236 /// Check if the struct contains any children.
237 pub fn is_empty(&self) -> bool {
238 self.num_children == 0
239 }
240
241 /// Remove the child at the provided nibble.
242 pub fn remove(&mut self, nibble: Nibble) {
243 let idx = nibble.as_usize();
244 if self.children[idx].is_some() {
245 self.num_children -= 1;
246 }
247 self.children[idx] = None;
248 }
249
250 /// Returns a (possibly unsorted) iterator over the children.
251 pub fn values(&self) -> impl Iterator<Item = &Child> {
252 self.children.iter().filter_map(|child| child.as_ref())
253 }
254
255 /// Returns a (possibly unsorted) iterator over the children and their respective Nibbles.
256 pub fn iter(&self) -> impl Iterator<Item = (Nibble, &Child)> {
257 self.iter_sorted()
258 }
259
260 /// Returns a (possibly unsorted) mutable iterator over the children, also yielding their respective nibbles.
261 pub fn iter_mut(&mut self) -> impl Iterator<Item = (Nibble, &mut Child)> {
262 self.children
263 .iter_mut()
264 .enumerate()
265 .filter_map(|(nibble, child)| {
266 if let Some(child) = child {
267 Some((Nibble::from(nibble as u8), child))
268 } else {
269 None
270 }
271 })
272 }
273
274 /// Returns the number of children.
275 pub fn num_children(&self) -> usize {
276 self.num_children
277 }
278
279 /// Returns an iterator that yields the children and their respective Nibbles in sorted order.
280 pub fn iter_sorted(&self) -> impl Iterator<Item = (Nibble, &Child)> {
281 self.children
282 .iter()
283 .enumerate()
284 .filter_map(|(nibble, child)| {
285 if let Some(child) = child {
286 Some((Nibble::from(nibble as u8), child))
287 } else {
288 None
289 }
290 })
291 }
292}
293
294/// Represents a 4-level subtree with 16 children at the bottom level. Theoretically, this reduces
295/// IOPS to query a tree by 4x since we compress 4 levels in a standard Merkle tree into 1 node.
296/// Though we choose the same internal node structure as that of Patricia Merkle tree, the root hash
297/// computation logic is similar to a 4-level sparse Merkle tree except for some customizations. See
298/// the `CryptoHash` trait implementation below for details.
299#[derive(
300 Clone,
301 Debug,
302 Eq,
303 PartialEq,
304 Serialize,
305 Deserialize,
306 borsh::BorshSerialize,
307 borsh::BorshDeserialize,
308)]
309pub struct InternalNode {
310 /// Up to 16 children.
311 children: Children,
312 /// Total number of leaves under this internal node
313 leaf_count: usize,
314}
315
316impl SparseMerkleInternalNode {
317 fn from<H: SimpleHasher>(internal_node: InternalNode) -> Self {
318 let bitmaps = internal_node.generate_bitmaps();
319 SparseMerkleInternalNode::new(
320 internal_node.merkle_hash::<H>(0, 8, bitmaps),
321 internal_node.merkle_hash::<H>(8, 8, bitmaps),
322 )
323 }
324}
325
326/// Computes the hash of internal node according to [`JellyfishTree`](crate::JellyfishTree)
327/// data structure in the logical view. `start` and `nibble_height` determine a subtree whose
328/// root hash we want to get. For an internal node with 16 children at the bottom level, we compute
329/// the root hash of it as if a full binary Merkle tree with 16 leaves as below:
330///
331/// ```text
332/// 4 -> +------ root hash ------+
333/// | |
334/// 3 -> +---- # ----+ +---- # ----+
335/// | | | |
336/// 2 -> # # # #
337/// / \ / \ / \ / \
338/// 1 -> # # # # # # # #
339/// / \ / \ / \ / \ / \ / \ / \ / \
340/// 0 -> 0 1 2 3 4 5 6 7 8 9 A B C D E F
341/// ^
342/// height
343/// ```
344///
345/// As illustrated above, at nibble height 0, `0..F` in hex denote 16 chidren hashes. Each `#`
346/// means the hash of its two direct children, which will be used to generate the hash of its
347/// parent with the hash of its sibling. Finally, we can get the hash of this internal node.
348///
349/// However, if an internal node doesn't have all 16 chidren exist at height 0 but just a few of
350/// them, we have a modified hashing rule on top of what is stated above:
351/// 1. From top to bottom, a node will be replaced by a leaf child if the subtree rooted at this
352/// node has only one child at height 0 and it is a leaf child.
353/// 2. From top to bottom, a node will be replaced by the placeholder node if the subtree rooted at
354/// this node doesn't have any child at height 0. For example, if an internal node has 3 leaf
355/// children at index 0, 3, 8, respectively, and 1 internal node at index C, then the computation
356/// graph will be like:
357///
358/// ```text
359/// 4 -> +------ root hash ------+
360/// | |
361/// 3 -> +---- # ----+ +---- # ----+
362/// | | | |
363/// 2 -> # @ 8 #
364/// / \ / \
365/// 1 -> 0 3 # @
366/// / \
367/// 0 -> C @
368/// ^
369/// height
370/// Note: @ denotes placeholder hash.
371/// ```
372#[cfg(any(test))]
373impl Arbitrary for InternalNode {
374 type Parameters = ();
375 type Strategy = BoxedStrategy<Self>;
376
377 fn arbitrary_with(_args: ()) -> Self::Strategy {
378 (any::<Children>().prop_filter(
379 "InternalNode constructor panics when its only child is a leaf.",
380 |children| {
381 !(children.num_children() == 1
382 && children.values().next().expect("Must exist.").is_leaf())
383 },
384 ))
385 .prop_map(InternalNode::new)
386 .boxed()
387 }
388}
389
390/// Helper for `InternalNode` implementations. Test if the leaf exaclty has one child within the width range specified
391fn has_only_child(width: u8, range_existence_bitmap: u16, range_leaf_bitmap: u16) -> bool {
392 width == 1 || (range_existence_bitmap.count_ones() == 1 && range_leaf_bitmap != 0)
393}
394
395/// Helper for `InternalNode` implementations. Test if the leaf exactly has one child *at the position n*
396/// within the width range specified
397fn has_child(
398 width: u8,
399 range_existence_bitmap: u16,
400 n_bitmap: u16,
401 range_leaf_bitmap: u16,
402) -> bool {
403 width == 1 || (range_existence_bitmap == n_bitmap && range_leaf_bitmap != 0)
404}
405
406impl InternalNode {
407 /// Creates a new Internal node.
408 pub fn new(children: Children) -> Self {
409 // Assert the internal node must have >= 1 children. If it only has one child, it cannot be
410 // a leaf node. Otherwise, the leaf node should be a child of this internal node's parent.
411 assert!(!children.is_empty(), "Children must not be empty");
412 if children.num_children() == 1 {
413 assert!(
414 !children
415 .values()
416 .next()
417 .expect("Must have 1 element")
418 .is_leaf(),
419 "If there's only one child, it must not be a leaf."
420 );
421 }
422
423 let leaf_count = Self::sum_leaf_count(&children);
424 Self {
425 children,
426 leaf_count,
427 }
428 }
429
430 fn sum_leaf_count(children: &Children) -> usize {
431 let mut leaf_count = 0;
432 for child in children.values() {
433 let n = child.leaf_count();
434 leaf_count += n;
435 }
436 leaf_count
437 }
438
439 pub fn leaf_count(&self) -> usize {
440 self.leaf_count
441 }
442
443 pub fn node_type(&self) -> NodeType {
444 NodeType::Internal {
445 leaf_count: self.leaf_count,
446 }
447 }
448
449 pub fn hash<H: SimpleHasher>(&self) -> [u8; 32] {
450 self.merkle_hash::<H>(
451 0, /* start index */
452 16, /* the number of leaves in the subtree of which we want the hash of root */
453 self.generate_bitmaps(),
454 )
455 }
456
457 pub fn children_sorted(&self) -> impl Iterator<Item = (Nibble, &Child)> {
458 // Previously this used `.sorted_by_key()` directly on the iterator but this does not appear
459 // to be available in itertools (it does not seem to ever have existed???) for unknown
460 // reasons. This satisfies the same behavior. ¯\_(ツ)_/¯
461 self.children.iter_sorted()
462 }
463
464 pub fn children_unsorted(&self) -> impl Iterator<Item = (Nibble, &Child)> {
465 self.children.iter()
466 }
467
468 /// Gets the `n`-th child.
469 pub fn child(&self, n: Nibble) -> Option<&Child> {
470 self.children.get(n).as_ref()
471 }
472
473 /// Generates `existence_bitmap` and `leaf_bitmap` as a pair of `u16`s: child at index `i`
474 /// exists if `existence_bitmap[i]` is set; child at index `i` is leaf node if
475 /// `leaf_bitmap[i]` is set.
476 pub fn generate_bitmaps(&self) -> (u16, u16) {
477 let mut existence_bitmap = 0;
478 let mut leaf_bitmap = 0;
479 for (nibble, child) in self.children.iter() {
480 let i = u8::from(nibble);
481 existence_bitmap |= 1u16 << i;
482 if child.is_leaf() {
483 leaf_bitmap |= 1u16 << i;
484 }
485 }
486 // `leaf_bitmap` must be a subset of `existence_bitmap`.
487 assert_eq!(existence_bitmap | leaf_bitmap, existence_bitmap);
488 (existence_bitmap, leaf_bitmap)
489 }
490
491 /// Given a range [start, start + width), returns the sub-bitmap of that range.
492 fn range_bitmaps(start: u8, width: u8, bitmaps: (u16, u16)) -> (u16, u16) {
493 assert!(start < 16 && width.count_ones() == 1 && start % width == 0);
494 assert!(width <= 16 && (start + width) <= 16);
495 // A range with `start == 8` and `width == 4` will generate a mask 0b0000111100000000.
496 // use as converting to smaller integer types when 'width == 16'
497 let mask = (((1u32 << width) - 1) << start) as u16;
498 (bitmaps.0 & mask, bitmaps.1 & mask)
499 }
500
501 /// [`build_sibling`] builds the sibling contained in the merkle tree between
502 /// [start; start+width) under the internal node (`self`) using the `TreeReader` as
503 /// a node reader to get the leaves/internal nodes at the bottom level of this internal node
504 fn build_sibling<H: SimpleHasher>(
505 &self,
506 tree_reader: &impl TreeReader,
507 node_key: &NodeKey,
508 start: u8,
509 width: u8,
510 (existence_bitmap, leaf_bitmap): (u16, u16),
511 ) -> SparseMerkleNode {
512 // Given a bit [start, 1 << nibble_height], return the value of that range.
513 let (range_existence_bitmap, range_leaf_bitmap) =
514 Self::range_bitmaps(start, width, (existence_bitmap, leaf_bitmap));
515 if range_existence_bitmap == 0 {
516 // No child under this subtree
517 SparseMerkleNode::Null
518 } else if has_only_child(width, range_existence_bitmap, range_leaf_bitmap) {
519 // Only 1 leaf child under this subtree or reach the lowest level
520 let only_child_index = Nibble::from(range_existence_bitmap.trailing_zeros() as u8);
521
522 let child = self
523 .child(only_child_index)
524 .with_context(|| {
525 format!(
526 "Corrupted internal node: existence_bitmap indicates \
527 the existence of a non-exist child at index {:x}",
528 only_child_index
529 )
530 })
531 .unwrap();
532
533 let child_node = tree_reader
534 .get_node(&node_key.gen_child_node_key(child.version, only_child_index))
535 .with_context(|| {
536 format!(
537 "Corruption error: the merkle tree reader supplied cannot find \
538 the child of version {:?} at index {:x}.",
539 child.version, only_child_index
540 )
541 })
542 .unwrap();
543
544 match child_node {
545 Node::Internal(node) => {
546 SparseMerkleNode::Internal(SparseMerkleInternalNode::from::<H>(node))
547 }
548 Node::Leaf(node) => SparseMerkleNode::Leaf(SparseMerkleLeafNode::from(node)),
549 Node::Null => unreachable!("Impossible to get a null node at this location"),
550 }
551 } else {
552 let left_child = self.merkle_hash::<H>(
553 start,
554 width / 2,
555 (range_existence_bitmap, range_leaf_bitmap),
556 );
557 let right_child = self.merkle_hash::<H>(
558 start + width / 2,
559 width / 2,
560 (range_existence_bitmap, range_leaf_bitmap),
561 );
562 SparseMerkleNode::Internal(SparseMerkleInternalNode::new(left_child, right_child))
563 }
564 }
565
566 fn merkle_hash<H: SimpleHasher>(
567 &self,
568 start: u8,
569 width: u8,
570 (existence_bitmap, leaf_bitmap): (u16, u16),
571 ) -> [u8; 32] {
572 // Given a bit [start, 1 << nibble_height], return the value of that range.
573 let (range_existence_bitmap, range_leaf_bitmap) =
574 Self::range_bitmaps(start, width, (existence_bitmap, leaf_bitmap));
575 if range_existence_bitmap == 0 {
576 // No child under this subtree
577 SPARSE_MERKLE_PLACEHOLDER_HASH
578 } else if has_only_child(width, range_existence_bitmap, range_leaf_bitmap) {
579 // Only 1 leaf child under this subtree or reach the lowest level
580 let only_child_index = Nibble::from(range_existence_bitmap.trailing_zeros() as u8);
581 self.child(only_child_index)
582 .with_context(|| {
583 format!(
584 "Corrupted internal node: existence_bitmap indicates \
585 the existence of a non-exist child at index {:x}",
586 only_child_index
587 )
588 })
589 .unwrap()
590 .hash
591 } else {
592 let left_child = self.merkle_hash::<H>(
593 start,
594 width / 2,
595 (range_existence_bitmap, range_leaf_bitmap),
596 );
597 let right_child = self.merkle_hash::<H>(
598 start + width / 2,
599 width / 2,
600 (range_existence_bitmap, range_leaf_bitmap),
601 );
602 SparseMerkleInternalNode::new(left_child, right_child).hash::<H>()
603 }
604 }
605
606 /// Gets the child without its corresponding siblings (like using
607 /// [`get_only_child_with_siblings`](InternalNode::get_only_child_with_siblings) and dropping the
608 /// siblings, but more efficient).
609 pub fn get_only_child_without_siblings(
610 &self,
611 node_key: &NodeKey,
612 n: Nibble,
613 ) -> Option<NodeKey> {
614 let (existence_bitmap, leaf_bitmap) = self.generate_bitmaps();
615
616 // Nibble height from 3 to 0.
617 for h in (0..4).rev() {
618 // Get the number of children of the internal node that each subtree at this height
619 // covers.
620 let width = 1 << h;
621 let child_half_start = get_child_half_start(n, h);
622
623 let (range_existence_bitmap, range_leaf_bitmap) =
624 Self::range_bitmaps(child_half_start, width, (existence_bitmap, leaf_bitmap));
625
626 if range_existence_bitmap == 0 {
627 // No child in this range.
628 return None;
629 } else if has_only_child(width, range_existence_bitmap, range_leaf_bitmap) {
630 // Return the only 1 leaf child under this subtree or reach the lowest level
631 // Even this leaf child is not the n-th child, it should be returned instead of
632 // `None` because it's existence indirectly proves the n-th child doesn't exist.
633 // Please read proof format for details.
634 let only_child_index = Nibble::from(range_existence_bitmap.trailing_zeros() as u8);
635
636 let only_child_version = self
637 .child(only_child_index)
638 // Should be guaranteed by the self invariants, but these are not easy to express at the moment
639 .with_context(|| {
640 format!(
641 "Corrupted internal node: child_bitmap indicates \
642 the existence of a non-exist child at index {:x}",
643 only_child_index
644 )
645 })
646 .unwrap()
647 .version;
648
649 return Some(node_key.gen_child_node_key(only_child_version, only_child_index));
650 }
651 }
652 unreachable!("Impossible to get here without returning even at the lowest level.")
653 }
654
655 /// Gets the child and its corresponding siblings that are necessary to generate the proof for
656 /// the `n`-th child. This function will **either** return the child that matches the nibble n or the only
657 /// child in the largest width range pointed by n. If it is an existence proof, the returned child must be the `n`-th
658 /// child; otherwise, the returned child may be another child in the same nibble pointed by n.
659 /// See inline explanation for details. When calling this function with n = 11
660 /// (node `b` in the following graph), the range at each level is illustrated as a pair of square brackets:
661 ///
662 /// ```text
663 /// 4 [f e d c b a 9 8 7 6 5 4 3 2 1 0] -> root level
664 /// ---------------------------------------------------------------
665 /// 3 [f e d c b a 9 8] [7 6 5 4 3 2 1 0] width = 8
666 /// chs <--┘ shs <--┘
667 /// 2 [f e d c] [b a 9 8] [7 6 5 4] [3 2 1 0] width = 4
668 /// shs <--┘ └--> chs
669 /// 1 [f e] [d c] [b a] [9 8] [7 6] [5 4] [3 2] [1 0] width = 2
670 /// chs <--┘ └--> shs
671 /// 0 [f] [e] [d] [c] [b] [a] [9] [8] [7] [6] [5] [4] [3] [2] [1] [0] width = 1
672 /// ^ chs <--┘ └--> shs
673 /// | MSB|<---------------------- uint 16 ---------------------------->|LSB
674 /// height chs: `child_half_start` shs: `sibling_half_start`
675 /// ```
676 fn get_child_with_siblings_helper<H: SimpleHasher>(
677 &self,
678 tree_reader: &impl TreeReader,
679 node_key: &NodeKey,
680 n: Nibble,
681 get_only_child: bool,
682 ) -> (Option<NodeKey>, Vec<SparseMerkleNode>) {
683 let mut siblings: Vec<SparseMerkleNode> = vec![];
684 let (existence_bitmap, leaf_bitmap) = self.generate_bitmaps();
685
686 let n_bitmap = 1 << n.as_usize();
687
688 // Nibble height from 3 to 0.
689 for h in (0..4).rev() {
690 // Get the number of children of the internal node that each subtree at this height
691 // covers.
692 let width = 1 << h;
693 let (child_half_start, sibling_half_start) = get_child_and_sibling_half_start(n, h);
694 // Compute the root hash of the subtree rooted at the sibling of `r`.
695 siblings.push(self.build_sibling::<H>(
696 tree_reader,
697 node_key,
698 sibling_half_start,
699 width,
700 (existence_bitmap, leaf_bitmap),
701 ));
702
703 let (range_existence_bitmap, range_leaf_bitmap) =
704 Self::range_bitmaps(child_half_start, width, (existence_bitmap, leaf_bitmap));
705
706 if range_existence_bitmap == 0 {
707 // No child in this range.
708 return (None, siblings);
709 } else if get_only_child
710 && (has_only_child(width, range_existence_bitmap, range_leaf_bitmap))
711 {
712 // Return the only 1 leaf child under this subtree or reach the lowest level
713 // Even this leaf child is not the n-th child, it should be returned instead of
714 // `None` because it's existence indirectly proves the n-th child doesn't exist.
715 // Please read proof format for details.
716 let only_child_index = Nibble::from(range_existence_bitmap.trailing_zeros() as u8);
717 return (
718 {
719 let only_child_version = self
720 .child(only_child_index)
721 // Should be guaranteed by the self invariants, but these are not easy to express at the moment
722 .with_context(|| {
723 format!(
724 "Corrupted internal node: child_bitmap indicates \
725 the existence of a non-exist child at index {:x}",
726 only_child_index
727 )
728 })
729 .unwrap()
730 .version;
731 Some(node_key.gen_child_node_key(only_child_version, only_child_index))
732 },
733 siblings,
734 );
735 } else if !get_only_child
736 && (has_child(width, range_existence_bitmap, n_bitmap, range_leaf_bitmap))
737 {
738 // Early return the child in that subtree iff it is the only child and the nibble points
739 // to it
740 return (
741 {
742 let only_child_version = self
743 .child(n)
744 // Should be guaranteed by the self invariants, but these are not easy to express at the moment
745 .with_context(|| {
746 format!(
747 "Corrupted internal node: child_bitmap indicates \
748 the existence of a non-exist child at index {:x}",
749 n
750 )
751 })
752 .unwrap()
753 .version;
754 Some(node_key.gen_child_node_key(only_child_version, n))
755 },
756 siblings,
757 );
758 }
759 }
760 unreachable!("Impossible to get here without returning even at the lowest level.")
761 }
762
763 /// [`get_child_with_siblings`] will return the child from this subtree that matches the nibble n in addition
764 /// to building the list of its sibblings. This function has the same behavior as [`child`].
765 pub(crate) fn get_child_with_siblings<H: SimpleHasher>(
766 &self,
767 tree_cache: &impl TreeReader,
768 node_key: &NodeKey,
769 n: Nibble,
770 ) -> (Option<NodeKey>, Vec<SparseMerkleNode>) {
771 self.get_child_with_siblings_helper::<H>(tree_cache, node_key, n, false)
772 }
773
774 /// [`get_only_child_with_siblings`] will **either** return the child that matches the nibble n or the only
775 /// child in the largest width range pointed by n (see the helper function [`get_child_with_siblings_helper`] for more information).
776 ///
777 /// Even this leaf child is not the n-th child, it should be returned instead of
778 /// `None` because it's existence indirectly proves the n-th child doesn't exist.
779 /// Please read proof format for details.
780 pub(crate) fn get_only_child_with_siblings<H: SimpleHasher>(
781 &self,
782 tree_reader: &impl TreeReader,
783 node_key: &NodeKey,
784 n: Nibble,
785 ) -> (Option<NodeKey>, Vec<SparseMerkleNode>) {
786 self.get_child_with_siblings_helper::<H>(tree_reader, node_key, n, true)
787 }
788
789 #[cfg(test)]
790 pub(crate) fn children(&self) -> &Children {
791 &self.children
792 }
793}
794
795/// Given a nibble, computes the start position of its `child_half_start` and `sibling_half_start`
796/// at `height` level.
797pub(crate) fn get_child_and_sibling_half_start(n: Nibble, height: u8) -> (u8, u8) {
798 // Get the index of the first child belonging to the same subtree whose root, let's say `r` is
799 // at `height` that the n-th child belongs to.
800 // Note: `child_half_start` will be always equal to `n` at height 0.
801 let child_half_start = (0xff << height) & u8::from(n);
802
803 // Get the index of the first child belonging to the subtree whose root is the sibling of `r`
804 // at `height`.
805 let sibling_half_start = child_half_start ^ (1 << height);
806
807 (child_half_start, sibling_half_start)
808}
809
810/// Given a nibble, computes the start position of its `child_half_start` at `height` level.
811pub(crate) fn get_child_half_start(n: Nibble, height: u8) -> u8 {
812 // Get the index of the first child belonging to the same subtree whose root, let's say `r` is
813 // at `height` that the n-th child belongs to.
814 // Note: `child_half_start` will be always equal to `n` at height 0.
815 (0xff << height) & u8::from(n)
816}
817
818/// Represents a key-value pair in the map.
819///
820/// Note: this does not store the key itself.
821#[derive(
822 Clone,
823 Debug,
824 Eq,
825 PartialEq,
826 Serialize,
827 Deserialize,
828 borsh::BorshSerialize,
829 borsh::BorshDeserialize,
830)]
831pub struct LeafNode {
832 /// The hash of the key for this entry.
833 key_hash: KeyHash,
834 /// The hash of the value for this entry.
835 value_hash: ValueHash,
836}
837
838impl LeafNode {
839 /// Creates a new leaf node.
840 pub fn new(key_hash: KeyHash, value_hash: ValueHash) -> Self {
841 Self {
842 key_hash,
843 value_hash,
844 }
845 }
846
847 /// Gets the key hash.
848 pub fn key_hash(&self) -> KeyHash {
849 self.key_hash
850 }
851
852 /// Gets the associated value hash.
853 pub(crate) fn value_hash(&self) -> ValueHash {
854 self.value_hash
855 }
856
857 pub fn hash<H: SimpleHasher>(&self) -> [u8; 32] {
858 SparseMerkleLeafNode::new(self.key_hash, self.value_hash).hash::<H>()
859 }
860}
861
862impl From<LeafNode> for SparseMerkleLeafNode {
863 fn from(leaf_node: LeafNode) -> Self {
864 Self::new(leaf_node.key_hash, leaf_node.value_hash)
865 }
866}
867
868#[repr(u8)]
869#[derive(FromPrimitive, ToPrimitive, BorshDeserialize, BorshSerialize)]
870#[borsh(use_discriminant = false)]
871enum NodeTag {
872 Null = 0,
873 Leaf = 1,
874 Internal = 2,
875}
876
877/// The concrete node type of [`JellyfishMerkleTree`](crate::JellyfishMerkleTree).
878#[derive(Clone, Debug, Eq, PartialEq, BorshSerialize, BorshDeserialize, Serialize, Deserialize)]
879pub enum Node {
880 /// Represents `null`.
881 Null,
882 /// A wrapper of [`InternalNode`].
883 Internal(InternalNode),
884 /// A wrapper of [`LeafNode`].
885 Leaf(LeafNode),
886}
887
888impl From<InternalNode> for Node {
889 fn from(node: InternalNode) -> Self {
890 Node::Internal(node)
891 }
892}
893
894impl From<InternalNode> for Children {
895 fn from(node: InternalNode) -> Self {
896 node.children
897 }
898}
899
900impl From<LeafNode> for Node {
901 fn from(node: LeafNode) -> Self {
902 Node::Leaf(node)
903 }
904}
905
906impl Node {
907 /// Creates the [`Null`](Node::Null) variant.
908 pub(crate) fn new_null() -> Self {
909 Node::Null
910 }
911
912 /// Creates the [`Internal`](Node::Internal) variant.
913 #[cfg(any(test))]
914 pub(crate) fn new_internal(children: Children) -> Self {
915 Node::Internal(InternalNode::new(children))
916 }
917
918 /// Creates the [`Leaf`](Node::Leaf) variant.
919 pub(crate) fn new_leaf(key_hash: KeyHash, value_hash: ValueHash) -> Self {
920 Node::Leaf(LeafNode::new(key_hash, value_hash))
921 }
922
923 /// Creates the [`Leaf`](Node::Leaf) variant by hashing a raw value.
924 #[cfg(any(test))]
925 pub(crate) fn leaf_from_value<H: SimpleHasher>(
926 key_hash: KeyHash,
927 value: impl AsRef<[u8]>,
928 ) -> Self {
929 Node::Leaf(LeafNode::new(key_hash, ValueHash::with::<H>(value)))
930 }
931
932 /// Returns `true` if the node is a leaf node.
933 pub(crate) fn is_leaf(&self) -> bool {
934 matches!(self, Node::Leaf(_))
935 }
936
937 /// Returns `NodeType`
938 pub(crate) fn node_type(&self) -> NodeType {
939 match self {
940 // The returning value will be used to construct a `Child` of a internal node, while an
941 // internal node will never have a child of Node::Null.
942 Self::Null => unreachable!(),
943 Self::Leaf(_) => NodeType::Leaf,
944 Self::Internal(n) => n.node_type(),
945 }
946 }
947
948 /// Returns leaf count if known
949 pub(crate) fn leaf_count(&self) -> usize {
950 match self {
951 Node::Null => 0,
952 Node::Leaf(_) => 1,
953 Node::Internal(internal_node) => internal_node.leaf_count,
954 }
955 }
956
957 /// Computes the hash of nodes.
958 pub(crate) fn hash<H: SimpleHasher>(&self) -> [u8; 32] {
959 match self {
960 Node::Null => SPARSE_MERKLE_PLACEHOLDER_HASH,
961 Node::Internal(internal_node) => internal_node.hash::<H>(),
962 Node::Leaf(leaf_node) => leaf_node.hash::<H>(),
963 }
964 }
965}