From a36ee7f947cbd762198c19412626c0ba0d3d6c8f Mon Sep 17 00:00:00 2001 From: suoyuan Date: Fri, 4 Sep 2020 19:01:49 +0800 Subject: [PATCH] [jellyfish_tree] update test for jellyfish tree (#1230) * [jellyfish_tree] update test for jellyfish tree * [jellyfish_tree] fix clippy check --- core/forkable-jellyfish-merkle/src/blob.rs | 4 + .../src/iterator/mod.rs | 27 +- .../src/jellyfish_merkle_test.rs | 1126 +++++++++-------- core/forkable-jellyfish-merkle/src/lib.rs | 14 +- .../src/node_type/node_type_test.rs | 2 - .../src/tree_cache/mod.rs | 10 + state/state-tree/src/state_tree.rs | 3 +- 7 files changed, 614 insertions(+), 572 deletions(-) diff --git a/core/forkable-jellyfish-merkle/src/blob.rs b/core/forkable-jellyfish-merkle/src/blob.rs index cb2b7ef329..c9ba5d2650 100644 --- a/core/forkable-jellyfish-merkle/src/blob.rs +++ b/core/forkable-jellyfish-merkle/src/blob.rs @@ -1,10 +1,14 @@ // Copyright (c) The Starcoin Core Contributors // SPDX-License-Identifier: Apache-2.0 +#[cfg(any(test, feature = "fuzzing"))] +use proptest_derive::Arbitrary; use serde::{Deserialize, Serialize}; use starcoin_crypto::hash::*; use std::fmt; + #[derive(Clone, Eq, PartialEq, Serialize, Deserialize, CryptoHasher, CryptoHash)] +#[cfg_attr(any(test, feature = "fuzzing"), derive(Arbitrary))] pub struct Blob { blob: Vec, } diff --git a/core/forkable-jellyfish-merkle/src/iterator/mod.rs b/core/forkable-jellyfish-merkle/src/iterator/mod.rs index b3fdd2461d..f60f3e7cff 100644 --- a/core/forkable-jellyfish-merkle/src/iterator/mod.rs +++ b/core/forkable-jellyfish-merkle/src/iterator/mod.rs @@ -21,7 +21,6 @@ use crate::{ }; use anyhow::{format_err, Result}; use starcoin_crypto::HashValue; -use std::sync::Arc; type Version = HashValue; /// `NodeVisitInfo` keeps track of the status of an internal node during the iteration process. It @@ -95,9 +94,9 @@ impl NodeVisitInfo { } /// The `JellyfishMerkleIterator` implementation. -pub struct JellyfishMerkleIterator { +pub struct JellyfishMerkleIterator<'a, R: 'a + TreeReader> { /// The storage engine from which we can read nodes using node keys. - reader: Arc, + reader: &'a R, /// The version of the tree this iterator is running on. state_root_hash: Version, @@ -111,14 +110,14 @@ pub struct JellyfishMerkleIterator { done: bool, } -impl JellyfishMerkleIterator +impl<'a, R> JellyfishMerkleIterator<'a, R> where - R: TreeReader, + R: 'a + TreeReader, { /// Constructs a new iterator. This puts the internal state in the correct position, so the /// following `next` call will yield the smallest key that is greater or equal to /// `starting_key`. - pub fn new(reader: Arc, state_root_hash: Version, starting_key: HashValue) -> Result { + pub fn new(reader: &'a R, state_root_hash: Version, starting_key: HashValue) -> Result { let mut parent_stack = vec![]; let mut done = false; @@ -195,11 +194,23 @@ where } } } + + #[cfg(test)] + pub fn print(&self) -> Result<()> { + let nodes = &self.parent_stack; + for node in nodes { + println!("internal node key: {:?}", node.node_key.short_str()); + if let Ok(Node::Internal(internal)) = self.reader.get_node(&node.node_key) { + println!("child: {:?}", internal.all_child()); + } + } + Ok(()) + } } -impl Iterator for JellyfishMerkleIterator +impl<'a, R> Iterator for JellyfishMerkleIterator<'a, R> where - R: TreeReader, + R: 'a + TreeReader, { type Item = Result<(HashValue, Blob)>; diff --git a/core/forkable-jellyfish-merkle/src/jellyfish_merkle_test.rs b/core/forkable-jellyfish-merkle/src/jellyfish_merkle_test.rs index d8acc8516c..6282aaa62e 100644 --- a/core/forkable-jellyfish-merkle/src/jellyfish_merkle_test.rs +++ b/core/forkable-jellyfish-merkle/src/jellyfish_merkle_test.rs @@ -8,9 +8,16 @@ use super::*; use crate::nibble::Nibble; use mock_tree_store::MockTreeStore; +use crate::node_type::SparseMerkleInternalNode; +use proptest::{ + collection::{btree_map, hash_map, vec}, + prelude::*, +}; use rand::{rngs::StdRng, Rng, SeedableRng}; -use starcoin_crypto::hash::*; +use starcoin_crypto::hash::{HashValue, *}; use std::collections::HashMap; +use std::ops::Bound; +use test_helper::{init_mock_db, plus_one}; fn update_nibble(original_key: &HashValue, n: usize, nibble: u8) -> HashValue { assert!(nibble < 16); @@ -155,8 +162,10 @@ fn test_insert_at_leaf_with_multiple_internals_created() { assert_eq!(tree.get(_root0_hash, key1,).unwrap().unwrap(), value1); assert!(tree.get(_root0_hash, key2,).unwrap().is_none()); assert_eq!(tree.get(_root1_hash, key2,).unwrap().unwrap(), value2); + assert_eq!(tree.get(_root1_hash, key1,).unwrap().unwrap(), value1); assert_eq!(db.num_nodes(), 4); + tree.print_tree(_root1_hash, key1).unwrap(); let leaf1 = Node::new_leaf(key1, value1); let leaf2 = Node::new_leaf(key2, value2.clone()); @@ -195,6 +204,7 @@ fn test_insert_at_leaf_with_multiple_internals_created() { value2_update ); + tree.print_tree(_root2_hash, key1).unwrap(); // Get # of nodes. assert_eq!(db.num_nodes(), 7); @@ -203,299 +213,311 @@ fn test_insert_at_leaf_with_multiple_internals_created() { db.purge_stale_nodes(_root1_hash).unwrap(); assert_eq!(db.num_nodes(), 7); db.purge_stale_nodes(_root2_hash).unwrap(); + tree.print_tree(_root2_hash, key1).unwrap(); assert_eq!(db.num_nodes(), 4); } -// #[test] -// fn test_batch_insertion() { -// // ```text -// // internal(root) -// // / \ -// // internal 2 <- nibble 0 -// // / | \ -// // internal 3 4 <- nibble 1 -// // | -// // internal <- nibble 2 -// // / \ -// // internal 6 <- nibble 3 -// // | -// // internal <- nibble 4 -// // / \ -// // 1 5 <- nibble 5 -// // -// // Total: 12 nodes -// // ``` -// let key1 = HashValue::new([0x00u8; HashValue::LENGTH]); -// let value1 = AccountStateBlob::from(vec![1u8]); -// -// let key2 = update_nibble(&key1, 0, 2); -// let value2 = AccountStateBlob::from(vec![2u8]); -// let value2_update = AccountStateBlob::from(vec![22u8]); -// -// let key3 = update_nibble(&key1, 1, 3); -// let value3 = AccountStateBlob::from(vec![3u8]); -// -// let key4 = update_nibble(&key1, 1, 4); -// let value4 = AccountStateBlob::from(vec![4u8]); -// -// let key5 = update_nibble(&key1, 5, 5); -// let value5 = AccountStateBlob::from(vec![5u8]); -// -// let key6 = update_nibble(&key1, 3, 6); -// let value6 = AccountStateBlob::from(vec![6u8]); -// -// let batches = vec![ -// vec![(key1, value1)], -// vec![(key2, value2)], -// vec![(key3, value3)], -// vec![(key4, value4)], -// vec![(key5, value5)], -// vec![(key6, value6)], -// vec![(key2, value2_update)], -// ]; -// let one_batch = batches.iter().flatten().cloned().collect::>(); -// -// let mut to_verify = one_batch.clone(); -// // key2 was updated so we remove it. -// to_verify.remove(1); -// let verify_fn = |tree: &JellyfishMerkleTree, version: Version| { -// to_verify -// .iter() -// .for_each(|(k, v)| assert_eq!(tree.get(*k, version).unwrap().unwrap(), *v)) -// }; -// -// // Insert as one batch. -// { -// let db = MockTreeStore::default(); -// let tree = JellyfishMerkleTree::new(&db); -// -// let (_root, batch) = tree.put_blob_set(one_batch, 0 /* version */).unwrap(); -// db.write_tree_update_batch(batch).unwrap(); -// verify_fn(&tree, 0); -// -// // get # of nodes -// assert_eq!(db.num_nodes(), 12); -// } -// -// // Insert in multiple batches. -// { -// let db = MockTreeStore::default(); -// let tree = JellyfishMerkleTree::new(&db); -// -// let (_roots, batch) = tree.put_blob_sets(batches, 0 /* first_version */).unwrap(); -// db.write_tree_update_batch(batch).unwrap(); -// verify_fn(&tree, 6); -// -// // get # of nodes -// assert_eq!(db.num_nodes(), 26 /* 1 + 3 + 4 + 3 + 8 + 5 + 2 */); -// -// // Purge retired nodes('p' means purged and 'a' means added). -// // The initial state of the tree at version 0 -// // ```test -// // 1(root) -// // ``` -// db.purge_stale_nodes(1).unwrap(); -// // ```text -// // 1 (p) internal(a) -// // -> / \ -// // 1(a) 2(a) -// // add 3, prune 1 -// // ``` -// assert_eq!(db.num_nodes(), 25); -// db.purge_stale_nodes(2).unwrap(); -// // ```text -// // internal(p) internal(a) -// // / \ / \ -// // 1(p) 2 -> internal(a) 2 -// // / \ -// // 1(a) 3(a) -// // add 4, prune 2 -// // ``` -// assert_eq!(db.num_nodes(), 23); -// db.purge_stale_nodes(3).unwrap(); -// // ```text -// // internal(p) internal(a) -// // / \ / \ -// // internal(p) 2 -> internal(a) 2 -// // / \ / | \ -// // 1 3 1 3 4(a) -// // add 3, prune 2 -// // ``` -// assert_eq!(db.num_nodes(), 21); -// db.purge_stale_nodes(4).unwrap(); -// // ```text -// // internal(p) internal(a) -// // / \ / \ -// // internal(p) 2 internal(a) 2 -// // / | \ / | \ -// // 1(p) 3 4 -> internal(a) 3 4 -// // | -// // internal(a) -// // | -// // internal(a) -// // | -// // internal(a) -// // / \ -// // 1(a) 5(a) -// // add 8, prune 3 -// // ``` -// assert_eq!(db.num_nodes(), 18); -// db.purge_stale_nodes(5).unwrap(); -// // ```text -// // internal(p) internal(a) -// // / \ / \ -// // internal(p) 2 internal(a) 2 -// // / | \ / | \ -// // internal(p) 3 4 internal(a) 3 4 -// // | | -// // internal(p) -> internal(a) -// // | / \ -// // internal internal 6(a) -// // | | -// // internal internal -// // / \ / \ -// // 1 5 1 5 -// // add 5, prune 4 -// // ``` -// assert_eq!(db.num_nodes(), 14); -// db.purge_stale_nodes(6).unwrap(); -// // ```text -// // internal(p) internal(a) -// // / \ / \ -// // internal 2(p) internal 2(a) -// // / | \ / | \ -// // internal 3 4 internal 3 4 -// // | | -// // internal -> internal -// // / \ / \ -// // internal 6 internal 6 -// // | | -// // internal internal -// // / \ / \ -// // 1 5 1 5 -// // add 2, prune 2 -// // ``` -// assert_eq!(db.num_nodes(), 12); -// verify_fn(&tree, 6); -// } -// } -// -// #[test] -// fn test_non_existence() { -// let db = MockTreeStore::default(); -// let tree = JellyfishMerkleTree::new(&db); -// // ```text -// // internal(root) -// // / \ -// // internal 2 -// // | -// // internal -// // / \ -// // 1 3 -// // Total: 7 nodes -// // ``` -// let key1 = HashValue::new([0x00u8; HashValue::LENGTH]); -// let value1 = AccountStateBlob::from(vec![1u8]); -// -// let key2 = update_nibble(&key1, 0, 15); -// let value2 = AccountStateBlob::from(vec![2u8]); -// -// let key3 = update_nibble(&key1, 2, 3); -// let value3 = AccountStateBlob::from(vec![3u8]); -// -// let (root, batch) = tree -// .put_blob_set( -// vec![ -// (key1, value1.clone()), -// (key2, value2.clone()), -// (key3, value3.clone()), -// ], -// 0, /* version */ -// ) -// .unwrap(); -// db.write_tree_update_batch(batch).unwrap(); -// assert_eq!(tree.get(key1, 0).unwrap().unwrap(), value1); -// assert_eq!(tree.get(key2, 0).unwrap().unwrap(), value2); -// assert_eq!(tree.get(key3, 0).unwrap().unwrap(), value3); -// // get # of nodes -// assert_eq!(db.num_nodes(), 6); -// -// // test non-existing nodes. -// // 1. Non-existing node at root node -// { -// let non_existing_key = update_nibble(&key1, 0, 1); -// let (value, proof) = tree.get_with_proof(non_existing_key, 0).unwrap(); -// assert_eq!(value, None); -// assert!(proof.verify(root, non_existing_key, None).is_ok()); -// } -// // 2. Non-existing node at non-root internal node -// { -// let non_existing_key = update_nibble(&key1, 1, 15); -// let (value, proof) = tree.get_with_proof(non_existing_key, 0).unwrap(); -// assert_eq!(value, None); -// assert!(proof.verify(root, non_existing_key, None).is_ok()); -// } -// // 3. Non-existing node at leaf node -// { -// let non_existing_key = update_nibble(&key1, 2, 4); -// let (value, proof) = tree.get_with_proof(non_existing_key, 0).unwrap(); -// assert_eq!(value, None); -// assert!(proof.verify(root, non_existing_key, None).is_ok()); -// } -// } -// -// #[test] -// fn test_put_blob_sets() { -// let mut keys = vec![]; -// let mut values = vec![]; -// let total_updates = 20; -// for _i in 0..total_updates { -// keys.push(HashValue::random()); -// values.push(AccountStateBlob::from(HashValue::random().to_vec())); -// } -// -// let mut root_hashes_one_by_one = vec![]; -// let mut batch_one_by_one = TreeUpdateBatch::default(); -// { -// let mut iter = keys.clone().into_iter().zip(values.clone().into_iter()); -// let db = MockTreeStore::default(); -// let tree = JellyfishMerkleTree::new(&db); -// for version in 0..10 { -// let mut keyed_blob_set = vec![]; -// for _ in 0..total_updates / 10 { -// keyed_blob_set.push(iter.next().unwrap()); -// } -// let (root, batch) = tree -// .put_blob_set(keyed_blob_set, version as Version) -// .unwrap(); -// db.write_tree_update_batch(batch.clone()).unwrap(); -// root_hashes_one_by_one.push(root); -// batch_one_by_one.node_batch.extend(batch.node_batch); -// batch_one_by_one -// .stale_node_index_batch -// .extend(batch.stale_node_index_batch); -// batch_one_by_one.num_new_leaves += batch.num_new_leaves; -// batch_one_by_one.num_stale_leaves += batch.num_stale_leaves; -// } -// } -// { -// let mut iter = keys.into_iter().zip(values.into_iter()); -// let db = MockTreeStore::default(); -// let tree = JellyfishMerkleTree::new(&db); -// let mut blob_sets = vec![]; -// for _ in 0..10 { -// let mut keyed_blob_set = vec![]; -// for _ in 0..total_updates / 10 { -// keyed_blob_set.push(iter.next().unwrap()); -// } -// blob_sets.push(keyed_blob_set); -// } -// let (root_hashes, batch) = tree.put_blob_sets(blob_sets, 0 /* version */).unwrap(); -// assert_eq!(root_hashes, root_hashes_one_by_one); -// assert_eq!(batch, batch_one_by_one); -// } -// } -// +#[test] +fn test_batch_insertion() { + // ```text + // internal(root) + // / \ + // internal 2 <- nibble 0 + // / | \ + // internal 3 4 <- nibble 1 + // | + // internal <- nibble 2 + // / \ + // internal 6 <- nibble 3 + // | + // internal <- nibble 4 + // / \ + // 1 5 <- nibble 5 + // + // Total: 12 nodes + // ``` + let key1 = HashValue::new([0x00u8; HashValue::LENGTH]); + let value1 = Blob::from(vec![1u8]); + + let key2 = update_nibble(&key1, 0, 2); + let value2 = Blob::from(vec![2u8]); + let value2_update = Blob::from(vec![22u8]); + + let key3 = update_nibble(&key1, 1, 3); + let value3 = Blob::from(vec![3u8]); + + let key4 = update_nibble(&key1, 1, 4); + let value4 = Blob::from(vec![4u8]); + + let key5 = update_nibble(&key1, 5, 5); + let value5 = Blob::from(vec![5u8]); + + let key6 = update_nibble(&key1, 3, 6); + let value6 = Blob::from(vec![6u8]); + + let batches = vec![ + vec![(key1, value1)], + vec![(key2, value2)], + vec![(key3, value3)], + vec![(key4, value4)], + vec![(key5, value5)], + vec![(key6, value6)], + vec![(key2, value2_update)], + ]; + let one_batch = batches.iter().flatten().cloned().collect::>(); + + let mut to_verify = one_batch.clone(); + // key2 was updated so we remove it. + to_verify.remove(1); + let verify_fn = |tree: &JellyfishMerkleTree, root: HashValue| { + to_verify + .iter() + .for_each(|(k, v)| assert_eq!(tree.get(root, *k).unwrap().unwrap(), *v)) + }; + + // Insert as one batch. + { + let db = MockTreeStore::default(); + let tree = JellyfishMerkleTree::new(&db); + + let (root, batch) = tree.put_blob_set(None, one_batch).unwrap(); + db.write_tree_update_batch(batch).unwrap(); + verify_fn(&tree, root); + + // get # of nodes + assert_eq!(db.num_nodes(), 12); + tree.print_tree(root, key1).unwrap(); + } + + // Insert in multiple batches. + { + let db = MockTreeStore::default(); + let tree = JellyfishMerkleTree::new(&db); + let mut batches2 = vec![]; + + for (_idx, sub_vec) in batches.iter().enumerate() { + for x in sub_vec { + batches2.push(vec![(x.0, Some(x.1.clone()))]); + } + } + let (mut roots, batch) = tree.puts(None, batches2).unwrap(); + db.write_tree_update_batch(batch).unwrap(); + let root_hash = roots.pop().unwrap(); + verify_fn(&tree, root_hash); + + // get # of nodes + assert_eq!(db.num_nodes(), 23 /* 1 + 2 + 3 + 3 + 7 + 5 + 2 */); + tree.print_tree(root_hash, key1).unwrap(); + + // Purge retired nodes('p' means purged and 'a' means added). + // The initial state of the tree at version 0 + // ```test + // 1(root) + // ``` + db.purge_stale_nodes(key1).unwrap(); + // ```text + // 1 (p) internal(a) + // -> / \ + // 1(a) 2(a) + // add 3, prune 1 + // ``` + assert_eq!(db.num_nodes(), 23); + db.purge_stale_nodes(key2).unwrap(); + // ```text + // internal(p) internal(a) + // / \ / \ + // 1(p) 2 -> internal(a) 2 + // / \ + // 1(a) 3(a) + // add 4, prune 2 + // ``` + assert_eq!(db.num_nodes(), 23); + db.purge_stale_nodes(key3).unwrap(); + // ```text + // internal(p) internal(a) + // / \ / \ + // internal(p) 2 -> internal(a) 2 + // / \ / | \ + // 1 3 1 3 4(a) + // add 3, prune 2 + // ``` + assert_eq!(db.num_nodes(), 23); + db.purge_stale_nodes(key4).unwrap(); + // ```text + // internal(p) internal(a) + // / \ / \ + // internal(p) 2 internal(a) 2 + // / | \ / | \ + // 1(p) 3 4 -> internal(a) 3 4 + // | + // internal(a) + // | + // internal(a) + // | + // internal(a) + // / \ + // 1(a) 5(a) + // add 8, prune 3 + // ``` + assert_eq!(db.num_nodes(), 23); + db.purge_stale_nodes(key5).unwrap(); + // ```text + // internal(p) internal(a) + // / \ / \ + // internal(p) 2 internal(a) 2 + // / | \ / | \ + // internal(p) 3 4 internal(a) 3 4 + // | | + // internal(p) -> internal(a) + // | / \ + // internal internal 6(a) + // | | + // internal internal + // / \ / \ + // 1 5 1 5 + // add 5, prune 4 + // ``` + assert_eq!(db.num_nodes(), 23); + db.purge_stale_nodes(key6).unwrap(); + // ```text + // internal(p) internal(a) + // / \ / \ + // internal 2(p) internal 2(a) + // / | \ / | \ + // internal 3 4 internal 3 4 + // | | + // internal -> internal + // / \ / \ + // internal 6 internal 6 + // | | + // internal internal + // / \ / \ + // 1 5 1 5 + // add 2, prune 2 + // ``` + assert_eq!(db.num_nodes(), 23); + verify_fn(&tree, root_hash); + } +} + +#[test] +fn test_non_existence() { + let db = MockTreeStore::default(); + let tree = JellyfishMerkleTree::new(&db); + // ```text + // internal(root) + // / \ + // internal 2 + // | + // internal + // / \ + // 1 3 + // Total: 7 nodes + // ``` + let key1 = HashValue::new([0x00u8; HashValue::LENGTH]); + let value1 = Blob::from(vec![1u8]); + + let key2 = update_nibble(&key1, 0, 15); + let value2 = Blob::from(vec![2u8]); + + let key3 = update_nibble(&key1, 2, 3); + let value3 = Blob::from(vec![3u8]); + + let (root, batch) = tree + .put_blob_set( + None, + vec![ + (key1, value1.clone()), + (key2, value2.clone()), + (key3, value3.clone()), + ], + ) + .unwrap(); + db.write_tree_update_batch(batch).unwrap(); + assert_eq!(tree.get(root, key1).unwrap().unwrap(), value1); + assert_eq!(tree.get(root, key2).unwrap().unwrap(), value2); + assert_eq!(tree.get(root, key3).unwrap().unwrap(), value3); + // get # of nodes + assert_eq!(db.num_nodes(), 6); + + // test non-existing nodes. + // 1. Non-existing node at root node + { + let non_existing_key = update_nibble(&key1, 0, 1); + let (value, proof) = tree.get_with_proof(root, non_existing_key).unwrap(); + assert_eq!(value, None); + assert!(proof.verify(root, non_existing_key, None).is_ok()); + } + // 2. Non-existing node at non-root internal node + { + let non_existing_key = update_nibble(&key1, 1, 15); + let (value, proof) = tree.get_with_proof(root, non_existing_key).unwrap(); + assert_eq!(value, None); + assert!(proof.verify(root, non_existing_key, None).is_ok()); + } + // 3. Non-existing node at leaf node + { + let non_existing_key = update_nibble(&key1, 2, 4); + let (value, proof) = tree.get_with_proof(root, non_existing_key).unwrap(); + assert_eq!(value, None); + assert!(proof.verify(root, non_existing_key, None).is_ok()); + } +} + +#[test] +fn test_put_blob_sets() { + let mut keys = vec![]; + let mut values = vec![]; + let total_updates = 20; + for _i in 0..total_updates { + keys.push(HashValue::random()); + values.push(Blob::from(HashValue::random().to_vec())); + } + + let mut root_hashes_one_by_one = vec![]; + let mut batch_one_by_one = TreeUpdateBatch::default(); + { + let mut iter = keys.clone().into_iter().zip(values.clone().into_iter()); + let db = MockTreeStore::default(); + let tree = JellyfishMerkleTree::new(&db); + + let mut temp_root = None; + for _version in 0..10 { + let mut keyed_blob_set = vec![]; + for _ in 0..2 { + keyed_blob_set.push(iter.next().unwrap()); + } + let (root, batch) = tree.put_blob_set(temp_root, keyed_blob_set).unwrap(); + db.write_tree_update_batch(batch.clone()).unwrap(); + temp_root = Some(root); + root_hashes_one_by_one.push(root); + batch_one_by_one.node_batch.extend(batch.node_batch); + batch_one_by_one + .stale_node_index_batch + .extend(batch.stale_node_index_batch); + batch_one_by_one.num_new_leaves += batch.num_new_leaves; + batch_one_by_one.num_stale_leaves += batch.num_stale_leaves; + } + } + { + let mut iter = keys.into_iter().zip(values.into_iter()); + let db = MockTreeStore::default(); + let tree = JellyfishMerkleTree::new(&db); + let mut blob_sets = vec![]; + for _ in 0..10 { + let mut keyed_blob_set = vec![]; + for _ in 0..2 { + let val = iter.next().unwrap(); + keyed_blob_set.push((val.0, Some(val.1))); + } + blob_sets.push(keyed_blob_set); + } + let (root_hashes, batch) = tree.puts(None, blob_sets).unwrap(); + assert_eq!(root_hashes, root_hashes_one_by_one); + assert_eq!(batch, batch_one_by_one); + } +} + fn many_keys_get_proof_and_verify_tree_root(seed: &[u8], num_keys: usize) { assert!(seed.len() < 32); let mut actual_seed = [0u8; 32]; @@ -590,266 +612,262 @@ fn test_1000_versions() { many_versions_get_proof_and_verify_tree_root(seed, 1000); } -// proptest! { -// #![proptest_config(ProptestConfig::with_cases(10))] -// -// #[test] -// fn test_get_with_proof1( -// (existent_kvs, nonexistent_keys) in hash_map( -// any::(), -// any::(), -// 1..1000, -// ) -// .prop_flat_map(|kvs| { -// let kvs_clone = kvs.clone(); -// ( -// Just(kvs), -// vec( -// any::().prop_filter( -// "Make sure these keys do not exist in the tree.", -// move |key| !kvs_clone.contains_key(key), -// ), -// 100, -// ), -// ) -// }) -// ) { -// let (db, version) = init_mock_db(&existent_kvs); -// let tree = JellyfishMerkleTree::new(&db); -// -// test_existent_keys_impl(&tree, version, &existent_kvs); -// test_nonexistent_keys_impl(&tree, version, &nonexistent_keys); -// } -// -// #[test] -// fn test_get_with_proof2( -// key1 in any::() -// .prop_filter( -// "Can't be 0xffffff...", -// |key| *key != HashValue::new([0xff; HashValue::LENGTH]), -// ), -// accounts in vec(any::(), 2), -// ) { -// let key2 = plus_one(key1); -// -// let mut kvs = HashMap::new(); -// kvs.insert(key1, accounts[0].clone()); -// kvs.insert(key2, accounts[1].clone()); -// -// let (db, version) = init_mock_db(&kvs); -// let tree = JellyfishMerkleTree::new(&db); -// -// test_existent_keys_impl(&tree, version, &kvs); -// } -// -// #[test] -// fn test_get_range_proof( -// (btree, n) in btree_map(any::(), any::(), 1..1000) -// .prop_flat_map(|btree| { -// let len = btree.len(); -// (Just(btree), 0..len) -// }) -// ) { -// let (db, version) = init_mock_db(&btree.clone().into_iter().collect()); -// let tree = JellyfishMerkleTree::new(&db); -// -// let nth_key = *btree.keys().nth(n).unwrap(); -// let proof = tree.get_range_proof(nth_key, version).unwrap(); -// verify_range_proof( -// tree.get_root_hash(version).unwrap(), -// btree.into_iter().take(n + 1).collect(), -// proof, -// ); -// } -// } -// -// fn test_existent_keys_impl<'a>( -// tree: &JellyfishMerkleTree<'a, MockTreeStore>, -// version: Version, -// existent_kvs: &HashMap, -// ) { -// let root_hash = tree.get_root_hash(version).unwrap(); -// -// for (key, value) in existent_kvs { -// let (account, proof) = tree.get_with_proof(*key, version).unwrap(); -// assert!(proof.verify(root_hash, *key, account.as_ref()).is_ok()); -// assert_eq!(account.unwrap(), *value); -// } -// } -// -// fn test_nonexistent_keys_impl<'a>( -// tree: &JellyfishMerkleTree<'a, MockTreeStore>, -// version: Version, -// nonexistent_keys: &[HashValue], -// ) { -// let root_hash = tree.get_root_hash(version).unwrap(); -// -// for key in nonexistent_keys { -// let (account, proof) = tree.get_with_proof(*key, version).unwrap(); -// assert!(proof.verify(root_hash, *key, account.as_ref()).is_ok()); -// assert!(account.is_none()); -// } -// } -// -// /// Checks if we can construct the expected root hash using the entries in the btree and the proof. -// fn verify_range_proof( -// expected_root_hash: HashValue, -// btree: BTreeMap, -// proof: SparseMerkleRangeProof, -// ) { -// // For example, given the following sparse Merkle tree: -// // -// // root -// // / \ -// // / \ -// // / \ -// // o o -// // / \ / \ -// // a o o h -// // / \ / \ -// // o d e X -// // / \ / \ -// // b c f g -// // -// // we transform the keys as follows: -// // a => 00, -// // b => 0100, -// // c => 0101, -// // d => 011, -// // e => 100, -// // X => 101 -// // h => 11 -// // -// // Basically, the suffixes that doesn't affect the common prefix of adjacent leaves are -// // discarded. In this example, we assume `btree` has the keys `a` to `e` and the proof has `X` -// // and `h` in the siblings. -// -// // Now we want to construct a set of key-value pairs that covers the entire set of leaves. For -// // `a` to `e` this is simple -- we just insert them directly into this set. For the rest of the -// // leaves, they are represented by the siblings, so we just make up some keys that make sense. -// // For example, for `X` we just use 101000... (more zeros omitted), because that is one key -// // that would cause `X` to end up in the above position. -// let mut btree1 = BTreeMap::new(); -// for (key, blob) in &btree { -// let leaf = LeafNode::new(*key, blob.clone()); -// btree1.insert(*key, leaf.hash()); -// } -// // Using the above example, `last_proven_key` is `e`. We look at the path from root to `e`. -// // For each 0-bit, there should be a sibling in the proof. And we use the path from root to -// // this position, plus a `1` as the key. -// let last_proven_key = *btree -// .keys() -// .last() -// .expect("We are proving at least one key."); -// for (i, sibling) in last_proven_key -// .iter_bits() -// .enumerate() -// .filter_map(|(i, bit)| if !bit { Some(i) } else { None }) -// .zip(proof.right_siblings().iter().rev()) -// { -// // This means the `i`-th bit is zero. We take `i` bits from `last_proven_key` and append a -// // one to make up the key for this sibling. -// let mut buf: Vec<_> = last_proven_key.iter_bits().take(i).collect(); -// buf.push(true); -// // The rest doesn't matter, because they don't affect the position of the node. We just -// // add zeros. -// buf.resize(HashValue::LENGTH_IN_BITS, false); -// let key = HashValue::from_bit_iter(buf.into_iter()).unwrap(); -// btree1.insert(key, *sibling); -// } -// -// // Now we do the transformation (removing the suffixes) described above. -// let mut kvs = vec![]; -// for (key, value) in &btree1 { -// // The length of the common prefix of the previous key and the current key. -// let prev_common_prefix_len = -// prev_key(&btree1, key).map(|pkey| pkey.common_prefix_bits_len(*key)); -// // The length of the common prefix of the next key and the current key. -// let next_common_prefix_len = -// next_key(&btree1, key).map(|nkey| nkey.common_prefix_bits_len(*key)); -// -// // We take the longest common prefix of the current key and its neighbors. That's how much -// // we need to keep. -// let len = match (prev_common_prefix_len, next_common_prefix_len) { -// (Some(plen), Some(nlen)) => std::cmp::max(plen, nlen), -// (Some(plen), None) => plen, -// (None, Some(nlen)) => nlen, -// (None, None) => 0, -// }; -// let transformed_key: Vec<_> = key.iter_bits().take(len + 1).collect(); -// kvs.push((transformed_key, *value)); -// } -// -// assert_eq!(compute_root_hash(kvs), expected_root_hash); -// } -// -// /// Computes the root hash of a sparse Merkle tree. `kvs` consists of the entire set of key-value -// /// pairs stored in the tree. -// fn compute_root_hash(kvs: Vec<(Vec, HashValue)>) -> HashValue { -// let mut kv_ref = vec![]; -// for (key, value) in &kvs { -// kv_ref.push((&key[..], *value)); -// } -// compute_root_hash_impl(kv_ref) -// } -// -// fn compute_root_hash_impl(kvs: Vec<(&[bool], HashValue)>) -> HashValue { -// assert!(!kvs.is_empty()); -// -// // If there is only one entry, it is the root. -// if kvs.len() == 1 { -// return kvs[0].1; -// } -// -// // Otherwise the tree has more than one leaves, which means we can find which ones are in the -// // left subtree and which ones are in the right subtree. So we find the first key that starts -// // with a 1-bit. -// let left_hash; -// let right_hash; -// match kvs.iter().position(|(key, _value)| key[0]) { -// Some(0) => { -// // Every key starts with a 1-bit, i.e., they are all in the right subtree. -// left_hash = *SPARSE_MERKLE_PLACEHOLDER_HASH; -// right_hash = compute_root_hash_impl(reduce(&kvs)); -// } -// Some(index) => { -// // Both left subtree and right subtree have some keys. -// left_hash = compute_root_hash_impl(reduce(&kvs[0..index])); -// right_hash = compute_root_hash_impl(reduce(&kvs[index..])); -// } -// None => { -// // Every key starts with a 0-bit, i.e., they are all in the left subtree. -// left_hash = compute_root_hash_impl(reduce(&kvs)); -// right_hash = *SPARSE_MERKLE_PLACEHOLDER_HASH; -// } -// } -// -// SparseMerkleInternalNode::new(left_hash, right_hash).hash() -// } -// -// /// Reduces the problem by removing the first bit of every key. -// fn reduce<'a>(kvs: &'a [(&[bool], HashValue)]) -> Vec<(&'a [bool], HashValue)> { -// kvs.iter().map(|(key, value)| (&key[1..], *value)).collect() -// } -// -// /// Returns the key immediately before `key` in `btree`. -// fn prev_key(btree: &BTreeMap, key: &K) -> Option -// where -// K: Clone + Ord, -// { -// btree -// .range((Bound::Unbounded, Bound::Excluded(key))) -// .next_back() -// .map(|(k, _v)| k.clone()) -// } -// -// /// Returns the key immediately after `key` in `btree`. -// fn next_key(btree: &BTreeMap, key: &K) -> Option -// where -// K: Clone + Ord, -// { -// btree -// .range((Bound::Excluded(key), Bound::Unbounded)) -// .next() -// .map(|(k, _v)| k.clone()) -// } +proptest! { + #![proptest_config(ProptestConfig::with_cases(10))] + + #[test] + fn test_get_with_proof1( + (existent_kvs, nonexistent_keys) in hash_map( + any::(), + any::(), + 1..1000, + ) + .prop_flat_map(|kvs| { + let kvs_clone = kvs.clone(); + ( + Just(kvs), + vec( + any::().prop_filter( + "Make sure these keys do not exist in the tree.", + move |key| !kvs_clone.contains_key(key), + ), + 100, + ), + ) + }) + ) { + let (db, root_hash_option) = init_mock_db(&existent_kvs); + let tree = JellyfishMerkleTree::new(&db); + let root_hash = root_hash_option.unwrap(); + test_existent_keys_impl(&tree, root_hash, &existent_kvs); + test_nonexistent_keys_impl(&tree, root_hash, &nonexistent_keys); + } + + #[test] + fn test_get_with_proof2( + key1 in any::() + .prop_filter( + "Can't be 0xffffff...", + |key| *key != HashValue::new([0xff; HashValue::LENGTH]), + ), + accounts in vec(any::(), 2), + ) { + let key2 = plus_one(key1); + + let mut kvs = HashMap::new(); + kvs.insert(key1, accounts[0].clone()); + kvs.insert(key2, accounts[1].clone()); + + let (db, root_hash_option) = init_mock_db(&kvs); + let tree = JellyfishMerkleTree::new(&db); + + test_existent_keys_impl(&tree, root_hash_option.unwrap(), &kvs); + } + + #[test] + fn test_get_range_proof( + (btree, n) in btree_map(any::(), any::(), 1..50) + .prop_flat_map(|btree| { + let len = btree.len(); + (Just(btree), 0..len) + }) + ) { + let (db, root_hash_option) = init_mock_db(&btree.clone().into_iter().collect()); + let tree = JellyfishMerkleTree::new(&db); + let root_hash = root_hash_option.unwrap(); + let nth_key = *btree.keys().nth(n).unwrap(); + let proof = tree.get_range_proof(root_hash, nth_key).unwrap(); + verify_range_proof( + root_hash, + btree.into_iter().take(n + 1).collect(), + proof, + ); + } +} + +fn test_existent_keys_impl<'a>( + tree: &JellyfishMerkleTree<'a, MockTreeStore>, + root_hash: HashValue, + existent_kvs: &HashMap, +) { + for (key, value) in existent_kvs { + let (account, proof) = tree.get_with_proof(root_hash, *key).unwrap(); + assert!(proof.verify(root_hash, *key, account.as_ref()).is_ok()); + assert_eq!(account.unwrap(), *value); + } +} + +fn test_nonexistent_keys_impl<'a>( + tree: &JellyfishMerkleTree<'a, MockTreeStore>, + root_hash: HashValue, + nonexistent_keys: &[HashValue], +) { + for key in nonexistent_keys { + let (account, proof) = tree.get_with_proof(root_hash, *key).unwrap(); + assert!(proof.verify(root_hash, *key, account.as_ref()).is_ok()); + assert!(account.is_none()); + } +} + +/// Checks if we can construct the expected root hash using the entries in the btree and the proof. +fn verify_range_proof( + expected_root_hash: HashValue, + btree: BTreeMap, + proof: SparseMerkleRangeProof, +) { + // For example, given the following sparse Merkle tree: + // + // root + // / \ + // / \ + // / \ + // o o + // / \ / \ + // a o o h + // / \ / \ + // o d e X + // / \ / \ + // b c f g + // + // we transform the keys as follows: + // a => 00, + // b => 0100, + // c => 0101, + // d => 011, + // e => 100, + // X => 101 + // h => 11 + // + // Basically, the suffixes that doesn't affect the common prefix of adjacent leaves are + // discarded. In this example, we assume `btree` has the keys `a` to `e` and the proof has `X` + // and `h` in the siblings. + + // Now we want to construct a set of key-value pairs that covers the entire set of leaves. For + // `a` to `e` this is simple -- we just insert them directly into this set. For the rest of the + // leaves, they are represented by the siblings, so we just make up some keys that make sense. + // For example, for `X` we just use 101000... (more zeros omitted), because that is one key + // that would cause `X` to end up in the above position. + let mut btree1 = BTreeMap::new(); + for (key, blob) in &btree { + let leaf = LeafNode::new(*key, blob.clone()); + btree1.insert(*key, leaf.hash()); + } + // Using the above example, `last_proven_key` is `e`. We look at the path from root to `e`. + // For each 0-bit, there should be a sibling in the proof. And we use the path from root to + // this position, plus a `1` as the key. + let last_proven_key = *btree + .keys() + .last() + .expect("We are proving at least one key."); + for (i, sibling) in last_proven_key + .iter_bits() + .enumerate() + .filter_map(|(i, bit)| if !bit { Some(i) } else { None }) + .zip(proof.right_siblings().iter().rev()) + { + // This means the `i`-th bit is zero. We take `i` bits from `last_proven_key` and append a + // one to make up the key for this sibling. + let mut buf: Vec<_> = last_proven_key.iter_bits().take(i).collect(); + buf.push(true); + // The rest doesn't matter, because they don't affect the position of the node. We just + // add zeros. + buf.resize(HashValue::LENGTH_IN_BITS, false); + let key = HashValue::from_bit_iter(buf.into_iter()).unwrap(); + btree1.insert(key, *sibling); + } + + // Now we do the transformation (removing the suffixes) described above. + let mut kvs = vec![]; + for (key, value) in &btree1 { + // The length of the common prefix of the previous key and the current key. + let prev_common_prefix_len = + prev_key(&btree1, key).map(|pkey| pkey.common_prefix_bits_len(*key)); + // The length of the common prefix of the next key and the current key. + let next_common_prefix_len = + next_key(&btree1, key).map(|nkey| nkey.common_prefix_bits_len(*key)); + + // We take the longest common prefix of the current key and its neighbors. That's how much + // we need to keep. + let len = match (prev_common_prefix_len, next_common_prefix_len) { + (Some(plen), Some(nlen)) => std::cmp::max(plen, nlen), + (Some(plen), None) => plen, + (None, Some(nlen)) => nlen, + (None, None) => 0, + }; + let transformed_key: Vec<_> = key.iter_bits().take(len + 1).collect(); + kvs.push((transformed_key, *value)); + } + + assert_eq!(compute_root_hash(kvs), expected_root_hash); +} + +/// Computes the root hash of a sparse Merkle tree. `kvs` consists of the entire set of key-value +/// pairs stored in the tree. +fn compute_root_hash(kvs: Vec<(Vec, HashValue)>) -> HashValue { + let mut kv_ref = vec![]; + for (key, value) in &kvs { + kv_ref.push((&key[..], *value)); + } + compute_root_hash_impl(kv_ref) +} + +fn compute_root_hash_impl(kvs: Vec<(&[bool], HashValue)>) -> HashValue { + assert!(!kvs.is_empty()); + + // If there is only one entry, it is the root. + if kvs.len() == 1 { + return kvs[0].1; + } + + // Otherwise the tree has more than one leaves, which means we can find which ones are in the + // left subtree and which ones are in the right subtree. So we find the first key that starts + // with a 1-bit. + let left_hash; + let right_hash; + match kvs.iter().position(|(key, _value)| key[0]) { + Some(0) => { + // Every key starts with a 1-bit, i.e., they are all in the right subtree. + left_hash = *SPARSE_MERKLE_PLACEHOLDER_HASH; + right_hash = compute_root_hash_impl(reduce(&kvs)); + } + Some(index) => { + // Both left subtree and right subtree have some keys. + left_hash = compute_root_hash_impl(reduce(&kvs[..index])); + right_hash = compute_root_hash_impl(reduce(&kvs[index..])); + } + None => { + // Every key starts with a 0-bit, i.e., they are all in the left subtree. + left_hash = compute_root_hash_impl(reduce(&kvs)); + right_hash = *SPARSE_MERKLE_PLACEHOLDER_HASH; + } + } + + SparseMerkleInternalNode::new(left_hash, right_hash).hash() +} + +/// Reduces the problem by removing the first bit of every key. +fn reduce<'a>(kvs: &'a [(&[bool], HashValue)]) -> Vec<(&'a [bool], HashValue)> { + kvs.iter().map(|(key, value)| (&key[1..], *value)).collect() +} + +/// Returns the key immediately before `key` in `btree`. +fn prev_key(btree: &BTreeMap, key: &K) -> Option +where + K: Clone + Ord, +{ + btree + .range((Bound::Unbounded, Bound::Excluded(key))) + .next_back() + .map(|(k, _v)| k.clone()) +} + +/// Returns the key immediately after `key` in `btree`. +fn next_key(btree: &BTreeMap, key: &K) -> Option +where + K: Clone + Ord, +{ + btree + .range((Bound::Excluded(key), Bound::Unbounded)) + .next() + .map(|(k, _v)| k.clone()) +} diff --git a/core/forkable-jellyfish-merkle/src/lib.rs b/core/forkable-jellyfish-merkle/src/lib.rs index 7fbc61e5a4..8ca1de3b2f 100644 --- a/core/forkable-jellyfish-merkle/src/lib.rs +++ b/core/forkable-jellyfish-merkle/src/lib.rs @@ -84,6 +84,8 @@ pub mod restore; pub mod test_helper; pub mod tree_cache; +#[cfg(test)] +use crate::iterator::JellyfishMerkleIterator; use anyhow::{bail, ensure, format_err, Result}; use blob::Blob; use nibble_path::{skip_common_prefix, NibbleIterator, NibblePath}; @@ -188,6 +190,12 @@ where self.updates(state_root_hash, blob_set) } + #[cfg(test)] + pub fn print_tree(&self, state_root_hash: HashValue, start_key: HashValue) -> Result<()> { + let iter = JellyfishMerkleIterator::new(self.reader, state_root_hash, start_key).unwrap(); + iter.print() + } + /// Delete a key from the tree. If the key is not found in the tree, nothing happens. #[cfg(test)] pub fn delete( @@ -671,10 +679,4 @@ where pub fn get(&self, state_root_hash: HashValue, key: HashValue) -> Result> { Ok(self.get_with_proof(state_root_hash, key)?.0) } - // #[cfg(any(test, feature = "fuzzing"))] - // pub fn get_root_hash(&self, version: Version) -> Result { - // let root_node_key = NodeKey::new_empty_path(version); - // let root_node = self.reader.get_node(&root_node_key)?; - // Ok(root_node.hash()) - // } } diff --git a/core/forkable-jellyfish-merkle/src/node_type/node_type_test.rs b/core/forkable-jellyfish-merkle/src/node_type/node_type_test.rs index c0e73262b0..f84fb29c93 100644 --- a/core/forkable-jellyfish-merkle/src/node_type/node_type_test.rs +++ b/core/forkable-jellyfish-merkle/src/node_type/node_type_test.rs @@ -30,10 +30,8 @@ fn random_63nibblepath() -> NibblePath { // nibble to be appended. fn gen_leaf_keys(nibble_path: &NibblePath, nibble: Nibble) -> HashValue { assert_eq!(nibble_path.num_nibbles(), 63); - dbg!(nibble_path.bytes()); let mut np = nibble_path.clone(); np.push(nibble); - dbg!(np.bytes()); HashValue::from_slice(np.bytes()).unwrap() } diff --git a/core/forkable-jellyfish-merkle/src/tree_cache/mod.rs b/core/forkable-jellyfish-merkle/src/tree_cache/mod.rs index 549f9708bd..1f336dc77a 100644 --- a/core/forkable-jellyfish-merkle/src/tree_cache/mod.rs +++ b/core/forkable-jellyfish-merkle/src/tree_cache/mod.rs @@ -179,6 +179,16 @@ where self.root_node_key = root_node_key; } + #[cfg(test)] + pub fn get_root_node(&self) -> Result { + self.get_node(&self.root_node_key) + } + + #[cfg(test)] + pub fn get_node_num(&self) -> usize { + self.node_cache.len() + } + /// Puts the node with given hash as key into node_cache. pub fn put_node(&mut self, node_key: NodeKey, new_node: Node) -> Result<()> { match self.node_cache.entry(node_key) { diff --git a/state/state-tree/src/state_tree.rs b/state/state-tree/src/state_tree.rs index 956a7b8134..4bbbefdd7d 100644 --- a/state/state-tree/src/state_tree.rs +++ b/state/state-tree/src/state_tree.rs @@ -214,8 +214,7 @@ impl StateTree { store: self.storage.as_ref(), cache, }; - let iterator = - JellyfishMerkleIterator::new(Arc::new(reader), cur_root_hash, HashValue::zero())?; + let iterator = JellyfishMerkleIterator::new(&reader, cur_root_hash, HashValue::zero())?; let mut states = vec![]; for item in iterator { let item = item?;