diem_jellyfish_merkle/lib.rs
1// Copyright (c) The Diem Core Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4// Copyright 2021 Conflux Foundation. All rights reserved.
5// Conflux is free software and distributed under GNU General Public License.
6// See http://www.gnu.org/licenses/
7
8#![forbid(unsafe_code)]
9
10//! This module implements [`JellyfishMerkleTree`] backed by storage module. The
11//! tree itself doesn't persist anything, but realizes the logic of R/W only.
12//! The write path will produce all the intermediate results in a batch for
13//! storage layer to commit and the read path will return results directly. The
14//! public APIs are only [`new`], [`put_value_sets`], [`put_value_set`] and
15//! [`get_with_proof`]. After each put with a `value_set` based on a known
16//! version, the tree will return a new root hash with a [`TreeUpdateBatch`]
17//! containing all the new nodes and indices of stale nodes.
18//!
19//! A Jellyfish Merkle Tree itself logically is a 256-bit sparse Merkle tree
20//! with an optimization that any subtree containing 0 or 1 leaf node will be
21//! replaced by that leaf node or a placeholder node with default hash value.
22//! With this optimization we can save CPU by avoiding hashing on many sparse
23//! levels in the tree. Physically, the tree is structurally similar to the
24//! modified Patricia Merkle tree of Ethereum but with some modifications. A
25//! standard Jellyfish Merkle tree will look like the following figure:
26//!
27//! ```text
28//! .──────────────────────.
29//! _.─────' `──────.
30//! _.──' `───.
31//! _.─' `──.
32//! _.─' `──.
33//! ,' `.
34//! ,─' '─.
35//! ,' `.
36//! ,' `.
37//! ╱ ╲
38//! ╱ ╲
39//! ╱ ╲
40//! ╱ ╲
41//! ; :
42//! ; :
43//! ; :
44//! │ │
45//! +──────────────────────────────────────────────────────────────────────────────────────────────+
46//! .''. .''. .''. .''. .''. .''. .''. .''. .''. .''. .''. .''. .''. .''. .''. .''.
47//! / \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \/ \
48//! +----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----+
49//! ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( (
50//! ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) )
51//! ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( (
52//! ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) )
53//! ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( (
54//! ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) )
55//! ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( (
56//! ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) ) )
57//! ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( (
58//! ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■ ■
59//! ■: the [`Value`] type this tree stores.
60//! ```
61//!
62//! A Jellyfish Merkle Tree consists of [`InternalNode`] and [`LeafNode`].
63//! [`InternalNode`] is like branch node in ethereum patricia merkle with 16
64//! children to represent a 4-level binary tree and [`LeafNode`] is similar to
65//! that in patricia merkle too. In the above figure, each `bell` in the
66//! jellyfish is an [`InternalNode`] while each tentacle is a [`LeafNode`]. It
67//! is noted that Jellyfish merkle doesn't have a counterpart for `extension`
68//! node of ethereum patricia merkle.
69//!
70//! [`JellyfishMerkleTree`]: struct.JellyfishMerkleTree.html
71//! [`new`]: struct.JellyfishMerkleTree.html#method.new
72//! [`put_value_sets`]: struct.JellyfishMerkleTree.html#method.put_value_sets
73//! [`put_value_set`]: struct.JellyfishMerkleTree.html#method.put_value_set
74//! [`get_with_proof`]: struct.JellyfishMerkleTree.html#method.get_with_proof
75//! [`TreeUpdateBatch`]: struct.TreeUpdateBatch.html
76//! [`InternalNode`]: node_type/struct.InternalNode.html
77//! [`LeafNode`]: node_type/struct.LeafNode.html
78
79pub mod iterator;
80#[cfg(test)]
81mod jellyfish_merkle_test;
82pub mod metrics;
83#[cfg(any(test, feature = "fuzzing"))]
84mod mock_tree_store;
85mod nibble_path;
86pub mod node_type;
87pub mod restore;
88#[cfg(any(test, feature = "fuzzing"))]
89pub mod test_helper;
90mod tree_cache;
91
92use anyhow::{bail, ensure, format_err, Result};
93use diem_crypto::{hash::CryptoHash, HashValue};
94use diem_nibble::Nibble;
95use diem_types::{
96 proof::{SparseMerkleProof, SparseMerkleRangeProof},
97 transaction::Version,
98};
99use nibble_path::{skip_common_prefix, NibbleIterator, NibblePath};
100use node_type::{Child, Children, InternalNode, LeafNode, Node, NodeKey};
101#[cfg(any(test, feature = "fuzzing"))]
102use proptest::arbitrary::Arbitrary;
103#[cfg(any(test, feature = "fuzzing"))]
104use proptest_derive::Arbitrary;
105use serde::{de::DeserializeOwned, Serialize};
106use std::{
107 collections::{BTreeMap, BTreeSet},
108 marker::PhantomData,
109};
110use thiserror::Error;
111use tree_cache::TreeCache;
112
113#[derive(Error, Debug)]
114#[error("Missing state root node at version {version}, probably pruned.")]
115pub struct MissingRootError {
116 pub version: Version,
117}
118
119/// The hardcoded maximum height of a [`JellyfishMerkleTree`] in nibbles.
120pub const ROOT_NIBBLE_HEIGHT: usize = HashValue::LENGTH * 2;
121
122/// `TreeReader` defines the interface between
123/// [`JellyfishMerkleTree`](struct.JellyfishMerkleTree.html)
124/// and underlying storage holding nodes.
125pub trait TreeReader<V> {
126 /// Gets node given a node key. Returns error if the node does not exist.
127 fn get_node(&self, node_key: &NodeKey) -> Result<Node<V>> {
128 self.get_node_option(node_key)?
129 .ok_or_else(|| format_err!("Missing node at {:?}.", node_key))
130 }
131
132 /// Gets node given a node key. Returns `None` if the node does not exist.
133 fn get_node_option(&self, node_key: &NodeKey) -> Result<Option<Node<V>>>;
134
135 /// Gets the rightmost leaf. Note that this assumes we are in the process of
136 /// restoring the tree and all nodes are at the same version.
137 fn get_rightmost_leaf(&self) -> Result<Option<(NodeKey, LeafNode<V>)>>;
138}
139
140pub trait TreeWriter<V> {
141 /// Writes a node batch into storage.
142 fn write_node_batch(&self, node_batch: &NodeBatch<V>) -> Result<()>;
143}
144
145/// `Value` defines the types of data that can be stored in a Jellyfish Merkle
146/// tree.
147pub trait Value: Clone + CryptoHash + Serialize + DeserializeOwned {}
148
149/// `TestValue` defines the types of data that can be stored in a Jellyfish
150/// Merkle tree and used in tests.
151#[cfg(any(test, feature = "fuzzing"))]
152pub trait TestValue:
153 Value + Arbitrary + Clone + std::fmt::Debug + Eq + PartialEq + 'static
154{
155}
156
157// This crate still depends on types for a few things, therefore we implement
158// `Value` and `TestValue` for `AccountStateBlob` here. Ideally the module that
159// defines the specific value like `AccountStateBlob` should import the `Value`
160// trait and implement it there.
161impl Value for diem_types::account_state_blob::AccountStateBlob {}
162#[cfg(any(test, feature = "fuzzing"))]
163impl TestValue for diem_types::account_state_blob::AccountStateBlob {}
164
165/// Node batch that will be written into db atomically with other batches.
166pub type NodeBatch<V> = BTreeMap<NodeKey, Node<V>>;
167/// [`StaleNodeIndex`](struct.StaleNodeIndex.html) batch that will be written
168/// into db atomically with other batches.
169pub type StaleNodeIndexBatch = BTreeSet<StaleNodeIndex>;
170
171#[derive(Clone, Debug, Default, Eq, PartialEq)]
172pub struct NodeStats {
173 pub new_nodes: usize,
174 pub new_leaves: usize,
175 pub stale_nodes: usize,
176 pub stale_leaves: usize,
177}
178
179/// Indicates a node becomes stale since `stale_since_version`.
180#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
181#[cfg_attr(any(test, feature = "fuzzing"), derive(Arbitrary))]
182pub struct StaleNodeIndex {
183 /// The version since when the node is overwritten and becomes stale.
184 pub stale_since_version: Version,
185 /// The [`NodeKey`](node_type/struct.NodeKey.html) identifying the node
186 /// associated with this record.
187 pub node_key: NodeKey,
188}
189
190/// This is a wrapper of [`NodeBatch`](type.NodeBatch.html),
191/// [`StaleNodeIndexBatch`](type.StaleNodeIndexBatch.html) and some stats of
192/// nodes that represents the incremental updates of a tree and pruning indices
193/// after applying a write set, which is a vector of `hashed_account_address`
194/// and `new_value` pairs.
195#[derive(Clone, Debug, Default, Eq, PartialEq)]
196pub struct TreeUpdateBatch<V> {
197 pub node_batch: NodeBatch<V>,
198 pub stale_node_index_batch: StaleNodeIndexBatch,
199 pub node_stats: Vec<NodeStats>,
200}
201
202/// The Jellyfish Merkle tree data structure. See [`crate`] for description.
203pub struct JellyfishMerkleTree<'a, R, V> {
204 reader: &'a R,
205 phantom_value: PhantomData<V>,
206}
207
208impl<'a, R, V> JellyfishMerkleTree<'a, R, V>
209where
210 R: 'a + TreeReader<V>,
211 V: Value,
212{
213 /// Creates a `JellyfishMerkleTree` backed by the given
214 /// [`TreeReader`](trait.TreeReader.html).
215 pub fn new(reader: &'a R) -> Self {
216 Self {
217 reader,
218 phantom_value: PhantomData,
219 }
220 }
221
222 /// This is a convenient function that calls
223 /// [`put_value_sets`](struct.JellyfishMerkleTree.html#method.
224 /// put_value_sets) with a single `keyed_value_set`.
225 #[cfg(any(test, feature = "fuzzing"))]
226 pub fn put_value_set(
227 &self, value_set: Vec<(HashValue, V)>, version: Version,
228 ) -> Result<(HashValue, TreeUpdateBatch<V>)> {
229 let (root_hashes, tree_update_batch) =
230 self.put_value_sets(vec![value_set], version)?;
231 assert_eq!(
232 root_hashes.len(),
233 1,
234 "root_hashes must consist of a single value.",
235 );
236 Ok((root_hashes[0], tree_update_batch))
237 }
238
239 /// Returns the new nodes and values in a batch after applying `value_set`.
240 /// For example, if after transaction `T_i` the committed state of tree
241 /// in the persistent storage looks like the following structure:
242 ///
243 /// ```text
244 /// S_i
245 /// / \
246 /// . .
247 /// . .
248 /// / \
249 /// o x
250 /// / \
251 /// A B
252 /// storage (disk)
253 /// ```
254 ///
255 /// where `A` and `B` denote the states of two adjacent accounts, and `x` is
256 /// a sibling subtree of the path from root to A and B in the tree. Then
257 /// a `value_set` produced by the next transaction `T_{i+1}` modifies
258 /// other accounts `C` and `D` exist in the subtree under `x`, a
259 /// new partial tree will be constructed in memory and the structure will
260 /// be:
261 ///
262 /// ```text
263 /// S_i | S_{i+1}
264 /// / \ | / \
265 /// . . | . .
266 /// . . | . .
267 /// / \ | / \
268 /// / x | / x'
269 /// o<-------------+- / \
270 /// / \ | C D
271 /// A B |
272 /// storage (disk) | cache (memory)
273 /// ```
274 ///
275 /// With this design, we are able to query the global state in persistent
276 /// storage and generate the proposed tree delta based on a specific
277 /// root hash and `value_set`. For example, if we want to execute
278 /// another transaction `T_{i+1}'`, we can use the tree `S_i` in storage
279 /// and apply the `value_set` of transaction `T_{i+1}`. Then if the storage
280 /// commits the returned batch, the state `S_{i+1}` is ready to be read
281 /// from the tree by calling \[`get_with_proof`\](struct.
282 /// JellyfishMerkleTree.html#method.get_with_proof). Anything inside the
283 /// batch is not reachable from public interfaces before being committed.
284 pub fn put_value_sets(
285 &self, value_sets: Vec<Vec<(HashValue, V)>>, first_version: Version,
286 ) -> Result<(Vec<HashValue>, TreeUpdateBatch<V>)> {
287 let mut tree_cache = TreeCache::new(self.reader, first_version)?;
288 for (idx, value_set) in value_sets.into_iter().enumerate() {
289 assert!(
290 !value_set.is_empty(),
291 "Transactions that output empty write set should not be included.",
292 );
293 let version = first_version + idx as u64;
294 value_set.into_iter().try_for_each(|(key, value)| {
295 Self::put(key, value, version, &mut tree_cache)
296 })?;
297 // Freezes the current cache to make all contents in the current
298 // cache immutable.
299 tree_cache.freeze();
300 }
301
302 Ok(tree_cache.into())
303 }
304
305 fn put(
306 key: HashValue, value: V, version: Version,
307 tree_cache: &mut TreeCache<R, V>,
308 ) -> Result<()> {
309 let nibble_path = NibblePath::new(key.to_vec());
310
311 // Get the root node. If this is the first operation, it would get the
312 // root node from the underlying db. Otherwise it most likely
313 // would come from `cache`.
314 let root_node_key = tree_cache.get_root_node_key();
315 let mut nibble_iter = nibble_path.nibbles();
316
317 // Start insertion from the root node.
318 let (new_root_node_key, _) = Self::insert_at(
319 root_node_key.clone(),
320 version,
321 &mut nibble_iter,
322 value,
323 tree_cache,
324 )?;
325
326 tree_cache.set_root_node_key(new_root_node_key);
327 Ok(())
328 }
329
330 /// Helper function for recursive insertion into the subtree that starts
331 /// from the current [`NodeKey`](node_type/struct.NodeKey.html). Returns
332 /// the newly inserted node. It is safe to use recursion here because
333 /// the max depth is limited by the key length which for this tree is
334 /// the length of the hash of account addresses.
335 fn insert_at(
336 node_key: NodeKey, version: Version, nibble_iter: &mut NibbleIterator,
337 value: V, tree_cache: &mut TreeCache<R, V>,
338 ) -> Result<(NodeKey, Node<V>)> {
339 let node = tree_cache.get_node(&node_key)?;
340 match node {
341 Node::Internal(internal_node) => Self::insert_at_internal_node(
342 node_key,
343 internal_node,
344 version,
345 nibble_iter,
346 value,
347 tree_cache,
348 ),
349 Node::Leaf(leaf_node) => Self::insert_at_leaf_node(
350 node_key,
351 leaf_node,
352 version,
353 nibble_iter,
354 value,
355 tree_cache,
356 ),
357 Node::Null => {
358 if node_key.nibble_path().num_nibbles() != 0 {
359 bail!(
360 "Null node exists for non-root node with node_key {:?}",
361 node_key
362 );
363 }
364 // delete the old null node if the at the same version.
365 if node_key.version() == version {
366 tree_cache.delete_node(&node_key, false /* is_leaf */);
367 }
368 Self::create_leaf_node(
369 NodeKey::new_empty_path(version),
370 &nibble_iter,
371 value,
372 tree_cache,
373 )
374 }
375 }
376 }
377
378 /// Helper function for recursive insertion into the subtree that starts
379 /// from the current `internal_node`. Returns the newly inserted node
380 /// with its [`NodeKey`](node_type/struct.NodeKey.html).
381 fn insert_at_internal_node(
382 mut node_key: NodeKey, internal_node: InternalNode, version: Version,
383 nibble_iter: &mut NibbleIterator, value: V,
384 tree_cache: &mut TreeCache<R, V>,
385 ) -> Result<(NodeKey, Node<V>)> {
386 // We always delete the existing internal node here because it will not
387 // be referenced anyway since this version.
388 tree_cache.delete_node(&node_key, false /* is_leaf */);
389
390 // Find the next node to visit following the next nibble as index.
391 let child_index = nibble_iter.next().expect("Ran out of nibbles");
392
393 // Traverse downwards from this internal node recursively to get the
394 // `node_key` of the child node at `child_index`.
395 let (_, new_child_node) = match internal_node.child(child_index) {
396 Some(child) => {
397 let child_node_key =
398 node_key.gen_child_node_key(child.version, child_index);
399 Self::insert_at(
400 child_node_key,
401 version,
402 nibble_iter,
403 value,
404 tree_cache,
405 )?
406 }
407 None => {
408 let new_child_node_key =
409 node_key.gen_child_node_key(version, child_index);
410 Self::create_leaf_node(
411 new_child_node_key,
412 nibble_iter,
413 value,
414 tree_cache,
415 )?
416 }
417 };
418
419 // Reuse the current `InternalNode` in memory to create a new internal
420 // node.
421 let mut children: Children = internal_node.into();
422 children.insert(
423 child_index,
424 Child::new(
425 new_child_node.hash(),
426 version,
427 new_child_node.is_leaf(),
428 ),
429 );
430 let new_internal_node = InternalNode::new(children);
431
432 node_key.set_version(version);
433
434 // Cache this new internal node.
435 tree_cache
436 .put_node(node_key.clone(), new_internal_node.clone().into())?;
437 Ok((node_key, new_internal_node.into()))
438 }
439
440 /// Helper function for recursive insertion into the subtree that starts
441 /// from the `existing_leaf_node`. Returns the newly inserted node with
442 /// its [`NodeKey`](node_type/struct.NodeKey.html).
443 fn insert_at_leaf_node(
444 mut node_key: NodeKey, existing_leaf_node: LeafNode<V>,
445 version: Version, nibble_iter: &mut NibbleIterator, value: V,
446 tree_cache: &mut TreeCache<R, V>,
447 ) -> Result<(NodeKey, Node<V>)> {
448 // We are on a leaf node but trying to insert another node, so we may
449 // diverge. We always delete the existing leaf node here because
450 // it will not be referenced anyway since this version.
451 tree_cache.delete_node(&node_key, true /* is_leaf */);
452
453 // 1. Make sure that the existing leaf nibble_path has the same prefix
454 // as the already visited part of the nibble iter of the
455 // incoming key and advances the existing leaf nibble iterator
456 // by the length of that prefix.
457 let mut visited_nibble_iter = nibble_iter.visited_nibbles();
458 let existing_leaf_nibble_path =
459 NibblePath::new(existing_leaf_node.account_key().to_vec());
460 let mut existing_leaf_nibble_iter = existing_leaf_nibble_path.nibbles();
461 skip_common_prefix(
462 &mut visited_nibble_iter,
463 &mut existing_leaf_nibble_iter,
464 );
465
466 // TODO(lightmark): Change this to corrupted error.
467 assert!(
468 visited_nibble_iter.is_finished(),
469 "Leaf nodes failed to share the same visited nibbles before index {}",
470 existing_leaf_nibble_iter.visited_nibbles().num_nibbles()
471 );
472
473 // 2. Determine the extra part of the common prefix that extends from
474 // the position where step 1 ends between this leaf node and the
475 // incoming key.
476 let mut existing_leaf_nibble_iter_below_internal =
477 existing_leaf_nibble_iter.remaining_nibbles();
478 let num_common_nibbles_below_internal = skip_common_prefix(
479 nibble_iter,
480 &mut existing_leaf_nibble_iter_below_internal,
481 );
482 let mut common_nibble_path =
483 nibble_iter.visited_nibbles().collect::<NibblePath>();
484
485 // 2.1. Both are finished. That means the incoming key already exists in
486 // the tree and we just need to update its value.
487 if nibble_iter.is_finished() {
488 assert!(existing_leaf_nibble_iter_below_internal.is_finished());
489 // The new leaf node will have the same nibble_path with a new
490 // version as node_key.
491 node_key.set_version(version);
492 // Create the new leaf node with the same address but the new value.
493 return Self::create_leaf_node(
494 node_key,
495 nibble_iter,
496 value,
497 tree_cache,
498 );
499 }
500
501 // 2.2. both are unfinished(They have keys with same length so it's
502 // impossible to have one finished and the other not). This
503 // means the incoming key forks at some point between the
504 // position where step 1 ends and the last nibble, inclusive. Then
505 // create a seris of internal nodes the number of which equals
506 // to the length of the extra part of the common prefix in step
507 // 2, a new leaf node for the incoming key, and update the
508 // [`NodeKey`] of existing leaf node. We create new internal nodes in a
509 // bottom-up order.
510 let existing_leaf_index = existing_leaf_nibble_iter_below_internal
511 .next()
512 .expect("Ran out of nibbles");
513 let new_leaf_index = nibble_iter.next().expect("Ran out of nibbles");
514 assert_ne!(existing_leaf_index, new_leaf_index);
515
516 let mut children = Children::new();
517 children.insert(
518 existing_leaf_index,
519 Child::new(
520 existing_leaf_node.hash(),
521 version,
522 true, /* is_leaf */
523 ),
524 );
525 node_key = NodeKey::new(version, common_nibble_path.clone());
526 tree_cache.put_node(
527 node_key.gen_child_node_key(version, existing_leaf_index),
528 existing_leaf_node.into(),
529 )?;
530
531 let (_, new_leaf_node) = Self::create_leaf_node(
532 node_key.gen_child_node_key(version, new_leaf_index),
533 nibble_iter,
534 value,
535 tree_cache,
536 )?;
537 children.insert(
538 new_leaf_index,
539 Child::new(new_leaf_node.hash(), version, true /* is_leaf */),
540 );
541
542 let internal_node = InternalNode::new(children);
543 let mut next_internal_node = internal_node.clone();
544 tree_cache.put_node(node_key.clone(), internal_node.into())?;
545
546 for _i in 0..num_common_nibbles_below_internal {
547 let nibble = common_nibble_path.pop().expect(
548 "Common nibble_path below internal node ran out of nibble",
549 );
550 node_key = NodeKey::new(version, common_nibble_path.clone());
551 let mut children = Children::new();
552 children.insert(
553 nibble,
554 Child::new(
555 next_internal_node.hash(),
556 version,
557 false, /* is_leaf */
558 ),
559 );
560 let internal_node = InternalNode::new(children);
561 next_internal_node = internal_node.clone();
562 tree_cache.put_node(node_key.clone(), internal_node.into())?;
563 }
564
565 Ok((node_key, next_internal_node.into()))
566 }
567
568 /// Helper function for creating leaf nodes. Returns the newly created leaf
569 /// node.
570 fn create_leaf_node(
571 node_key: NodeKey, nibble_iter: &NibbleIterator, value: V,
572 tree_cache: &mut TreeCache<R, V>,
573 ) -> Result<(NodeKey, Node<V>)> {
574 // Get the underlying bytes of nibble_iter which must be a key, i.e.,
575 // hashed account address with `HashValue::LENGTH` bytes.
576 let new_leaf_node = Node::new_leaf(
577 HashValue::from_slice(nibble_iter.get_nibble_path().bytes())
578 .expect("LeafNode must have full nibble path."),
579 value,
580 );
581
582 tree_cache.put_node(node_key.clone(), new_leaf_node.clone())?;
583 Ok((node_key, new_leaf_node))
584 }
585
586 /// Returns the value (if applicable) and the corresponding merkle proof.
587 pub fn get_with_proof(
588 &self, key: HashValue, version: Version,
589 ) -> Result<(Option<V>, SparseMerkleProof<V>)> {
590 // Empty tree just returns proof with no sibling hash.
591 let mut next_node_key = NodeKey::new_empty_path(version);
592 let mut siblings = vec![];
593 let nibble_path = NibblePath::new(key.to_vec());
594 let mut nibble_iter = nibble_path.nibbles();
595
596 // We limit the number of loops here deliberately to avoid potential
597 // cyclic graph bugs in the tree structure.
598 for nibble_depth in 0..=ROOT_NIBBLE_HEIGHT {
599 let next_node =
600 self.reader.get_node(&next_node_key).map_err(|err| {
601 if nibble_depth == 0 {
602 MissingRootError { version }.into()
603 } else {
604 err
605 }
606 })?;
607 match next_node {
608 Node::Internal(internal_node) => {
609 let queried_child_index = nibble_iter
610 .next()
611 .ok_or_else(|| format_err!("ran out of nibbles"))?;
612 let (child_node_key, mut siblings_in_internal) =
613 internal_node.get_child_with_siblings(
614 &next_node_key,
615 queried_child_index,
616 );
617 siblings.append(&mut siblings_in_internal);
618 next_node_key = match child_node_key {
619 Some(node_key) => node_key,
620 None => {
621 return Ok((
622 None,
623 SparseMerkleProof::new(None, {
624 siblings.reverse();
625 siblings
626 }),
627 ))
628 }
629 };
630 }
631 Node::Leaf(leaf_node) => {
632 return Ok((
633 if leaf_node.account_key() == key {
634 Some(leaf_node.value().clone())
635 } else {
636 None
637 },
638 SparseMerkleProof::new(Some(leaf_node.into()), {
639 siblings.reverse();
640 siblings
641 }),
642 ));
643 }
644 Node::Null => {
645 if nibble_depth == 0 {
646 return Ok((
647 None,
648 SparseMerkleProof::new(None, vec![]),
649 ));
650 } else {
651 bail!(
652 "Non-root null node exists with node key {:?}",
653 next_node_key
654 );
655 }
656 }
657 }
658 }
659 bail!("Jellyfish Merkle tree has cyclic graph inside.");
660 }
661
662 /// Gets the proof that shows a list of keys up to `rightmost_key_to_prove`
663 /// exist at `version`.
664 pub fn get_range_proof(
665 &self, rightmost_key_to_prove: HashValue, version: Version,
666 ) -> Result<SparseMerkleRangeProof> {
667 let (account, proof) =
668 self.get_with_proof(rightmost_key_to_prove, version)?;
669 ensure!(account.is_some(), "rightmost_key_to_prove must exist.");
670
671 let siblings = proof
672 .siblings()
673 .iter()
674 .rev()
675 .zip(rightmost_key_to_prove.iter_bits())
676 .filter_map(|(sibling, bit)| {
677 // We only need to keep the siblings on the right.
678 if !bit {
679 Some(*sibling)
680 } else {
681 None
682 }
683 })
684 .rev()
685 .collect();
686 Ok(SparseMerkleRangeProof::new(siblings))
687 }
688
689 #[cfg(test)]
690 pub fn get(&self, key: HashValue, version: Version) -> Result<Option<V>> {
691 Ok(self.get_with_proof(key, version)?.0)
692 }
693
694 pub fn get_root_hash(&self, version: Version) -> Result<HashValue> {
695 self.get_root_hash_option(version)?.ok_or_else(|| {
696 format_err!("Root node not found for version {}.", version)
697 })
698 }
699
700 pub fn get_root_hash_option(
701 &self, version: Version,
702 ) -> Result<Option<HashValue>> {
703 let root_node_key = NodeKey::new_empty_path(version);
704 Ok(self
705 .reader
706 .get_node_option(&root_node_key)?
707 .map(|root_node| root_node.hash()))
708 }
709}
710
711trait NibbleExt {
712 fn get_nibble(&self, index: usize) -> Nibble;
713 fn common_prefix_nibbles_len(&self, other: HashValue) -> usize;
714}
715
716impl NibbleExt for HashValue {
717 /// Returns the `index`-th nibble.
718 fn get_nibble(&self, index: usize) -> Nibble {
719 mirai_annotations::precondition!(index < HashValue::LENGTH);
720 Nibble::from(
721 if index % 2 == 0 {
722 self[index / 2] >> 4
723 } else {
724 self[index / 2] & 0x0F
725 },
726 )
727 }
728
729 /// Returns the length of common prefix of `self` and `other` in nibbles.
730 fn common_prefix_nibbles_len(&self, other: HashValue) -> usize {
731 self.common_prefix_bits_len(other) / 4
732 }
733}
734
735#[cfg(test)]
736mod test {
737 use super::NibbleExt;
738 use diem_crypto::hash::{HashValue, TestOnlyHash};
739 use diem_nibble::Nibble;
740
741 #[test]
742 fn test_common_prefix_nibbles_len() {
743 {
744 let hash1 = b"hello".test_only_hash();
745 let hash2 = b"HELLO".test_only_hash();
746 assert_eq!(hash1[0], 0b0011_0011);
747 assert_eq!(hash2[0], 0b1011_1000);
748 assert_eq!(hash1.common_prefix_nibbles_len(hash2), 0);
749 }
750 {
751 let hash1 = b"hello".test_only_hash();
752 let hash2 = b"world".test_only_hash();
753 assert_eq!(hash1[0], 0b0011_0011);
754 assert_eq!(hash2[0], 0b0100_0010);
755 assert_eq!(hash1.common_prefix_nibbles_len(hash2), 0);
756 }
757 {
758 let hash1 = b"hello".test_only_hash();
759 let hash2 = b"100011001000".test_only_hash();
760 assert_eq!(hash1[0], 0b0011_0011);
761 assert_eq!(hash2[0], 0b0011_0011);
762 assert_eq!(hash1[1], 0b0011_1000);
763 assert_eq!(hash2[1], 0b0010_0010);
764 assert_eq!(hash1.common_prefix_nibbles_len(hash2), 2);
765 }
766 {
767 let hash1 = b"hello".test_only_hash();
768 let hash2 = b"hello".test_only_hash();
769 assert_eq!(
770 hash1.common_prefix_nibbles_len(hash2),
771 HashValue::LENGTH * 2
772 );
773 }
774 }
775
776 #[test]
777 fn test_get_nibble() {
778 let hash = b"hello".test_only_hash();
779 assert_eq!(hash.get_nibble(0), Nibble::from(3));
780 assert_eq!(hash.get_nibble(1), Nibble::from(3));
781 assert_eq!(hash.get_nibble(2), Nibble::from(3));
782 assert_eq!(hash.get_nibble(3), Nibble::from(8));
783 assert_eq!(hash.get_nibble(62), Nibble::from(9));
784 assert_eq!(hash.get_nibble(63), Nibble::from(2));
785 }
786}