fix(store): query only latest vault rows when applying account delta (#1789)

Author: TomasArrachea

CHANGELOG.md

@@ -117,7 +117,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 of account updates ([#1567](https://github.com/0xMiden/miden-node/pull/1567)).
+- Improve speed of account updates ([#1567](https://github.com/0xMiden/miden-node/pull/1567), [#1789](https://github.com/0xMiden/node/pull/1789)).
 - Ensure store terminates on nullifier tree or account tree root vs header mismatch (#[#1569](https://github.com/0xMiden/miden-node/pull/1569)).
 - Added support for foreign accounts to `NtxDataStore` and add `GetAccount` endpoint to NTX Builder gRPC store client ([#1521](https://github.com/0xMiden/miden-node/pull/1521)).
 - Use paged queries for tree rebuilding to reduce memory usage during startup ([#1536](https://github.com/0xMiden/miden-node/pull/1536)).

crates/store/src/account_state_forest/mod.rs

@@ -600,7 +600,7 @@ impl AccountStateForest {
         slot_name: &StorageSlotName,
     ) -> Word {
         let lineage = Self::storage_lineage_id(account_id, slot_name);
-        self.forest.latest_root(lineage).map_or_else(Self::empty_smt_root, |root| root)
+        self.forest.latest_root(lineage).unwrap_or_else(Self::empty_smt_root)
     }
 
     /// Updates the forest with storage map changes from a delta.

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

@@ -57,6 +57,7 @@ use delta::{
     AccountStateForInsert,
     PartialAccountState,
     apply_storage_delta,
+    select_latest_vault_assets,
     select_minimal_account_state_headers,
     select_vault_balances_by_faucet_ids,
 };
@@ -1071,7 +1072,6 @@ fn prepare_partial_account_update(
     update: &BlockAccountUpdate,
     account_id: AccountId,
     delta: &miden_protocol::account::delta::AccountDelta,
-    block_num: BlockNumber,
 ) -> Result<(AccountStateForInsert, PendingStorageInserts, PendingAssetInserts), DatabaseError> {
     // Build the minimal account state needed for partial delta application.
     // Only load the storage map entries and vault balances that will receive updates.
@@ -1141,9 +1141,7 @@ fn prepare_partial_account_update(
 
     // --- Update the vault root by constructing the asset vault from DB.
     let new_vault_root = {
-        let (_last_block, assets) =
-            select_account_vault_assets(conn, account_id, BlockNumber::GENESIS..=block_num)?;
-        let assets: Vec<Asset> = assets.into_iter().filter_map(|entry| entry.asset).collect();
+        let assets = select_latest_vault_assets(conn, account_id)?;
         let mut vault = AssetVault::new(&assets)?;
         vault.apply_delta(delta.vault())?;
         vault.root()
@@ -1243,7 +1241,7 @@ pub(crate) fn upsert_accounts(
 
             // Update of an existing account
             AccountUpdateDetails::Delta(delta) => {
-                prepare_partial_account_update(conn, update, account_id, delta, block_num)?
+                prepare_partial_account_update(conn, update, account_id, delta)?
             },
         };
 

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

@@ -189,6 +189,36 @@ pub(super) fn select_vault_balances_by_faucet_ids(
     Ok(balances)
 }
 
+/// Selects the latest vault assets for an account.
+///
+/// # Raw SQL
+///
+/// ```sql
+/// SELECT vault_key, asset
+/// FROM account_vault_assets
+/// WHERE account_id = ?1 AND is_latest = 1
+/// ```
+pub(super) fn select_latest_vault_assets(
+    conn: &mut SqliteConnection,
+    account_id: AccountId,
+) -> Result<Vec<Asset>, DatabaseError> {
+    use schema::account_vault_assets as vault;
+
+    let entries: Vec<(Vec<u8>, Option<Vec<u8>>)> =
+        SelectDsl::select(vault::table, (vault::vault_key, vault::asset))
+            .filter(vault::account_id.eq(account_id.to_bytes()))
+            .filter(vault::is_latest.eq(true))
+            .load(conn)?;
+
+    entries
+        .into_iter()
+        .filter_map(|(_vault_key_bytes, maybe_asset_bytes)| {
+            maybe_asset_bytes.map(|bytes| Asset::read_from_bytes(&bytes))
+        })
+        .collect::<Result<Vec<_>, _>>()
+        .map_err(Into::into)
+}
+
 // HELPER FUNCTIONS
 // ================================================================================================
 

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

@@ -311,13 +311,19 @@ fn optimized_delta_matches_full_account_method() {
 }
 
 #[test]
+#[expect(
+    clippy::too_many_lines,
+    reason = "test exercises vault deltas across multiple blocks"
+)]
 fn optimized_delta_updates_non_empty_vault() {
     const ACCOUNT_SEED: [u8; 32] = [40u8; 32];
     const BLOCK_NUM_1: u32 = 1;
     const BLOCK_NUM_2: u32 = 2;
+    const BLOCK_NUM_3: u32 = 3;
     const NONCE_DELTA: u64 = 1;
     const INITIAL_AMOUNT: u64 = 700;
-    const ADDED_AMOUNT: u64 = 250;
+    const ADDED_AMOUNT_BLOCK_2: u64 = 250;
+    const ADDED_AMOUNT_BLOCK_3: u64 = 150;
     const SLOT_INDEX: usize = 0;
 
     let mut conn = setup_test_db();
@@ -355,9 +361,12 @@ fn optimized_delta_updates_non_empty_vault() {
 
     let block_1 = BlockNumber::from(BLOCK_NUM_1);
     let block_2 = BlockNumber::from(BLOCK_NUM_2);
+    let block_3 = BlockNumber::from(BLOCK_NUM_3);
     insert_block_header(&mut conn, block_1);
     insert_block_header(&mut conn, block_2);
+    insert_block_header(&mut conn, block_3);
 
+    // Block 1: insert full-state delta (initial account with 700 tokens of faucet_id)
     let delta_initial = AccountDelta::try_from(account.clone()).unwrap();
     let account_update_initial = BlockAccountUpdate::new(
         account.id(),
@@ -369,9 +378,10 @@ fn optimized_delta_updates_non_empty_vault() {
     let full_account_before =
         select_full_account(&mut conn, account.id()).expect("Failed to load full account");
 
+    // Block 2: partial delta — remove faucet_id (700), add faucet_id_1 (250)
     let mut vault_delta = AccountVaultDelta::default();
     vault_delta
-        .add_asset(Asset::Fungible(FungibleAsset::new(faucet_id_1, ADDED_AMOUNT).unwrap()))
+        .add_asset(Asset::Fungible(FungibleAsset::new(faucet_id_1, ADDED_AMOUNT_BLOCK_2).unwrap()))
         .unwrap();
     vault_delta
         .remove_asset(Asset::Fungible(FungibleAsset::new(faucet_id, INITIAL_AMOUNT).unwrap()))
@@ -403,14 +413,53 @@ fn optimized_delta_updates_non_empty_vault() {
     assert_eq!(vault_assets_after.len(), 1, "Should have 1 vault asset");
     assert_matches!(&vault_assets_after[0], Asset::Fungible(f) => {
         assert_eq!(f.faucet_id(), faucet_id_1, "Faucet ID should match");
-        assert_eq!(f.amount(), ADDED_AMOUNT, "Amount should match");
+        assert_eq!(f.amount(), ADDED_AMOUNT_BLOCK_2, "Amount should match");
     });
 
     let full_account_after = select_full_account(&mut conn, account.id())
         .expect("Failed to load full account after update");
 
     assert_eq!(full_account_after.vault().root(), expected_vault_root);
     assert_eq!(full_account_after.to_commitment(), expected_commitment);
+
+    // Block 3: partial delta — add more of faucet_id_1 (150 more, total = 400)
+    let mut vault_delta_3 = AccountVaultDelta::default();
+    vault_delta_3
+        .add_asset(Asset::Fungible(FungibleAsset::new(faucet_id_1, ADDED_AMOUNT_BLOCK_3).unwrap()))
+        .unwrap();
+
+    let partial_delta_3 = AccountDelta::new(
+        account.id(),
+        AccountStorageDelta::new(),
+        vault_delta_3,
+        Felt::new(NONCE_DELTA),
+    )
+    .unwrap();
+
+    let mut expected_after_3 = full_account_after.clone();
+    expected_after_3.apply_delta(&partial_delta_3).unwrap();
+    let commitment_3 = expected_after_3.to_commitment();
+    let expected_vault_root_3 = expected_after_3.vault().root();
+
+    let account_update_3 = BlockAccountUpdate::new(
+        account.id(),
+        commitment_3,
+        AccountUpdateDetails::Delta(partial_delta_3),
+    );
+    upsert_accounts(&mut conn, &[account_update_3], block_3).expect("Block 3 upsert failed");
+
+    let full_account_final =
+        select_full_account(&mut conn, account.id()).expect("Failed to load after block 3");
+
+    let final_assets: Vec<Asset> = full_account_final.vault().assets().collect();
+    assert_eq!(final_assets.len(), 1, "Should have exactly 1 vault asset");
+    assert_matches!(&final_assets[0], Asset::Fungible(f) => {
+        assert_eq!(f.faucet_id(), faucet_id_1);
+        assert_eq!(f.amount(), ADDED_AMOUNT_BLOCK_2 + ADDED_AMOUNT_BLOCK_3, "Expected total of 400");
+    });
+
+    assert_eq!(full_account_final.vault().root(), expected_vault_root_3);
+    assert_eq!(full_account_final.to_commitment(), commitment_3);
 }
 
 #[test]