perf/store: optimize updating existing accounts

CHANGELOG.md

@@ -31,6 +31,7 @@
 - Pin tool versions in CI ([#1523](https://github.com/0xMiden/miden-node/pull/1523)).
 - Add `GetVaultAssetWitnesses` and `GetStorageMapWitness` RPC endpoints to store ([#1529](https://github.com/0xMiden/miden-node/pull/1529)).
 - Add check to ensure tree store state is in sync with database storage ([#1532](https://github.com/0xMiden/miden-node/issues/1534)).
+- Improve speed account updates ([#1567](https://github.com/0xMiden/miden-node/pull/1567)).
 - Ensure store terminates on nullifier tree or account tree root vs header mismatch (#[#1569](https://github.com/0xMiden/miden-node/pull/1569)).
 
 ### Changes

alt_approach.txt

@@ -0,0 +1,44 @@
+# Alternative Approach: Pre-compute in InnerForest
+
+## Current Flow (this PR)
+```
+db.apply_block() -> upsert_accounts() -> [ad-hoc SmtForest computes roots]
+                                                    |
+State::apply_block() -----> InnerForest::apply_block_updates() -> [InnerForest recomputes same roots]
+```
+
+Duplicate SMT computation in both DB layer and InnerForest.
+
+## Alternative Flow
+```
+InnerForest::apply_block_updates() -> [computes roots, stores results]
+                                                    |
+                                                    v
+db.apply_block(precomputed_roots) -> upsert_accounts() -> [uses pre-computed roots]
+```
+
+## Required Changes
+
+1. **Extend InnerForest** to also track:
+   - Value slot updates (currently only tracks map roots)
+   - Full `AccountStorageHeader` per (account_id, block_num)
+
+2. **Add extraction method**:
+   ```rust
+   InnerForest::get_precomputed_state(account_id, block_num) -> (AccountStorageHeader, vault_root)
+   ```
+
+3. **Reorder apply_block**:
+   - Update InnerForest BEFORE db.apply_block
+   - Pass pre-computed roots to upsert_accounts
+
+4. **Remove** `apply_storage_delta_to_header()` and `compute_vault_root_after_delta()`
+
+## Trade-offs
+
+| Pros | Cons |
+|------|------|
+| Single SMT computation | Tighter coupling between State and DB |
+| InnerForest as single source of truth | Must rollback InnerForest if DB fails |
+| | More memory (store full headers) |
+| | Complex locking during apply_block |

crates/store/src/db/models/queries/accounts.rs

@@ -23,12 +23,10 @@ use miden_node_utils::limiter::{
     QueryParamAccountIdLimit,
     QueryParamLimiter,
 };
-use miden_protocol::Word;
 use miden_protocol::account::delta::AccountUpdateDetails;
 use miden_protocol::account::{
     Account,
     AccountCode,
-    AccountDelta,
     AccountId,
     AccountStorage,
     AccountStorageHeader,
@@ -42,6 +40,7 @@ use miden_protocol::account::{
 use miden_protocol::asset::{Asset, AssetVault, AssetVaultKey, FungibleAsset};
 use miden_protocol::block::{BlockAccountUpdate, BlockNumber};
 use miden_protocol::utils::{Deserializable, Serializable};
+use miden_protocol::{Felt, Word};
 
 use crate::COMPONENT;
 use crate::db::models::conv::{
@@ -60,6 +59,16 @@ pub(crate) use at_block::{
     select_account_vault_at_block,
 };
 
+mod delta;
+use delta::{
+    AccountStateForInsert,
+    PartialAccountState,
+    apply_storage_delta_to_header,
+    compute_vault_root_after_delta,
+    select_account_state_for_delta,
+    select_vault_balances_by_faucet_ids,
+};
+
 #[cfg(test)]
 mod tests;
 
@@ -158,7 +167,7 @@ pub(crate) fn select_account(
 /// `State` which contains an `SmtForest` to serve the latest and most recent
 /// historical data.
 // TODO: remove eventually once refactoring is complete
-fn select_full_account(
+pub(crate) fn select_full_account(
     conn: &mut SqliteConnection,
     account_id: AccountId,
 ) -> Result<Account, DatabaseError> {
@@ -954,9 +963,9 @@ pub(crate) fn upsert_accounts(
         // written. The storage and vault tables have FKs pointing to `accounts (account_id,
         // block_num)`, so inserting them earlier would violate those constraints when inserting a
         // brand-new account.
-        let (full_account, pending_storage_inserts, pending_asset_inserts) = match update.details()
+        let (account_state, pending_storage_inserts, pending_asset_inserts) = match update.details()
         {
-            AccountUpdateDetails::Private => (None, vec![], vec![]),
+            AccountUpdateDetails::Private => (AccountStateForInsert::Private, vec![], vec![]),
 
             AccountUpdateDetails::Delta(delta) if delta.is_full_state() => {
                 let account = Account::try_from(delta)?;
@@ -992,12 +1001,14 @@ pub(crate) fn upsert_accounts(
                     }
                 }
 
-                (Some(account), storage, assets)
+                (AccountStateForInsert::FullAccount(account), storage, assets)
             },
 
             AccountUpdateDetails::Delta(delta) => {
-                // Reconstruct the full account from database tables
-                let account = select_full_account(conn, account_id)?;
+                // OPTIMIZATION: Load only the minimal data needed for delta updates.
+                // Avoids loading full code bytes, all storage map entries, and all vault
+                // assets.
+                let state = select_account_state_for_delta(conn, account_id)?;
 
                 // --- collect storage map updates ----------------------------
 
@@ -1008,20 +1019,27 @@ pub(crate) fn upsert_accounts(
                     }
                 }
 
-                // apply delta to the account; we need to do this before we process asset updates
-                // because we currently need to get the current value of fungible assets from the
-                // account
-                let account_after = apply_delta(account, delta, &update.final_state_commitment())?;
-
                 // --- process asset updates ----------------------------------
+                // Only query balances for faucet_ids that are being updated
+                let faucet_ids: Vec<AccountId> =
+                    Vec::from_iter(delta.vault().fungible().iter().map(|(id, _)| *id));
+                let prev_balances =
+                    select_vault_balances_by_faucet_ids(conn, account_id, &faucet_ids)?;
 
                 let mut assets = Vec::new();
 
-                for (faucet_id, _) in delta.vault().fungible().iter() {
-                    let current_amount = account_after.vault().get_balance(*faucet_id).unwrap();
-                    let asset: Asset = FungibleAsset::new(*faucet_id, current_amount)?.into();
-                    let update_or_remove = if current_amount == 0 { None } else { Some(asset) };
-
+                for (faucet_id, amount_delta) in delta.vault().fungible().iter() {
+                    let prev_balance = prev_balances.get(faucet_id).copied().unwrap_or(0);
+                    let new_balance = (i128::from(prev_balance) + i128::from(*amount_delta))
+                        .try_into()
+                        .map_err(|_| {
+                            DatabaseError::DataCorrupted(format!(
+                                "Balance underflow for account {account_id}, faucet {faucet_id}"
+                            ))
+                        })?;
+
+                    let asset: Asset = FungibleAsset::new(*faucet_id, new_balance)?.into();
+                    let update_or_remove = if new_balance == 0 { None } else { Some(asset) };
                     assets.push((account_id, asset.vault_key(), update_or_remove));
                 }
 
@@ -1033,11 +1051,36 @@ pub(crate) fn upsert_accounts(
                     assets.push((account_id, asset.vault_key(), asset_update));
                 }
 
-                (Some(account_after), storage, assets)
+                // --- compute updated account state for the accounts row ---
+                // Apply nonce delta
+                let new_nonce = Felt::new(state.nonce.as_int() + delta.nonce_delta().as_int());
+
+                // Apply storage value updates to header
+                let new_storage_header =
+                    apply_storage_delta_to_header(&state.storage_header, delta.storage())?;
+
+                // Compute new vault root using SMT operations
+                let new_vault_root = compute_vault_root_after_delta(
+                    state.vault_root,
+                    delta.vault(),
+                    &prev_balances,
+                )?;
+
+                // Create minimal account state data for the row insert
+                let account_state = PartialAccountState {
+                    nonce: new_nonce,
+                    code_commitment: state.code_commitment,
+                    storage_header: new_storage_header,
+                    vault_root: new_vault_root,
+                };
+
+                (AccountStateForInsert::PartialState(account_state), storage, assets)
             },
         };
 
-        if let Some(code) = full_account.as_ref().map(Account::code) {
+        // Insert account code for full accounts (new account creation)
+        if let AccountStateForInsert::FullAccount(ref account) = account_state {
+            let code = account.code();
             let code_value = AccountCodeRowInsert {
                 code_commitment: code.commitment().to_bytes(),
                 code: code.to_bytes(),
@@ -1059,22 +1102,30 @@ pub(crate) fn upsert_accounts(
             .set(schema::accounts::is_latest.eq(false))
             .execute(conn)?;
 
-        let account_value = AccountRowInsert {
-            account_id: account_id_bytes,
-            network_account_id_prefix: network_account_id.map(network_account_id_to_prefix_sql),
-            account_commitment: update.final_state_commitment().to_bytes(),
-            block_num: block_num_raw,
-            nonce: full_account.as_ref().map(|account| nonce_to_raw_sql(account.nonce())),
-            code_commitment: full_account
-                .as_ref()
-                .map(|account| account.code().commitment().to_bytes()),
-            // Store only the header (slot metadata + map roots), not full storage with map contents
-            storage_header: full_account
-                .as_ref()
-                .map(|account| account.storage().to_header().to_bytes()),
-            vault_root: full_account.as_ref().map(|account| account.vault().root().to_bytes()),
-            is_latest: true,
-            created_at_block,
+        let account_value = match &account_state {
+            AccountStateForInsert::Private => AccountRowInsert::new_private(
+                account_id_bytes,
+                network_account_id.map(network_account_id_to_prefix_sql),
+                update.final_state_commitment().to_bytes(),
+                block_num_raw,
+                created_at_block,
+            ),
+            AccountStateForInsert::FullAccount(account) => AccountRowInsert::new_from_account(
+                account_id_bytes,
+                network_account_id.map(network_account_id_to_prefix_sql),
+                update.final_state_commitment().to_bytes(),
+                block_num_raw,
+                created_at_block,
+                account,
+            ),
+            AccountStateForInsert::PartialState(state) => AccountRowInsert::new_from_partial(
+                account_id_bytes,
+                network_account_id.map(network_account_id_to_prefix_sql),
+                update.final_state_commitment().to_bytes(),
+                block_num_raw,
+                created_at_block,
+                state,
+            ),
         };
 
         diesel::insert_into(schema::accounts::table)
@@ -1096,25 +1147,6 @@ pub(crate) fn upsert_accounts(
     Ok(count)
 }
 
-/// Deserializes account and applies account delta.
-pub(crate) fn apply_delta(
-    mut account: Account,
-    delta: &AccountDelta,
-    final_state_commitment: &Word,
-) -> crate::db::Result<Account, DatabaseError> {
-    account.apply_delta(delta)?;
-
-    let actual_commitment = account.commitment();
-    if &actual_commitment != final_state_commitment {
-        return Err(DatabaseError::AccountCommitmentsMismatch {
-            calculated: actual_commitment,
-            expected: *final_state_commitment,
-        });
-    }
-
-    Ok(account)
-}
-
 #[derive(Insertable, Debug, Clone)]
 #[diesel(table_name = schema::account_codes)]
 pub(crate) struct AccountCodeRowInsert {
@@ -1137,6 +1169,76 @@ pub(crate) struct AccountRowInsert {
     pub(crate) created_at_block: i64,
 }
 
+impl AccountRowInsert {
+    /// Creates an insert row for a private account (no public state).
+    fn new_private(
+        account_id: Vec<u8>,
+        network_account_id_prefix: Option<i64>,
+        account_commitment: Vec<u8>,
+        block_num: i64,
+        created_at_block: i64,
+    ) -> Self {
+        Self {
+            account_id,
+            network_account_id_prefix,
+            account_commitment,
+            block_num,
+            nonce: None,
+            code_commitment: None,
+            storage_header: None,
+            vault_root: None,
+            is_latest: true,
+            created_at_block,
+        }
+    }
+
+    /// Creates an insert row from a full account (new account creation).
+    fn new_from_account(
+        account_id: Vec<u8>,
+        network_account_id_prefix: Option<i64>,
+        account_commitment: Vec<u8>,
+        block_num: i64,
+        created_at_block: i64,
+        account: &Account,
+    ) -> Self {
+        Self {
+            account_id,
+            network_account_id_prefix,
+            account_commitment,
+            block_num,
+            nonce: Some(nonce_to_raw_sql(account.nonce())),
+            code_commitment: Some(account.code().commitment().to_bytes()),
+            storage_header: Some(account.storage().to_header().to_bytes()),
+            vault_root: Some(account.vault().root().to_bytes()),
+            is_latest: true,
+            created_at_block,
+        }
+    }
+
+    /// Creates an insert row from a partial account state (delta update).
+    fn new_from_partial(
+        account_id: Vec<u8>,
+        network_account_id_prefix: Option<i64>,
+        account_commitment: Vec<u8>,
+        block_num: i64,
+        created_at_block: i64,
+        state: &PartialAccountState,
+    ) -> Self {
+        Self {
+            account_id,
+            network_account_id_prefix,
+            account_commitment,
+            block_num,
+            nonce: Some(nonce_to_raw_sql(state.nonce)),
+            code_commitment: Some(state.code_commitment.to_bytes()),
+            storage_header: Some(state.storage_header.to_bytes()),
+            vault_root: Some(state.vault_root.to_bytes()),
+            is_latest: true,
+            created_at_block,
+        }
+    }
+}
+
 #[derive(Insertable, AsChangeset, Debug, Clone)]
 #[diesel(table_name = schema::account_vault_assets)]
 pub(crate) struct AccountAssetRowInsert {