[jf+exec] add batch_update method to Jellyfish merkle to leverage exe…

…cution result

Closes: aptos-labs#8574

Cargo.toml

@@ -19,7 +19,6 @@ members = [
     "common/logger/derive",
     "common/metrics",
     "common/metrics-core",
-    "common/nibble",
     "common/num-variants",
     "common/proptest-helpers",
     "common/proxy",

crypto/crypto/src/hash.rs

@@ -212,6 +212,14 @@ impl HashValue {
         (self.hash[pos] >> bit) & 1 != 0
     }
 
+    /// Returns the `index`-th nibble in the bytes.
+    pub fn nibble(&self, index: usize) -> u8 {
+        assume!(index < Self::LENGTH * 2); // assumed precondition
+        let pos = index / 2;
+        let shift = if index % 2 == 0 { 4 } else { 0 };
+        (self.hash[pos] >> shift) & 0x0f
+    }
+
     /// Returns a `HashValueBitIterator` over all the bits that represent this `HashValue`.
     pub fn iter_bits(&self) -> HashValueBitIterator<'_> {
         HashValueBitIterator::new(self)

crypto/crypto/src/unit_tests/hash_test.rs

@@ -138,6 +138,21 @@ fn test_fmt_binary() {
     }
 }
 
+#[test]
+fn test_get_nibble() {
+    let mut bytes = [0u8; HashValue::LENGTH];
+    let mut nibbles = vec![];
+    for byte in bytes.iter_mut().take(HashValue::LENGTH) {
+        *byte = rand::thread_rng().gen();
+        nibbles.push(*byte >> 4);
+        nibbles.push(*byte & 0x0f);
+    }
+    let hash = HashValue::new(bytes);
+    for (i, nibble) in nibbles.iter().enumerate().take(HashValue::LENGTH * 2) {
+        assert_eq!(hash.nibble(i), *nibble);
+    }
+}
+
 #[test]
 fn test_common_prefix_bits_len() {
     {

execution/executor/src/lib.rs

@@ -292,7 +292,7 @@ where
             })
             .collect::<Result<Vec<_>>>()?;
 
-        let (txn_state_roots, current_state_tree) = parent_trees
+        let (roots_with_node_hashes, current_state_tree) = parent_trees
             .state_tree()
             .serial_update(
                 txn_blobs
@@ -307,9 +307,9 @@ where
             )
             .expect("Failed to update state tree.");
 
-        for ((vm_output, txn), (state_tree_hash, blobs)) in itertools::zip_eq(
+        for ((vm_output, txn), ((state_tree_hash, new_node_hashes), blobs)) in itertools::zip_eq(
             itertools::zip_eq(vm_outputs.into_iter(), transactions.iter()).take(transaction_count),
-            itertools::zip_eq(txn_state_roots, txn_blobs),
+            itertools::zip_eq(roots_with_node_hashes, txn_blobs),
         ) {
             let event_tree = {
                 let event_hashes: Vec<_> =
@@ -353,6 +353,7 @@ where
 
             txn_data.push(TransactionData::new(
                 blobs,
+                new_node_hashes,
                 vm_output.events().to_vec(),
                 vm_output.status().clone(),
                 state_tree_hash,
@@ -368,6 +369,7 @@ where
             txn_data.resize(
                 transactions.len(),
                 TransactionData::new(
+                    HashMap::new(),
                     HashMap::new(),
                     vec![],
                     TransactionStatus::Retry,
@@ -558,6 +560,7 @@ where
             txns_to_commit.push(TransactionToCommit::new(
                 txn,
                 txn_data.account_blobs().clone(),
+                Some(txn_data.jf_node_hashes().clone()),
                 txn_data.events().to_vec(),
                 txn_data.gas_used(),
                 recorded_status,
@@ -891,6 +894,7 @@ impl<V: VMExecutor> BlockExecutor for Executor<V> {
                 txns_to_keep.push(TransactionToCommit::new(
                     txn.clone(),
                     txn_data.account_blobs().clone(),
+                    Some(txn_data.jf_node_hashes().clone()),
                     txn_data.events().to_vec(),
                     txn_data.gas_used(),
                     recorded_status.clone(),
@@ -945,7 +949,6 @@ impl<V: VMExecutor> BlockExecutor for Executor<V> {
 
         // Skip duplicate txns that are already persistent.
         let txns_to_commit = &txns_to_keep[num_txns_to_skip as usize..];
-
         let num_txns_to_commit = txns_to_commit.len() as u64;
         {
             let _timer = DIEM_EXECUTOR_SAVE_TRANSACTIONS_SECONDS.start_timer();
@@ -957,6 +960,7 @@ impl<V: VMExecutor> BlockExecutor for Executor<V> {
                     "Injected error in commit_blocks"
                 )))
             });
+
             self.db.writer.save_transactions(
                 txns_to_commit,
                 first_version_to_commit,

execution/executor/src/types.rs

@@ -9,6 +9,7 @@ use diem_types::{
     account_state_blob::AccountStateBlob,
     contract_event::ContractEvent,
     epoch_state::EpochState,
+    nibble::nibble_path::NibblePath,
     on_chain_config,
     proof::accumulator::InMemoryAccumulator,
     transaction::{TransactionStatus, Version},
@@ -26,6 +27,10 @@ pub struct TransactionData {
     /// new blob.
     account_blobs: HashMap<AccountAddress, AccountStateBlob>,
 
+    /// Each entry in this map represents the the hash of a newly generated jellyfish node
+    /// and its corresponding nibble path.
+    jf_node_hashes: HashMap<NibblePath, HashValue>,
+
     /// The list of events emitted during this transaction.
     events: Vec<ContractEvent>,
 
@@ -48,6 +53,7 @@ pub struct TransactionData {
 impl TransactionData {
     pub fn new(
         account_blobs: HashMap<AccountAddress, AccountStateBlob>,
+        jf_node_hashes: HashMap<NibblePath, HashValue>,
         events: Vec<ContractEvent>,
         status: TransactionStatus,
         state_root_hash: HashValue,
@@ -57,6 +63,7 @@ impl TransactionData {
     ) -> Self {
         TransactionData {
             account_blobs,
+            jf_node_hashes,
             events,
             status,
             state_root_hash,
@@ -70,6 +77,10 @@ impl TransactionData {
         &self.account_blobs
     }
 
+    pub fn jf_node_hashes(&self) -> &HashMap<NibblePath, HashValue> {
+        &self.jf_node_hashes
+    }
+
     pub fn events(&self) -> &[ContractEvent] {
         &self.events
     }

storage/diemdb-benchmark/src/lib.rs

@@ -66,6 +66,7 @@ fn gen_txn_to_commit<R: Rng>(
     TransactionToCommit::new(
         txn,
         states,
+        None,
         vec![], /* events */
         0,      /* gas_used */
         KeptVMStatus::Executed,

storage/diemdb/src/lib.rs

@@ -539,9 +539,17 @@ impl DiemDB {
             .iter()
             .map(|txn_to_commit| txn_to_commit.account_states().clone())
             .collect::<Vec<_>>();
-        let state_root_hashes =
-            self.state_store
-                .put_account_state_sets(account_state_sets, first_version, &mut cs)?;
+
+        let node_hashes = txns_to_commit
+            .iter()
+            .map(|txn_to_commit| txn_to_commit.jf_node_hashes())
+            .collect::<Option<Vec<_>>>();
+        let state_root_hashes = self.state_store.put_account_state_sets(
+            account_state_sets,
+            node_hashes,
+            first_version,
+            &mut cs,
+        )?;
 
         // Event updates. Gather event accumulator root hashes.
         let event_root_hashes = zip_eq(first_version..=last_version, txns_to_commit)

storage/diemdb/src/pruner/test.rs

@@ -18,6 +18,7 @@ fn put_account_state_set(
     let root = state_store
         .put_account_state_sets(
             vec![account_state_set.into_iter().collect::<HashMap<_, _>>()],
+            None,
             version,
             &mut cs,
         )

storage/diemdb/src/state_store/mod.rs

@@ -15,12 +15,11 @@ use crate::{
 };
 use anyhow::Result;
 use diem_crypto::HashValue;
-use diem_jellyfish_merkle::{
-    node_type::NodeKey, JellyfishMerkleTree, TreeReader, TreeWriter, ROOT_NIBBLE_HEIGHT,
-};
+use diem_jellyfish_merkle::{node_type::NodeKey, JellyfishMerkleTree, TreeReader, TreeWriter};
 use diem_types::{
     account_address::{AccountAddress, HashAccountAddress},
     account_state_blob::AccountStateBlob,
+    nibble::{nibble_path::NibblePath, ROOT_NIBBLE_HEIGHT},
     proof::{SparseMerkleProof, SparseMerkleRangeProof},
     transaction::Version,
 };
@@ -67,6 +66,7 @@ impl StateStore {
     pub fn put_account_state_sets(
         &self,
         account_state_sets: Vec<HashMap<AccountAddress, AccountStateBlob>>,
+        node_hashes: Option<Vec<&HashMap<NibblePath, HashValue>>>,
         first_version: Version,
         cs: &mut ChangeSet,
     ) -> Result<Vec<HashValue>> {
@@ -80,8 +80,8 @@ impl StateStore {
             })
             .collect::<Vec<_>>();
 
-        let (new_root_hash_vec, tree_update_batch) =
-            JellyfishMerkleTree::new(self).put_value_sets(blob_sets, first_version)?;
+        let (new_root_hash_vec, tree_update_batch) = JellyfishMerkleTree::new(self)
+            .batch_put_value_sets(blob_sets, node_hashes, first_version)?;
 
         let num_versions = new_root_hash_vec.len();
         assert_eq!(num_versions, tree_update_batch.node_stats.len());

storage/diemdb/src/state_store/state_store_test.rs

@@ -24,6 +24,7 @@ fn put_account_state_set(
     let root = store
         .put_account_state_sets(
             vec![account_state_set.into_iter().collect::<HashMap<_, _>>()],
+            None,
             version,
             &mut cs,
         )
@@ -368,7 +369,7 @@ fn init_store(store: &StateStore, input: impl Iterator<Item = (AccountAddress, A
         let mut cs = ChangeSet::new();
         let account_state_set: HashMap<_, _> = std::iter::once((key, value)).collect();
         store
-            .put_account_state_sets(vec![account_state_set], i as Version, &mut cs)
+            .put_account_state_sets(vec![account_state_set], None, i as Version, &mut cs)
             .unwrap();
         store.db.write_schemas(cs.batch).unwrap();
     }

storage/diemdb/src/test_helper.rs

@@ -33,6 +33,7 @@ fn to_blocks_to_commit(
                 let txn_hash = txn_to_commit.transaction().hash();
                 let state_root_hash = db.state_store.put_account_state_sets(
                     vec![txn_to_commit.account_states().clone()],
+                    None,
                     cur_ver,
                     &mut cs,
                 )?[0];

storage/jellyfish-merkle/Cargo.toml

@@ -12,6 +12,7 @@ edition = "2018"
 [dependencies]
 anyhow = "1.0.38"
 byteorder = "1.4.3"
+itertools = { version = "0.10.0", default-features = false }
 mirai-annotations = "1.10.1"
 num-derive = "0.3.3"
 num-traits = "0.2.14"
@@ -27,7 +28,6 @@ diem-crypto = { path = "../../crypto/crypto" }
 diem-crypto-derive = { path = "../../crypto/crypto-derive" }
 diem-infallible = { path = "../../common/infallible" }
 diem-metrics = { path = "../../common/metrics" }
-diem-nibble = { path = "../../common/nibble" }
 diem-types = { path = "../../types" }
 diem-workspace-hack = { path = "../../common/workspace-hack" }
 
@@ -37,9 +37,8 @@ proptest = "1.0.0"
 proptest-derive = "0.3.0"
 
 diem-crypto = { path = "../../crypto/crypto", features = ["fuzzing"] }
-diem-nibble = { path = "../../common/nibble", features = ["fuzzing"] }
 diem-types = { path = "../../types", features = ["fuzzing"] }
 
 [features]
 default = []
-fuzzing = ["proptest", "rand", "proptest-derive", "diem-crypto/fuzzing", "diem-types/fuzzing", "diem-nibble/fuzzing"]
+fuzzing = ["proptest", "rand", "proptest-derive", "diem-crypto/fuzzing", "diem-types/fuzzing"]

storage/jellyfish-merkle/src/iterator/mod.rs

@@ -10,14 +10,15 @@
 mod iterator_test;
 
 use crate::{
-    nibble_path::NibblePath,
     node_type::{InternalNode, Node, NodeKey},
     TreeReader,
 };
 use anyhow::{format_err, Result};
 use diem_crypto::HashValue;
-use diem_nibble::Nibble;
-use diem_types::transaction::Version;
+use diem_types::{
+    nibble::{nibble_path::NibblePath, Nibble},
+    transaction::Version,
+};
 use std::{marker::PhantomData, sync::Arc};
 
 /// `NodeVisitInfo` keeps track of the status of an internal node during the iteration process. It