feat(node): avoid cloning in-memory data structures on apply_block

CHANGELOG.md

@@ -15,6 +15,7 @@
 - Changed the `BlockWitness` to pass the inputs to the VM using only advice provider (#516).
 - [BREAKING] Improved store API errors (return "not found" instead of "internal error" status if requested account(s) not found) (#518).
 - [BREAKING] Migrated to v0.11 version of Miden VM (#528).
+- Reduce cloning in the store's `apply_block` (#532).
 
 ## 0.5.1 (2024-09-12)
 

crates/store/src/nullifier_tree.rs

@@ -1,7 +1,7 @@
 use miden_objects::{
     crypto::{
         hash::rpo::RpoDigest,
-        merkle::{Smt, SmtProof},
+        merkle::{MutationSet, Smt, SmtProof, SMT_DEPTH},
     },
     notes::Nullifier,
     Felt, FieldElement, Word,
@@ -27,7 +27,7 @@ impl NullifierTree {
         Ok(Self(inner))
     }
 
-    /// Get SMT root.
+    /// Returns the root of the nullifier SMT.
     pub fn root(&self) -> RpoDigest {
         self.0.root()
     }
@@ -37,26 +37,6 @@ impl NullifierTree {
         self.0.open(&nullifier.inner())
     }
 
-    /// Inserts block number in which nullifier was consumed.
-    pub fn insert(
-        &mut self,
-        nullifier: &Nullifier,
-        block_num: BlockNumber,
-    ) -> Result<(), NullifierTreeError> {
-        let key = nullifier.inner();
-        let prev_value = self.0.get_value(&key);
-        if prev_value != Smt::EMPTY_VALUE {
-            return Err(NullifierTreeError::NullifierAlreadyExists {
-                nullifier: *nullifier,
-                block_num: Self::leaf_value_to_block_num(prev_value),
-            });
-        }
-
-        self.0.insert(key, Self::block_num_to_leaf_value(block_num));
-
-        Ok(())
-    }
-
     /// Returns block number stored for the given nullifier or `None` if the nullifier wasn't
     /// consumed.
     pub fn get_block_num(&self, nullifier: &Nullifier) -> Option<BlockNumber> {
@@ -68,6 +48,27 @@ impl NullifierTree {
         Some(Self::leaf_value_to_block_num(value))
     }
 
+    /// Computes mutations for the nullifier SMT.
+    pub fn compute_mutations(
+        &self,
+        kv_pairs: impl IntoIterator<Item = (Nullifier, BlockNumber)>,
+    ) -> MutationSet<SMT_DEPTH, RpoDigest, Word> {
+        self.0.compute_mutations(kv_pairs.into_iter().map(|(nullifier, block_num)| {
+            (nullifier.inner(), Self::block_num_to_leaf_value(block_num))
+        }))
+    }
+
+    /// Applies mutations to the nullifier SMT.
+    pub fn apply_mutations(
+        &mut self,
+        mutations: MutationSet<SMT_DEPTH, RpoDigest, Word>,
+    ) -> Result<(), NullifierTreeError> {
+        self.0.apply_mutations(mutations).map_err(Into::into)
+    }
+
+    // HELPER FUNCTIONS
+    // --------------------------------------------------------------------------------------------
+
     /// Returns the nullifier's leaf value in the SMT by its block number.
     fn block_num_to_leaf_value(block: BlockNumber) -> Word {
         [Felt::from(block), Felt::ZERO, Felt::ZERO, Felt::ZERO]

crates/store/src/state.rs

@@ -203,15 +203,21 @@ impl State {
         // block store. Thus, such block should be considered as block candidates, but not
         // finalized blocks. So we should check for the latest block when getting block from
         // the store.
-        let store = self.block_store.clone();
+        let store = Arc::clone(&self.block_store);
         let block_save_task =
             tokio::spawn(async move { store.save_block(block_num, &block_data).await });
 
-        // scope to read in-memory data, validate the request, and compute intermediary values
-        let (account_tree, chain_mmr, nullifier_tree, notes) = {
+        // scope to read in-memory data, compute mutations required for updating account
+        // and nullifier trees, and validate the request
+        let (
+            nullifier_tree_old_root,
+            account_tree_old_root,
+            account_tree_update,
+            nullifier_tree_update,
+        ) = {
             let inner = self.inner.read().await;
 
-            let span = info_span!(target: COMPONENT, "update_in_memory_structs").entered();
+            let _span = info_span!(target: COMPONENT, "update_in_memory_structs").entered();
 
             // nullifiers can be produced only once
             let duplicate_nullifiers: Vec<_> = block
@@ -224,83 +230,72 @@ impl State {
                 return Err(ApplyBlockError::DuplicatedNullifiers(duplicate_nullifiers));
             }
 
-            // update the in-memory data structures and compute the new block header. Important, the
-            // structures are not yet committed
+            // compute updates for the in-memory data structures
 
-            // update chain MMR
-            let chain_mmr = {
-                let mut chain_mmr = inner.chain_mmr.clone();
-
-                // new_block.chain_root must be equal to the chain MMR root prior to the update
-                let peaks = chain_mmr.peaks();
-                if peaks.hash_peaks() != header.chain_root() {
-                    return Err(ApplyBlockError::NewBlockInvalidChainRoot);
-                }
-
-                chain_mmr.add(block_hash);
-                chain_mmr
-            };
-
-            // update nullifier tree
-            let nullifier_tree = {
-                let mut nullifier_tree = inner.nullifier_tree.clone();
-                for nullifier in block.nullifiers() {
-                    nullifier_tree
-                        .insert(nullifier, block_num)
-                        .map_err(ApplyBlockError::FailedToUpdateNullifierTree)?;
-                }
+            // new_block.chain_root must be equal to the chain MMR root prior to the update
+            let peaks = inner.chain_mmr.peaks();
+            if peaks.hash_peaks() != header.chain_root() {
+                return Err(ApplyBlockError::NewBlockInvalidChainRoot);
+            }
 
-                if nullifier_tree.root() != header.nullifier_root() {
-                    return Err(ApplyBlockError::NewBlockInvalidNullifierRoot);
-                }
-                nullifier_tree
-            };
+            // compute update for nullifier tree
+            let nullifier_tree_update = inner.nullifier_tree.compute_mutations(
+                block.nullifiers().iter().map(|nullifier| (*nullifier, block_num)),
+            );
 
-            // update account tree
-            let mut account_tree = inner.account_tree.clone();
-            for update in block.updated_accounts() {
-                account_tree.insert(
-                    LeafIndex::new_max_depth(update.account_id().into()),
-                    update.new_state_hash().into(),
-                );
+            if nullifier_tree_update.root() != header.nullifier_root() {
+                return Err(ApplyBlockError::NewBlockInvalidNullifierRoot);
             }
 
-            if account_tree.root() != header.account_root() {
+            // compute update for account tree
+            let account_tree_update = inner.account_tree.compute_mutations(
+                block.updated_accounts().iter().map(|update| {
+                    (
+                        LeafIndex::new_max_depth(update.account_id().into()),
+                        update.new_state_hash().into(),
+                    )
+                }),
+            );
+
+            if account_tree_update.root() != header.account_root() {
                 return Err(ApplyBlockError::NewBlockInvalidAccountRoot);
             }
 
-            // build note tree
-            let note_tree = block.build_note_tree();
-            if note_tree.root() != header.note_root() {
-                return Err(ApplyBlockError::NewBlockInvalidNoteRoot);
-            }
+            (
+                inner.nullifier_tree.root(),
+                inner.account_tree.root(),
+                account_tree_update,
+                nullifier_tree_update,
+            )
+        };
 
-            drop(span);
-
-            let notes = block
-                .notes()
-                .map(|(note_index, note)| {
-                    let details = match note {
-                        OutputNote::Full(note) => Some(note.to_bytes()),
-                        OutputNote::Header(_) => None,
-                        note => return Err(ApplyBlockError::InvalidOutputNoteType(note.clone())),
-                    };
-
-                    let merkle_path = note_tree.get_note_path(note_index);
-
-                    Ok(NoteRecord {
-                        block_num,
-                        note_index,
-                        note_id: note.id().into(),
-                        metadata: *note.metadata(),
-                        details,
-                        merkle_path,
-                    })
-                })
-                .collect::<Result<Vec<NoteRecord>, ApplyBlockError>>()?;
+        // build note tree
+        let note_tree = block.build_note_tree();
+        if note_tree.root() != header.note_root() {
+            return Err(ApplyBlockError::NewBlockInvalidNoteRoot);
+        }
 
-            (account_tree, chain_mmr, nullifier_tree, notes)
-        };
+        let notes = block
+            .notes()
+            .map(|(note_index, note)| {
+                let details = match note {
+                    OutputNote::Full(note) => Some(note.to_bytes()),
+                    OutputNote::Header(_) => None,
+                    note => return Err(ApplyBlockError::InvalidOutputNoteType(note.clone())),
+                };
+
+                let merkle_path = note_tree.get_note_path(note_index);
+
+                Ok(NoteRecord {
+                    block_num,
+                    note_index,
+                    note_id: note.id().into(),
+                    metadata: *note.metadata(),
+                    details,
+                    merkle_path,
+                })
+            })
+            .collect::<Result<Vec<NoteRecord>, ApplyBlockError>>()?;
 
         // Signals the transaction is ready to be committed, and the write lock can be acquired
         let (allow_acquire, acquired_allowed) = oneshot::channel::<()>();
@@ -311,7 +306,7 @@ impl State {
         // overlapping. Namely, the DB transaction only proceeds after this task acquires the
         // in-memory write lock. This requires the DB update to run concurrently, so a new task is
         // spawned.
-        let db = self.db.clone();
+        let db = Arc::clone(&self.db);
         let db_update_task =
             tokio::spawn(
                 async move { db.apply_block(allow_acquire, acquire_done, block, notes).await },
@@ -332,6 +327,16 @@ impl State {
             // successfully.
             let mut inner = self.inner.write().await;
 
+            // We need to check that neither the nullifier tree root, nor the account tree root
+            // haven't changed while we were waiting for the DB preparation task to complete.
+            // If it has changed, we mustn't proceed with both in-memory and database updates,
+            // since it might lead to inconsistent state.
+            if inner.nullifier_tree.root() != nullifier_tree_old_root
+                || inner.account_tree.root() != account_tree_old_root
+            {
+                return Err(ApplyBlockError::ConcurrentWrite);
+            }
+
             // Notify the DB update task that the write lock has been acquired, so it can commit
             // the DB transaction
             let _ = inform_acquire_done.send(());
@@ -343,9 +348,15 @@ impl State {
             db_update_task.await??;
 
             // Update the in-memory data structures after successful commit of the DB transaction
-            inner.chain_mmr = chain_mmr;
-            inner.nullifier_tree = nullifier_tree;
-            inner.account_tree = account_tree;
+            inner
+                .nullifier_tree
+                .apply_mutations(nullifier_tree_update)
+                .expect("Unreachable: old nullifier tree root must be checked before this step");
+            inner
+                .account_tree
+                .apply_mutations(account_tree_update)
+                .expect("Unreachable: old account tree root must be checked before this step");
+            inner.chain_mmr.add(block_hash);
         }
 
         info!(%block_hash, block_num, COMPONENT, "apply_block successful");