temp

storage/aptosdb/src/state_store/state_merkle_batch_committer.rs

@@ -5,20 +5,19 @@
 
 use crate::{
     metrics::{LATEST_SNAPSHOT_VERSION, OTHER_TIMERS_SECONDS},
+    schema::jellyfish_merkle_node::JellyfishMerkleNodeSchema,
+    state_merkle_db::Node,
     state_store::{buffered_state::CommitMessage, StateDb},
 };
 use anyhow::{anyhow, ensure, Result};
 use aptos_crypto::HashValue;
+use aptos_jellyfish_merkle::node_type::NodeKey;
 use aptos_logger::{info, trace};
 use aptos_metrics_core::TimerHelper;
 use aptos_scratchpad::SmtAncestors;
 use aptos_storage_interface::state_delta::StateDelta;
-use aptos_types::state_store::state_value::StateValue;
+use aptos_types::state_store::{state_key::StateKey, state_value::StateValue};
 use std::sync::{mpsc::Receiver, Arc};
-use aptos_jellyfish_merkle::node_type::{Children, NodeKey};
-use aptos_types::state_store::state_key::StateKey;
-use crate::schema::jellyfish_merkle_node::JellyfishMerkleNodeSchema;
-use crate::state_merkle_db::Node;
 
 pub struct StateMerkleBatch {
     // pub top_levels_batch: SchemaBatch,
@@ -66,8 +65,7 @@ impl StateMerkleBatchCommitter {
                         .with_label_values(&["commit_jellyfish_merkle_nodes"])
                         .start_timer();
 
-                    self
-                        .state_db
+                    self.state_db
                         .state_merkle_db
                         .metadata_db()
                         .put::<JellyfishMerkleNodeSchema>(
@@ -76,8 +74,9 @@ impl StateMerkleBatchCommitter {
                                 root_hash,
                                 root_hash,
                                 (StateKey::raw(b"root"), current_version),
-                            )
-                        ).unwrap();
+                            ),
+                        )
+                        .unwrap();
 
                     /*
                     // commit jellyfish merkle nodes

storage/aptosdb/src/state_store/state_snapshot_committer.rs

@@ -13,7 +13,7 @@ use crate::{
 };
 use aptos_logger::trace;
 use aptos_scratchpad::SmtAncestors;
-use aptos_storage_interface::{state_delta::StateDelta};
+use aptos_storage_interface::state_delta::StateDelta;
 use aptos_types::state_store::state_value::StateValue;
 use static_assertions::const_assert;
 use std::{

storage/scratchpad/src/sparse_merkle/mod.rs

@@ -100,7 +100,8 @@ use aptos_types::{
     nibble::nibble_path::NibblePath, proof::SparseMerkleLeafNode,
     state_store::state_storage_usage::StateStorageUsage,
 };
-use fastcrypto::hash::MultisetHash;
+use fastcrypto::hash::{EllipticCurveMultisetHash, MultisetHash};
+use rayon::prelude::*;
 use std::{
     collections::{BTreeMap, HashMap},
     sync::Arc,
@@ -422,6 +423,36 @@ where
             return Ok(self.clone());
         }
 
+        let inc_hash_diff = updates
+            .par_iter()
+            .fold(
+                EllipticCurveMultisetHash::default,
+                |mut inc_hash, (key, val_opt)| {
+                    if let Some(old_val) = proof_reader.get_proof(key.0) {
+                        let old_hash = SparseMerkleLeafNode::new(key.0, old_val)
+                            .hash()
+                            .into_inner();
+                        inc_hash.remove(old_hash);
+                    }
+                    if let Some(Some(new_val)) = val_opt {
+                        let new_hash = SparseMerkleLeafNode::new(key.0, new_val.hash())
+                            .hash()
+                            .into_inner();
+                        inc_hash.insert(new_hash);
+                    }
+                    inc_hash
+                },
+            )
+            .reduce(
+                EllipticCurveMultisetHash::default,
+                |mut inc_hash, other_inc_hash| {
+                    inc_hash.union(&other_inc_hash);
+                    inc_hash
+                },
+            );
+        let mut inc_hash = self.smt.inner.root().inc_hash.clone();
+        inc_hash.union(&inc_hash_diff);
+
         let content = self
             .smt
             .inner
@@ -430,24 +461,6 @@ where
             .view_layers_since(&self.base_smt.inner.root().content)
             .new_layer(&updates[..]);
 
-        let hashes_to_remove = updates.iter().filter_map(|(k, _)| {
-            proof_reader.get_proof(k.0).map(|value_hash| {
-                SparseMerkleLeafNode::new(k.0, value_hash)
-                    .hash()
-                    .into_inner()
-            })
-        });
-        let hashes_to_insert = updates.iter().filter_map(|(k, v)| {
-            v.as_ref().and_then(|val| {
-                val.as_ref()
-                    .map(|v| SparseMerkleLeafNode::new(k.0, v.hash()).hash().into_inner())
-            })
-        });
-
-        let mut inc_hash = self.smt.inner.root().inc_hash.clone();
-        inc_hash.remove_all(hashes_to_remove);
-        inc_hash.insert_all(hashes_to_insert);
-
         let root = Root::new(inc_hash, content);
 
         Ok(self.spawn(root, usage))

storage/storage-interface/Cargo.toml

@@ -24,7 +24,7 @@ aptos-types = { workspace = true }
 aptos-vm = { workspace = true }
 bcs = { workspace = true }
 crossbeam-channel = { workspace = true }
-dashmap = { workspace = true }
+dashmap = { workspace = true, features = ["rayon"] }
 once_cell = { workspace = true }
 parking_lot = { workspace = true }
 proptest = { workspace = true }

storage/storage-interface/src/cached_state_view.rs

@@ -6,6 +6,7 @@ use crate::{
 };
 use aptos_crypto::{hash::CryptoHash, HashValue};
 use aptos_experimental_runtimes::thread_manager::THREAD_MANAGER;
+use aptos_logger::info;
 use aptos_scratchpad::{FrozenSparseMerkleTree, SparseMerkleTree, StateStoreStatus};
 use aptos_types::{
     state_store::{
@@ -25,7 +26,6 @@ use std::{
     fmt::{Debug, Formatter},
     sync::Arc,
 };
-use aptos_logger::info;
 
 static IO_POOL: Lazy<rayon::ThreadPool> = Lazy::new(|| {
     rayon::ThreadPoolBuilder::new()
@@ -168,10 +168,7 @@ impl CachedStateView {
         let snapshot = reader
             .get_state_snapshot_before(next_version)
             .map_err(Into::<StateviewError>::into)?;
-        info!(
-            snapshot = snapshot,
-            "alden alden alden."
-        );
+        info!(snapshot = snapshot, "alden alden alden.");
 
         Ok(Self::new_impl(id, snapshot, speculative_state, reader))
     }
@@ -212,15 +209,17 @@ impl CachedStateView {
     }
 
     pub fn into_state_cache(self) -> StateCache {
-        // FIXME(aldenhu): make faster
         let proofs = self
             .sharded_state_cache
             .par_iter()
             .map(|shard| {
-                shard.iter().map(|dashmap_ref| {
-                    let (key, (_val_ver_opt, val_opt)) = dashmap_ref.pair();
-                    (key.hash(), val_opt.as_ref().map(CryptoHash::hash))
-                })
+                shard
+                    .par_iter()
+                    .map(|dashmap_ref| {
+                        let (key, (_val_ver_opt, val_opt)) = dashmap_ref.pair();
+                        (key.hash(), val_opt.as_ref().map(CryptoHash::hash))
+                    })
+                    .collect::<Vec<_>>()
             })
             .collect::<Vec<_>>()
             .into_iter()