diff --git a/CHANGELOG.md b/CHANGELOG.md index dbf050f30..a782c68c8 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/alt_approach.txt b/alt_approach.txt new file mode 100644 index 000000000..4027ca33f --- /dev/null +++ b/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 | diff --git a/crates/store/src/db/models/queries/accounts.rs b/crates/store/src/db/models/queries/accounts.rs index 5c049916e..ae10e83b0 100644 --- a/crates/store/src/db/models/queries/accounts.rs +++ b/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 { @@ -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 = + 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,24 +1102,49 @@ 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, + ), }; + if let AccountStateForInsert::PartialState(state) = &account_state { + let account_header = miden_protocol::account::AccountHeader::new( + account_id, + state.nonce, + state.vault_root, + state.storage_header.to_commitment(), + state.code_commitment, + ); + + if account_header.commitment() != update.final_state_commitment() { + return Err(DatabaseError::AccountCommitmentsMismatch { + calculated: account_header.commitment(), + expected: update.final_state_commitment(), + }); + } + } + diesel::insert_into(schema::accounts::table) .values(&account_value) .execute(conn)?; @@ -1096,25 +1164,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.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 +1186,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, + network_account_id_prefix: Option, + account_commitment: Vec, + 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, + network_account_id_prefix: Option, + account_commitment: Vec, + 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, + network_account_id_prefix: Option, + account_commitment: Vec, + 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 { diff --git a/crates/store/src/db/models/queries/accounts/delta.rs b/crates/store/src/db/models/queries/accounts/delta.rs new file mode 100644 index 000000000..e8e4a4342 --- /dev/null +++ b/crates/store/src/db/models/queries/accounts/delta.rs @@ -0,0 +1,326 @@ +//! Optimized delta update support for account updates. +//! +//! Provides functions and types for applying partial delta updates to accounts +//! without loading the full account state. Avoids loading: +//! - Full account code bytes +//! - All storage map entries +//! - All vault assets +//! +//! Instead, only the minimal data needed for the update is fetched. + +use std::collections::BTreeMap; + +use diesel::query_dsl::methods::SelectDsl; +use diesel::{ExpressionMethods, OptionalExtension, QueryDsl, RunQueryDsl, SqliteConnection}; +use miden_protocol::account::delta::{AccountStorageDelta, AccountVaultDelta}; +use miden_protocol::account::{ + Account, + AccountId, + AccountStorageHeader, + NonFungibleDeltaAction, + StorageSlotName, +}; +use miden_protocol::asset::{Asset, FungibleAsset}; +use miden_protocol::utils::{Deserializable, Serializable}; +use miden_protocol::{Felt, Word}; + +use crate::db::models::conv::raw_sql_to_nonce; +use crate::db::schema; +use crate::errors::DatabaseError; + +#[cfg(test)] +mod tests; + +// TYPES +// ================================================================================================ + +/// Raw row type for account state delta queries. +/// +/// Fields: (`nonce`, `code_commitment`, `storage_header`, `vault_root`) +#[derive(diesel::prelude::Queryable)] +struct AccountStateDeltaRow { + nonce: Option, + code_commitment: Option>, + storage_header: Option>, + vault_root: Option>, +} + +/// Data needed for applying a delta update to an existing account. +/// Fetches only the minimal data required, avoiding loading full code and storage. +#[derive(Debug, Clone)] +pub(super) struct AccountStateForDelta { + pub nonce: Felt, + pub code_commitment: Word, + pub storage_header: AccountStorageHeader, + pub vault_root: Word, +} + +/// Minimal account state computed from a partial delta update. +/// Contains only the fields needed for the accounts table row insert. +#[derive(Debug, Clone)] +pub(super) struct PartialAccountState { + pub nonce: Felt, + pub code_commitment: Word, + pub storage_header: AccountStorageHeader, + pub vault_root: Word, +} + +/// Represents the account state to be inserted, either from a full account +/// or from a partial delta update. +pub(super) enum AccountStateForInsert { + /// Private account - no public state stored + Private, + /// Full account state (from full-state delta, i.e., new account) + FullAccount(Account), + /// Partial account state (from partial delta, i.e., existing account update) + PartialState(PartialAccountState), +} + +// QUERIES +// ================================================================================================ + +/// Selects the minimal account state needed for applying a delta update. +/// +/// Optimized query that only fetches: +/// - `nonce` (to add `nonce_delta`) +/// - `code_commitment` (unchanged in partial deltas) +/// - `storage_header` (to apply storage delta) +/// - `vault_root` (to apply vault delta) +/// +/// # Raw SQL +/// +/// ```sql +/// SELECT nonce, code_commitment, storage_header, vault_root +/// FROM accounts +/// WHERE account_id = ?1 AND is_latest = 1 +/// ``` +pub(super) fn select_account_state_for_delta( + conn: &mut SqliteConnection, + account_id: AccountId, +) -> Result { + let row: AccountStateDeltaRow = SelectDsl::select( + schema::accounts::table, + ( + schema::accounts::nonce, + schema::accounts::code_commitment, + schema::accounts::storage_header, + schema::accounts::vault_root, + ), + ) + .filter(schema::accounts::account_id.eq(account_id.to_bytes())) + .filter(schema::accounts::is_latest.eq(true)) + .get_result(conn) + .optional()? + .ok_or(DatabaseError::AccountNotFoundInDb(account_id))?; + + let nonce = raw_sql_to_nonce(row.nonce.ok_or_else(|| { + DatabaseError::DataCorrupted(format!("No nonce found for account {account_id}")) + })?); + + let code_commitment = row + .code_commitment + .map(|bytes| Word::read_from_bytes(&bytes)) + .transpose()? + .ok_or_else(|| { + DatabaseError::DataCorrupted(format!( + "No code_commitment found for account {account_id}" + )) + })?; + + let storage_header = match row.storage_header { + Some(bytes) => AccountStorageHeader::read_from_bytes(&bytes)?, + None => AccountStorageHeader::new(Vec::new())?, + }; + + let vault_root = row + .vault_root + .map(|bytes| Word::read_from_bytes(&bytes)) + .transpose()? + .unwrap_or(Word::default()); + + Ok(AccountStateForDelta { + nonce, + code_commitment, + storage_header, + vault_root, + }) +} + +/// Selects vault balances for specific faucet IDs. +/// +/// Optimized query that only fetches balances for the faucet IDs +/// that are being updated by a delta, rather than loading all vault assets. +/// +/// Returns a map from `faucet_id` to the current balance (0 if not found). +/// +/// # Raw SQL +/// +/// ```sql +/// SELECT vault_key, asset +/// FROM account_vault_assets +/// WHERE account_id = ?1 AND is_latest = 1 AND vault_key IN (?2, ?3, ...) +/// ``` +pub(super) fn select_vault_balances_by_faucet_ids( + conn: &mut SqliteConnection, + account_id: AccountId, + faucet_ids: &[AccountId], +) -> Result, DatabaseError> { + use schema::account_vault_assets as vault; + + if faucet_ids.is_empty() { + return Ok(BTreeMap::new()); + } + + let account_id_bytes = account_id.to_bytes(); + + // Compute vault keys for each faucet ID + let vault_keys: Vec> = Result::from_iter(faucet_ids.iter().map(|faucet_id| { + let asset = FungibleAsset::new(*faucet_id, 0) + .map_err(|_| DatabaseError::DataCorrupted(format!("Invalid faucet id {faucet_id}")))?; + let key: Word = asset.vault_key().into(); + Ok::<_, DatabaseError>(key.to_bytes()) + }))?; + + let entries: Vec<(Vec, Option>)> = + SelectDsl::select(vault::table, (vault::vault_key, vault::asset)) + .filter(vault::account_id.eq(&account_id_bytes)) + .filter(vault::is_latest.eq(true)) + .filter(vault::vault_key.eq_any(&vault_keys)) + .load(conn)?; + + let mut balances = BTreeMap::new(); + + for (_vault_key_bytes, maybe_asset_bytes) in entries { + if let Some(asset_bytes) = maybe_asset_bytes { + let asset = Asset::read_from_bytes(&asset_bytes)?; + if let Asset::Fungible(fungible) = asset { + balances.insert(fungible.faucet_id(), fungible.amount()); + } + } + } + + Ok(balances) +} + +// HELPER FUNCTIONS +// ================================================================================================ + +/// Applies storage delta to an existing storage header to produce the new header. +/// +/// For value slots, updates the slot value directly. +/// For map slots, computes the new SMT root using batch insert operations. +pub(super) fn apply_storage_delta_to_header( + header: &AccountStorageHeader, + delta: &AccountStorageDelta, +) -> Result { + use miden_protocol::account::StorageSlotHeader; + use miden_protocol::crypto::merkle::smt::SmtForest; + + // Build a map of slot updates from the delta + let mut value_updates: BTreeMap<&StorageSlotName, Word> = BTreeMap::new(); + let mut map_updates: BTreeMap<&StorageSlotName, Word> = BTreeMap::new(); + + // Collect value updates + for (slot_name, new_value) in delta.values() { + value_updates.insert(slot_name, *new_value); + } + + // Compute new map roots from map deltas + let mut forest = SmtForest::new(); + for (slot_name, map_delta) in delta.maps() { + // Find the previous root from the header + let prev_root = header + .slots() + .find(|s| s.name() == slot_name) + .map(StorageSlotHeader::value) + .unwrap_or_default(); + + let entries = + Vec::from_iter(map_delta.entries().iter().map(|(key, value)| ((*key).into(), *value))); + + if !entries.is_empty() { + let new_root = + forest.batch_insert(prev_root, entries.iter().copied()).map_err(|e| { + DatabaseError::DataCorrupted(format!( + "Failed to compute storage map root: {e:?}" + )) + })?; + map_updates.insert(slot_name, new_root); + } + } + + // Create new slots with updates applied + let new_slots = Vec::from_iter(header.slots().map(|slot| { + let slot_name = slot.name(); + if let Some(&new_value) = value_updates.get(slot_name) { + StorageSlotHeader::new(slot_name.clone(), slot.slot_type(), new_value) + } else if let Some(&new_root) = map_updates.get(slot_name) { + StorageSlotHeader::new(slot_name.clone(), slot.slot_type(), new_root) + } else { + slot.clone() + } + })); + + AccountStorageHeader::new(new_slots).map_err(|e| { + DatabaseError::DataCorrupted(format!("Failed to create storage header: {e:?}")) + }) +} + +/// Computes the new vault root after applying a vault delta. +/// +/// Uses SMT operations to incrementally update the vault root from the +/// previous root and the delta entries. +pub(super) fn compute_vault_root_after_delta( + prev_root: Word, + vault_delta: &AccountVaultDelta, + prev_balances: &BTreeMap, +) -> Result { + use miden_protocol::EMPTY_WORD; + use miden_protocol::crypto::merkle::smt::SmtForest; + + if vault_delta.is_empty() { + return Ok(prev_root); + } + + let mut entries: Vec<(Word, Word)> = Vec::new(); + let mut forest = SmtForest::new(); + + // Process fungible asset changes + for (faucet_id, amount_delta) in vault_delta.fungible().iter() { + let prev_balance = prev_balances.get(faucet_id).copied().unwrap_or(0); + let new_balance: u64 = + (i128::from(prev_balance) + i128::from(*amount_delta)).try_into().map_err(|_| { + DatabaseError::DataCorrupted(format!("Balance underflow for faucet {faucet_id}")) + })?; + + let key: Word = + FungibleAsset::new(*faucet_id, 0).expect("valid faucet id").vault_key().into(); + + let value = if new_balance == 0 { + EMPTY_WORD + } else { + let asset: Asset = FungibleAsset::new(*faucet_id, new_balance) + .expect("valid fungible asset") + .into(); + Word::from(asset) + }; + + entries.push((key, value)); + } + + // Process non-fungible asset changes + for (asset, action) in vault_delta.non_fungible().iter() { + let key: Word = asset.vault_key().into(); + let value = match action { + NonFungibleDeltaAction::Add => Word::from(Asset::NonFungible(*asset)), + NonFungibleDeltaAction::Remove => EMPTY_WORD, + }; + entries.push((key, value)); + } + + let new_root = forest.batch_insert(prev_root, entries.iter().copied()).map_err(|e| { + DatabaseError::DataCorrupted(format!("Failed to compute vault root: {e:?}")) + })?; + + Ok(new_root) +} diff --git a/crates/store/src/db/models/queries/accounts/delta/tests.rs b/crates/store/src/db/models/queries/accounts/delta/tests.rs new file mode 100644 index 000000000..341b6b165 --- /dev/null +++ b/crates/store/src/db/models/queries/accounts/delta/tests.rs @@ -0,0 +1,556 @@ +//! +//! Tests for delta update functionality. + +use std::collections::BTreeMap; + +use assert_matches::assert_matches; +use diesel::{Connection, ExpressionMethods, QueryDsl, RunQueryDsl, SqliteConnection}; +use diesel_migrations::MigrationHarness; +use miden_node_utils::fee::test_fee_params; +use miden_protocol::account::auth::PublicKeyCommitment; +use miden_protocol::account::delta::{ + AccountStorageDelta, + AccountUpdateDetails, + AccountVaultDelta, + StorageMapDelta, + StorageSlotDelta, +}; +use miden_protocol::account::{ + AccountBuilder, + AccountComponent, + AccountDelta, + AccountId, + AccountStorageMode, + AccountType, + StorageMap, + StorageSlot, + StorageSlotName, +}; +use miden_protocol::asset::{Asset, FungibleAsset}; +use miden_protocol::block::{BlockAccountUpdate, BlockHeader, BlockNumber}; +use miden_protocol::crypto::dsa::ecdsa_k256_keccak::SecretKey; +use miden_protocol::testing::account_id::ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET; +use miden_protocol::utils::Serializable; +use miden_protocol::{EMPTY_WORD, Felt, Word}; +use miden_standards::account::auth::AuthFalcon512Rpo; +use miden_standards::code_builder::CodeBuilder; + +use crate::db::migrations::MIGRATIONS; +use crate::db::models::queries::accounts::{ + select_account_header_with_storage_header_at_block, + select_account_vault_at_block, + select_full_account, + upsert_accounts, +}; +use crate::db::schema::accounts; + +fn setup_test_db() -> SqliteConnection { + let mut conn = + SqliteConnection::establish(":memory:").expect("Failed to create in-memory database"); + + conn.run_pending_migrations(MIGRATIONS).expect("Failed to run migrations"); + + conn +} + +fn insert_block_header(conn: &mut SqliteConnection, block_num: BlockNumber) { + use crate::db::schema::block_headers; + + let block_header = BlockHeader::new( + 1_u8.into(), + Word::default(), + block_num, + Word::default(), + Word::default(), + Word::default(), + Word::default(), + Word::default(), + Word::default(), + SecretKey::new().public_key(), + test_fee_params(), + 0_u8.into(), + ); + + diesel::insert_into(block_headers::table) + .values(( + block_headers::block_num.eq(i64::from(block_num.as_u32())), + block_headers::block_header.eq(block_header.to_bytes()), + )) + .execute(conn) + .expect("Failed to insert block header"); +} + +/// Tests that the optimized delta update path produces the same results as the old +/// method that loads the full account. +/// +/// Covers partial deltas that update: +/// - Nonce (via `nonce_delta`) +/// - Value storage slots +/// - Vault assets (fungible) starting from empty vault +/// +/// The test ensures the optimized code path in `upsert_accounts` produces correct results +/// by comparing the final account state against a manually constructed expected state. +#[test] +#[allow(clippy::too_many_lines)] +fn optimized_delta_matches_full_account_method() { + // Use deterministic account seed to keep account IDs stable. + const ACCOUNT_SEED: [u8; 32] = [10u8; 32]; + // Use fixed block numbers to ensure deterministic ordering. + const BLOCK_NUM_1: u32 = 1; + const BLOCK_NUM_2: u32 = 2; + // Use explicit slot indices to avoid magic numbers. + const SLOT_INDEX_PRIMARY: u8 = 0; + const SLOT_INDEX_SECONDARY: u8 = 1; + // Use fixed values to verify storage delta updates. + const INITIAL_SLOT_VALUES: [u64; 4] = [100, 200, 300, 400]; + const UPDATED_SLOT_VALUES: [u64; 4] = [111, 222, 333, 444]; + // Use fixed delta values to validate nonce and vault changes. + const NONCE_DELTA: u64 = 5; + const VAULT_AMOUNT: u64 = 500; + + let mut conn = setup_test_db(); + + // Create an account with value slots only (no map slots to avoid SmtForest complexity) + let slot_value_initial = Word::from([ + Felt::new(INITIAL_SLOT_VALUES[0]), + Felt::new(INITIAL_SLOT_VALUES[1]), + Felt::new(INITIAL_SLOT_VALUES[2]), + Felt::new(INITIAL_SLOT_VALUES[3]), + ]); + + let component_storage = vec![ + StorageSlot::with_value(StorageSlotName::mock(SLOT_INDEX_PRIMARY), slot_value_initial), + StorageSlot::with_value(StorageSlotName::mock(SLOT_INDEX_SECONDARY), EMPTY_WORD), + ]; + + let account_component_code = CodeBuilder::default() + .compile_component_code("test::interface", "pub proc foo push.1 end") + .unwrap(); + + let component = AccountComponent::new(account_component_code, component_storage) + .unwrap() + .with_supported_type(AccountType::RegularAccountImmutableCode); + + let account = AccountBuilder::new(ACCOUNT_SEED) + .account_type(AccountType::RegularAccountImmutableCode) + .storage_mode(AccountStorageMode::Public) + .with_component(component) + .with_auth_component(AuthFalcon512Rpo::new(PublicKeyCommitment::from(EMPTY_WORD))) + .build_existing() + .unwrap(); + + let block_1 = BlockNumber::from(BLOCK_NUM_1); + let block_2 = BlockNumber::from(BLOCK_NUM_2); + insert_block_header(&mut conn, block_1); + insert_block_header(&mut conn, block_2); + + // Insert the initial account at block 1 (full state) - no vault assets + let delta_initial = AccountDelta::try_from(account.clone()).unwrap(); + let account_update_initial = BlockAccountUpdate::new( + account.id(), + account.commitment(), + AccountUpdateDetails::Delta(delta_initial), + ); + upsert_accounts(&mut conn, &[account_update_initial], block_1).expect("Initial upsert failed"); + + // Verify initial state + let full_account_before = + select_full_account(&mut conn, account.id()).expect("Failed to load full account"); + assert_eq!(full_account_before.nonce(), account.nonce()); + assert!( + full_account_before.vault().assets().next().is_none(), + "Vault should be empty initially" + ); + + // Create a partial delta to apply: + // - Increment nonce by 5 + // - Update the first value slot + // - Add 500 tokens to the vault (starting from empty) + + let new_slot_value = Word::from([ + Felt::new(UPDATED_SLOT_VALUES[0]), + Felt::new(UPDATED_SLOT_VALUES[1]), + Felt::new(UPDATED_SLOT_VALUES[2]), + Felt::new(UPDATED_SLOT_VALUES[3]), + ]); + let faucet_id = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap(); + + // Find the slot name from the account's storage + let value_slot_name = + full_account_before.storage().slots().iter().next().unwrap().name().clone(); + + // Build the storage delta (value slot update only) + let storage_delta = { + let deltas = BTreeMap::from_iter([( + value_slot_name.clone(), + StorageSlotDelta::Value(new_slot_value), + )]); + AccountStorageDelta::from_raw(deltas) + }; + + // Build the vault delta (add 500 tokens to empty vault) + let vault_delta = { + let mut delta = AccountVaultDelta::default(); + let asset = Asset::Fungible(FungibleAsset::new(faucet_id, VAULT_AMOUNT).unwrap()); + delta.add_asset(asset).unwrap(); + delta + }; + + // Create a partial delta + let nonce_delta = Felt::new(NONCE_DELTA); + let partial_delta = AccountDelta::new( + full_account_before.id(), + storage_delta.clone(), + vault_delta.clone(), + nonce_delta, + ) + .unwrap(); + assert!(!partial_delta.is_full_state(), "Delta should be partial, not full state"); + + // Construct the expected final account by applying the delta + let expected_nonce = Felt::new(full_account_before.nonce().as_int() + nonce_delta.as_int()); + let expected_code_commitment = full_account_before.code().commitment(); + + let mut expected_account = full_account_before.clone(); + expected_account.apply_delta(&partial_delta).unwrap(); + let final_account_for_commitment = expected_account; + + let final_commitment = final_account_for_commitment.commitment(); + let expected_storage_commitment = final_account_for_commitment.storage().to_commitment(); + let expected_vault_root = final_account_for_commitment.vault().root(); + + // ----- Apply the partial delta via upsert_accounts (optimized path) ----- + let account_update = BlockAccountUpdate::new( + account.id(), + final_commitment, + AccountUpdateDetails::Delta(partial_delta), + ); + upsert_accounts(&mut conn, &[account_update], block_2).expect("Partial delta upsert failed"); + + // ----- VERIFY: Query the DB and check that optimized path produced correct results ----- + + let (header_after, storage_header_after) = + select_account_header_with_storage_header_at_block(&mut conn, account.id(), block_2) + .expect("Query should succeed") + .expect("Account should exist"); + + // Verify nonce + assert_eq!( + header_after.nonce(), + expected_nonce, + "Nonce mismatch: optimized={:?}, expected={:?}", + header_after.nonce(), + expected_nonce + ); + + // Verify code commitment (should be unchanged) + assert_eq!( + header_after.code_commitment(), + expected_code_commitment, + "Code commitment mismatch" + ); + + // Verify storage header commitment + assert_eq!( + storage_header_after.to_commitment(), + expected_storage_commitment, + "Storage header commitment mismatch" + ); + + // Verify vault assets + let vault_assets_after = select_account_vault_at_block(&mut conn, account.id(), block_2) + .expect("Query vault should succeed"); + + 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, "Faucet ID should match"); + assert_eq!(f.amount(), VAULT_AMOUNT, "Amount should be 500"); + }); + + // Verify the account commitment matches + assert_eq!( + header_after.commitment(), + final_commitment, + "Account commitment should match the expected final state" + ); + + // Also verify we can load the full account and it has correct state + let full_account_after = select_full_account(&mut conn, account.id()) + .expect("Failed to load full account after update"); + + assert_eq!(full_account_after.nonce(), expected_nonce, "Full account nonce mismatch"); + assert_eq!( + full_account_after.storage().to_commitment(), + expected_storage_commitment, + "Full account storage commitment mismatch" + ); + assert_eq!( + full_account_after.vault().root(), + expected_vault_root, + "Full account vault root mismatch" + ); +} + +#[test] +fn optimized_delta_updates_storage_map_header() { + // Use deterministic account seed to keep account IDs stable. + const ACCOUNT_SEED: [u8; 32] = [30u8; 32]; + // Use fixed block numbers to ensure deterministic ordering. + const BLOCK_NUM_1: u32 = 1; + const BLOCK_NUM_2: u32 = 2; + // Use explicit slot index to avoid magic numbers. + const SLOT_INDEX_MAP: u8 = 3; + // Use fixed map values to validate root updates. + const MAP_KEY_VALUES: [u64; 4] = [7, 0, 0, 0]; + const MAP_VALUE_INITIAL: [u64; 4] = [10, 20, 30, 40]; + const MAP_VALUE_UPDATED: [u64; 4] = [50, 60, 70, 80]; + // Use zero nonce delta to isolate storage updates. + const NONCE_DELTA: u64 = 0; + + let mut conn = setup_test_db(); + + let map_key = Word::from([ + Felt::new(MAP_KEY_VALUES[0]), + Felt::new(MAP_KEY_VALUES[1]), + Felt::new(MAP_KEY_VALUES[2]), + Felt::new(MAP_KEY_VALUES[3]), + ]); + let map_value_initial = Word::from([ + Felt::new(MAP_VALUE_INITIAL[0]), + Felt::new(MAP_VALUE_INITIAL[1]), + Felt::new(MAP_VALUE_INITIAL[2]), + Felt::new(MAP_VALUE_INITIAL[3]), + ]); + let map_value_updated = Word::from([ + Felt::new(MAP_VALUE_UPDATED[0]), + Felt::new(MAP_VALUE_UPDATED[1]), + Felt::new(MAP_VALUE_UPDATED[2]), + Felt::new(MAP_VALUE_UPDATED[3]), + ]); + + let storage_map = StorageMap::with_entries(vec![(map_key, map_value_initial)]).unwrap(); + let component_storage = + vec![StorageSlot::with_map(StorageSlotName::mock(SLOT_INDEX_MAP), storage_map)]; + + let account_component_code = CodeBuilder::default() + .compile_component_code("test::interface", "pub proc map push.1 end") + .unwrap(); + + let component = AccountComponent::new(account_component_code, component_storage) + .unwrap() + .with_supported_type(AccountType::RegularAccountImmutableCode); + + let account = AccountBuilder::new(ACCOUNT_SEED) + .account_type(AccountType::RegularAccountImmutableCode) + .storage_mode(AccountStorageMode::Public) + .with_component(component) + .with_auth_component(AuthFalcon512Rpo::new(PublicKeyCommitment::from(EMPTY_WORD))) + .build_existing() + .unwrap(); + + let block_1 = BlockNumber::from(BLOCK_NUM_1); + let block_2 = BlockNumber::from(BLOCK_NUM_2); + insert_block_header(&mut conn, block_1); + insert_block_header(&mut conn, block_2); + + let delta_initial = AccountDelta::try_from(account.clone()).unwrap(); + let account_update_initial = BlockAccountUpdate::new( + account.id(), + account.commitment(), + AccountUpdateDetails::Delta(delta_initial), + ); + upsert_accounts(&mut conn, &[account_update_initial], block_1).expect("Initial upsert failed"); + + let full_account_before = + select_full_account(&mut conn, account.id()).expect("Failed to load full account"); + + let mut map_delta = StorageMapDelta::default(); + map_delta.insert(map_key, map_value_updated); + let storage_delta = AccountStorageDelta::from_raw(BTreeMap::from_iter([( + StorageSlotName::mock(SLOT_INDEX_MAP), + StorageSlotDelta::Map(map_delta), + )])); + + let partial_delta = AccountDelta::new( + account.id(), + storage_delta, + AccountVaultDelta::default(), + Felt::new(NONCE_DELTA), + ) + .unwrap(); + + let mut expected_account = full_account_before.clone(); + expected_account.apply_delta(&partial_delta).unwrap(); + let expected_commitment = expected_account.commitment(); + let expected_storage_commitment = expected_account.storage().to_commitment(); + + let account_update = BlockAccountUpdate::new( + account.id(), + expected_commitment, + AccountUpdateDetails::Delta(partial_delta), + ); + upsert_accounts(&mut conn, &[account_update], block_2).expect("Partial delta upsert failed"); + + let (header_after, storage_header_after) = + select_account_header_with_storage_header_at_block(&mut conn, account.id(), block_2) + .expect("Query should succeed") + .expect("Account should exist"); + + assert_eq!( + storage_header_after.to_commitment(), + expected_storage_commitment, + "Storage commitment should match after map delta" + ); + assert_eq!( + header_after.commitment(), + expected_commitment, + "Account commitment should match after map delta" + ); +} + +/// Tests that a private account update (no public state) is handled correctly. +/// +/// Private accounts store only the account commitment, not the full state. +#[test] +fn upsert_private_account() { + use miden_protocol::account::{AccountIdVersion, AccountStorageMode, AccountType}; + + // Use deterministic account seed to keep account IDs stable. + const ACCOUNT_ID_SEED: [u8; 15] = [20u8; 15]; + // Use fixed block number to keep test ordering deterministic. + const BLOCK_NUM: u32 = 1; + // Use fixed commitment values to validate storage behavior. + const COMMITMENT_WORDS: [u64; 4] = [1, 2, 3, 4]; + + let mut conn = setup_test_db(); + + let block_num = BlockNumber::from(BLOCK_NUM); + insert_block_header(&mut conn, block_num); + + // Create a private account ID + let account_id = AccountId::dummy( + ACCOUNT_ID_SEED, + AccountIdVersion::Version0, + AccountType::RegularAccountImmutableCode, + AccountStorageMode::Private, + ); + + let account_commitment = Word::from([ + Felt::new(COMMITMENT_WORDS[0]), + Felt::new(COMMITMENT_WORDS[1]), + Felt::new(COMMITMENT_WORDS[2]), + Felt::new(COMMITMENT_WORDS[3]), + ]); + + // Insert as private account + let account_update = + BlockAccountUpdate::new(account_id, account_commitment, AccountUpdateDetails::Private); + + upsert_accounts(&mut conn, &[account_update], block_num) + .expect("Private account upsert failed"); + + // Verify the account exists and commitment matches + + let (stored_commitment, stored_nonce, stored_code): (Vec, Option, Option>) = + accounts::table + .filter(accounts::account_id.eq(account_id.to_bytes())) + .filter(accounts::is_latest.eq(true)) + .select((accounts::account_commitment, accounts::nonce, accounts::code_commitment)) + .first(&mut conn) + .expect("Account should exist in DB"); + + assert_eq!( + stored_commitment, + account_commitment.to_bytes(), + "Stored commitment should match" + ); + + // Private accounts have NULL for nonce, code_commitment, storage_header, vault_root + assert!(stored_nonce.is_none(), "Private account should have NULL nonce"); + assert!(stored_code.is_none(), "Private account should have NULL code_commitment"); +} + +/// Tests that a full-state delta (new account creation) is handled correctly. +/// +/// Full-state deltas contain the complete account state including code. +#[test] +fn upsert_full_state_delta() { + // Use deterministic account seed to keep account IDs stable. + const ACCOUNT_SEED: [u8; 32] = [20u8; 32]; + // Use fixed block number to keep test ordering deterministic. + const BLOCK_NUM: u32 = 1; + // Use fixed slot values to validate storage behavior. + const SLOT_VALUES: [u64; 4] = [10, 20, 30, 40]; + // Use explicit slot index to avoid magic numbers. + const SLOT_INDEX: u8 = 0; + + let mut conn = setup_test_db(); + + let block_num = BlockNumber::from(BLOCK_NUM); + insert_block_header(&mut conn, block_num); + + // Create an account with storage + let slot_value = Word::from([ + Felt::new(SLOT_VALUES[0]), + Felt::new(SLOT_VALUES[1]), + Felt::new(SLOT_VALUES[2]), + Felt::new(SLOT_VALUES[3]), + ]); + let component_storage = + vec![StorageSlot::with_value(StorageSlotName::mock(SLOT_INDEX), slot_value)]; + + let account_component_code = CodeBuilder::default() + .compile_component_code("test::interface", "pub proc bar push.2 end") + .unwrap(); + + let component = AccountComponent::new(account_component_code, component_storage) + .unwrap() + .with_supported_type(AccountType::RegularAccountImmutableCode); + + let account = AccountBuilder::new(ACCOUNT_SEED) + .account_type(AccountType::RegularAccountImmutableCode) + .storage_mode(AccountStorageMode::Public) + .with_component(component) + .with_auth_component(AuthFalcon512Rpo::new(PublicKeyCommitment::from(EMPTY_WORD))) + .build_existing() + .unwrap(); + + // Create a full-state delta from the account + let delta = AccountDelta::try_from(account.clone()).unwrap(); + assert!(delta.is_full_state(), "Delta should be full state"); + + let account_update = BlockAccountUpdate::new( + account.id(), + account.commitment(), + AccountUpdateDetails::Delta(delta), + ); + + upsert_accounts(&mut conn, &[account_update], block_num) + .expect("Full-state delta upsert failed"); + + // Verify the account state was stored correctly + let (header, storage_header) = + select_account_header_with_storage_header_at_block(&mut conn, account.id(), block_num) + .expect("Query should succeed") + .expect("Account should exist"); + + assert_eq!(header.nonce(), account.nonce(), "Nonce should match"); + assert_eq!( + header.code_commitment(), + account.code().commitment(), + "Code commitment should match" + ); + assert_eq!( + storage_header.to_commitment(), + account.storage().to_commitment(), + "Storage commitment should match" + ); + + // Verify we can load the full account back + let loaded_account = + select_full_account(&mut conn, account.id()).expect("Should load full account"); + + assert_eq!(loaded_account.nonce(), account.nonce()); + assert_eq!(loaded_account.code().commitment(), account.code().commitment()); + assert_eq!(loaded_account.storage().to_commitment(), account.storage().to_commitment()); +}