diff --git a/Cargo.lock b/Cargo.lock index 561e8f8fa..414a690fc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3407,6 +3407,7 @@ dependencies = [ "miden-standards", "rand 0.10.2", "rayon", + "tempfile", "tokio", ] diff --git a/bin/stress-test/Cargo.toml b/bin/stress-test/Cargo.toml index 508642ba2..eb82ce982 100644 --- a/bin/stress-test/Cargo.toml +++ b/bin/stress-test/Cargo.toml @@ -28,3 +28,6 @@ miden-standards = { workspace = true } rand = { workspace = true } rayon = { workspace = true } tokio = { workspace = true } + +[dev-dependencies] +tempfile = { workspace = true } diff --git a/bin/stress-test/README.md b/bin/stress-test/README.md index 87ba85487..d2908f154 100644 --- a/bin/stress-test/README.md +++ b/bin/stress-test/README.md @@ -17,6 +17,29 @@ point-in-time comparison rather than current guarantees. Use the binary help output for the current command and configuration surface. The help output is the source of truth for flags and environment variables. +### Large public account storage map + +`seed-store` can create an exact number of accounts with a deterministic storage map on every public account. The +following command creates one public account containing 250,000 entries in one map, then applies one partial account +update that changes a single entry in that map: + +```sh +cargo run --release --locked -p miden-node-stress-test -- \ + seed-store \ + --data-directory /tmp/miden-large-storage-map \ + --num-accounts 1 \ + --public-accounts-percentage 100 \ + --storage-map-entries 250000 \ + --vault-entries 1 \ + --account-update-blocks 1 +``` + +Account creation and account updates each require a preceding note-emission block. The block metrics label these phases +separately, so `account-update` rows measure the partial public-account updates rather than their setup work. The +generated account IDs are written to `accounts.txt` in the data directory. If a full public account state would exceed +the protocol's transaction account-update size limit, `seed-store` inserts that account at genesis automatically; its +subsequent partial updates still use normal blocks and appear as `account-update` rows. + ## Benchmark Results The following reference results were obtained using a store with 100k accounts, half of which are public. diff --git a/bin/stress-test/src/main.rs b/bin/stress-test/src/main.rs index 705a67d3a..870055bbc 100644 --- a/bin/stress-test/src/main.rs +++ b/bin/stress-test/src/main.rs @@ -1,3 +1,4 @@ +use std::num::NonZeroUsize; use std::path::PathBuf; use clap::{Parser, Subcommand}; @@ -34,22 +35,28 @@ pub enum Command { /// Number of accounts to create. #[arg(short, long, value_name = "NUM_ACCOUNTS")] - num_accounts: usize, + num_accounts: NonZeroUsize, /// Percentage of accounts that will be created as public accounts. The rest will be private /// accounts. - #[arg(short, long, value_name = "PUBLIC_ACCOUNTS_PERCENTAGE", default_value = "0")] + #[arg( + short, + long, + value_name = "PUBLIC_ACCOUNTS_PERCENTAGE", + default_value = "0", + value_parser = clap::value_parser!(u8).range(0..=100) + )] public_accounts_percentage: u8, /// Number of entries to add to a deterministic storage map on every public account. #[arg(long, value_name = "STORAGE_MAP_ENTRIES", default_value = "0")] - storage_map_entries: usize, + storage_map_entries: u32, /// Number of distinct vault assets to add to every public account. #[arg(long, value_name = "VAULT_ENTRIES", default_value = "1")] - vault_entries: usize, + vault_entries: NonZeroUsize, - /// Number of post-initialization blocks to generate with random account updates. + /// Number of post-initialization blocks containing public account updates. #[arg(long, value_name = "ACCOUNT_UPDATE_BLOCKS", default_value = "0")] account_update_blocks: usize, }, @@ -124,10 +131,10 @@ async fn main() { } => { Box::pin(seed_store( data_directory, - num_accounts, + num_accounts.get(), public_accounts_percentage, storage_map_entries, - vault_entries, + vault_entries.get(), account_update_blocks, )) .await; diff --git a/bin/stress-test/src/seeding/metrics.rs b/bin/stress-test/src/seeding/metrics.rs index 56e89e4a9..58c80b302 100644 --- a/bin/stress-test/src/seeding/metrics.rs +++ b/bin/stress-test/src/seeding/metrics.rs @@ -25,14 +25,41 @@ const SQLITE_INDEXES: &[&str] = &[ "idx_transactions_block_num", ]; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum BlockKind { + AccountNoteEmission, + AccountCreation, + UpdateNoteEmission, + AccountUpdate, +} + +impl Display for BlockKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let name = match self { + Self::AccountNoteEmission => "account-note-emission", + Self::AccountCreation => "account-creation", + Self::UpdateNoteEmission => "update-note-emission", + Self::AccountUpdate => "account-update", + }; + f.write_str(name) + } +} + +struct BlockMetric { + kind: BlockKind, + insertion_time: Duration, + get_block_inputs_time: Duration, + block_size: usize, + transaction_count: usize, + public_account_updates: usize, +} + /// Metrics struct to show the results of the stress test pub struct SeedingMetrics { - insertion_time_per_block: Vec, - get_block_inputs_time_per_block: Vec, + blocks: Vec, + pending_get_block_inputs_time: Option, get_batch_inputs_time_per_block: Vec, - bytes_per_block: Vec, - num_insertions: u32, - store_file_sizes: Vec, + store_file_sizes: Vec<(usize, u64)>, initial_store_size: u64, store_file: PathBuf, } @@ -42,11 +69,9 @@ impl SeedingMetrics { pub fn new(store_file: PathBuf) -> Self { let initial_store_size = get_store_size(&store_file); Self { - insertion_time_per_block: Vec::new(), - get_block_inputs_time_per_block: Vec::new(), + blocks: Vec::new(), + pending_get_block_inputs_time: None, get_batch_inputs_time_per_block: Vec::new(), - bytes_per_block: Vec::new(), - num_insertions: 0, store_file_sizes: Vec::new(), initial_store_size, store_file, @@ -54,20 +79,37 @@ impl SeedingMetrics { } /// Tracks a new block insertion to the metrics, with the insertion time and size of the block. - pub fn track_block_insertion(&mut self, insertion_time: Duration, block_size: usize) { - self.insertion_time_per_block.push(insertion_time); - self.bytes_per_block.push(block_size); - self.num_insertions += 1; + pub fn track_block_insertion( + &mut self, + kind: BlockKind, + insertion_time: Duration, + block_size: usize, + transaction_count: usize, + public_account_updates: usize, + ) { + self.blocks.push(BlockMetric { + kind, + insertion_time, + get_block_inputs_time: self.pending_get_block_inputs_time.take().unwrap_or_default(), + block_size, + transaction_count, + public_account_updates, + }); } /// Tracks the size of the store file. pub fn record_store_size(&mut self) { - self.store_file_sizes.push(get_store_size(&self.store_file)); + if let Some(block_index) = self.blocks.len().checked_sub(1) { + self.store_file_sizes.push((block_index, get_store_size(&self.store_file))); + } } /// Tracks the time it takes to query the block inputs. pub fn add_get_block_inputs(&mut self, query_time: Duration) { - self.get_block_inputs_time_per_block.push(query_time); + assert!( + self.pending_get_block_inputs_time.replace(query_time).is_none(), + "block input query time should be consumed before the next query" + ); } /// Tracks the time it takes to query the batch inputs. @@ -79,60 +121,67 @@ impl SeedingMetrics { #[expect(clippy::cast_precision_loss)] fn print_block_metrics(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { writeln!(f, "\nBlock metrics:")?; - writeln!(f, "Note: Each block contains 256 transactions (16 batches * 16 transactions).")?; - writeln!( - f, - "The first transaction sends assets from the faucet to 255 accounts, and the remaining 255 transactions consume notes created in the previous block." - )?; writeln!( f, - "{:<10} {:<20} {:<30} {:<30} {:<20} {:<20}", + "{:<10} {:<24} {:<8} {:<16} {:<18} {:<24} {:<16} {:<16}", "Block #", - "Insert Time (ms)", - "Get Block Inputs Time (ms)", - "Get Batch Inputs Time (ms)", + "Kind", + "Txs", + "Public updates", + "Insert time (ms)", + "Get block inputs (ms)", "Block Size (KB)", "DB Size (MB)" )?; - writeln!(f, "{}", "-".repeat(135))?; - for (i, store_size) in self.store_file_sizes.iter().enumerate() { - let block_index = i * 50; - let insertion_time = self - .insertion_time_per_block - .get(block_index) - .unwrap_or(&Duration::default()) - .as_millis(); - let block_query_time = self - .get_block_inputs_time_per_block - .get(block_index) - .unwrap_or(&Duration::default()) - .as_millis(); - let batch_query_time = self - .get_batch_inputs_time_per_block - .get(block_index) - .unwrap_or(&Duration::default()) - .as_millis(); - let block_size_kb = - *self.bytes_per_block.get(block_index).unwrap_or(&0) as f64 / 1024.0; - let store_size_mb = *store_size as f64 / (1024.0 * 1024.0); + writeln!(f, "{}", "-".repeat(150))?; + for (block_index, metric) in self.blocks.iter().enumerate().filter(|(index, metric)| { + self.blocks.len() <= 20 + || index % 50 == 0 + || index + 1 == self.blocks.len() + || metric.kind == BlockKind::AccountUpdate + }) { + let block_size_kb = metric.block_size as f64 / 1024.0; + let store_size_mb = + self.store_file_sizes.iter().rev().find_map(|(sample_index, size)| { + (*sample_index == block_index).then_some(*size as f64 / (1024.0 * 1024.0)) + }); + let store_size = + store_size_mb.map_or_else(|| "-".to_string(), |size| format!("{size:.1}")); writeln!( f, - "{block_index:<10} {insertion_time:<20} {block_query_time:<30} {batch_query_time:<30} {block_size_kb:<20.1} {store_size_mb:<20.1}", + "{block_index:<10} {:<24} {:<8} {:<16} {:<18} {:<24} {block_size_kb:<16.1} {store_size:<16}", + metric.kind, + metric.transaction_count, + metric.public_account_updates, + metric.insertion_time.as_millis(), + metric.get_block_inputs_time.as_millis(), )?; } + + if !self.get_batch_inputs_time_per_block.is_empty() { + let average = self + .get_batch_inputs_time_per_block + .iter() + .map(Duration::as_micros) + .sum::() + / self.get_batch_inputs_time_per_block.len() as u128; + writeln!(f, "Average get-batch-inputs time: {average} us")?; + } Ok(()) } /// Prints the database table stats. fn print_table_stats(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { writeln!(f, "\nDatabase stats:")?; - writeln!( - f, - "Note: Database contains {} accounts and {} notes across all blocks.", - self.num_insertions * 255, - self.num_insertions * 255 - )?; + let account_count = self.table_row_count("accounts"); + let note_count = self.table_row_count("notes"); + if let (Some(account_count), Some(note_count)) = (account_count, note_count) { + writeln!( + f, + "Note: Database contains {account_count} accounts and {note_count} notes across all blocks." + )?; + } writeln!(f, "{:<35} {:<15} {:<15}", "Table", "Size (MB)", "KB/Entry")?; writeln!(f, "{}", "-".repeat(70))?; for table in SQLITE_TABLES { @@ -162,6 +211,15 @@ impl SeedingMetrics { Ok(()) } + fn table_row_count(&self, table: &str) -> Option { + let output = Command::new("sqlite3") + .arg(&self.store_file) + .arg(format!("SELECT COUNT(*) FROM {table};")) + .output() + .ok()?; + String::from_utf8(output.stdout).ok()?.trim().parse().ok() + } + /// Prints the index stats. fn print_index_stats(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { writeln!(f, "\nIndex stats:")?; @@ -194,9 +252,9 @@ impl Display for SeedingMetrics { writeln!( f, "Inserted {} blocks with avg insertion time {} ms", - self.num_insertions, - (self.insertion_time_per_block.iter().map(Duration::as_millis).sum::() - / self.insertion_time_per_block.len() as u128) + self.blocks.len(), + self.blocks.iter().map(|metric| metric.insertion_time.as_millis()).sum::() + / self.blocks.len().max(1) as u128 )?; writeln!( f, @@ -205,9 +263,12 @@ impl Display for SeedingMetrics { )?; // Print out the average growth rate of the store file - let final_size = self.store_file_sizes.last().unwrap(); - let growth_rate_mb = (final_size - self.initial_store_size) as f64 - / f64::from(self.num_insertions) + let final_size = self + .store_file_sizes + .last() + .map_or_else(|| get_store_size(&self.store_file), |(_, size)| *size); + let growth_rate_mb = final_size.saturating_sub(self.initial_store_size) as f64 + / self.blocks.len().max(1) as f64 / (1024.0 * 1024.0); writeln!(f, "Average DB growth rate: {growth_rate_mb:.1} MB per block")?; diff --git a/bin/stress-test/src/seeding/mod.rs b/bin/stress-test/src/seeding/mod.rs index e20bf0a0d..0b9c817cb 100644 --- a/bin/stress-test/src/seeding/mod.rs +++ b/bin/stress-test/src/seeding/mod.rs @@ -55,7 +55,7 @@ use miden_protocol::transaction::{ }; use miden_protocol::utils::serde::Serializable; use miden_protocol::vm::ExecutionProof; -use miden_protocol::{Felt, ONE, Word}; +use miden_protocol::{ACCOUNT_UPDATE_MAX_SIZE, Felt, ONE, Word}; use miden_standards::account::auth::{Approver, AuthSingleSig}; use miden_standards::account::faucets::{FungibleFaucet, TokenName}; use miden_standards::account::policies::{BurnPolicy, MintPolicy, TokenPolicyManager}; @@ -78,11 +78,58 @@ mod tests; const BATCHES_PER_BLOCK: usize = 16; const TRANSACTIONS_PER_BATCH: usize = 16; +const ACCOUNT_UPDATES_PER_BLOCK: usize = BATCHES_PER_BLOCK * TRANSACTIONS_PER_BATCH - 1; +const STORAGE_MAP_ENTRY_SERIALIZED_SIZE: u64 = 64; +const ACCOUNT_UPDATE_SIZE_HEADROOM: u64 = 32 * 1024; pub const ACCOUNTS_FILENAME: &str = "accounts.txt"; pub const BENCHMARK_STORAGE_MAP_SLOT_NAME: &str = "miden::mock::stress_test::map"; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct AccountBatch { + public: usize, + private: usize, +} + +fn rounded_percentage(count: usize, percentage: u8) -> usize { + let numerator = (count as u128) * u128::from(percentage) + 50; + usize::try_from(numerator / 100).expect("percentage of usize fits into usize") +} + +fn plan_account_batches( + num_accounts: usize, + public_accounts_percentage: u8, + batch_capacity: usize, +) -> Vec { + assert!( + public_accounts_percentage <= 100, + "public account percentage must be at most 100" + ); + assert!(batch_capacity > 0, "account batch capacity must be non-zero"); + + (0..num_accounts) + .step_by(batch_capacity) + .map(|start| { + let end = start.saturating_add(batch_capacity).min(num_accounts); + let public = rounded_percentage(end, public_accounts_percentage) + - rounded_percentage(start, public_accounts_percentage); + AccountBatch { public, private: end - start - public } + }) + .collect() +} + +fn account_update_may_exceed_protocol_limit( + storage_map_entries: u32, + vault_entries: usize, +) -> bool { + let estimated_size = u64::from(storage_map_entries) * STORAGE_MAP_ENTRY_SERIALIZED_SIZE + + u64::try_from(vault_entries).expect("vault entry count fits into u64") + * STORAGE_MAP_ENTRY_SERIALIZED_SIZE + + ACCOUNT_UPDATE_SIZE_HEADROOM; + estimated_size > u64::from(ACCOUNT_UPDATE_MAX_SIZE) +} + // SEED STORE // ================================================================================================ @@ -91,16 +138,25 @@ pub async fn seed_store( data_directory: PathBuf, num_accounts: usize, public_accounts_percentage: u8, - storage_map_entries: usize, + storage_map_entries: u32, vault_entries: usize, account_update_blocks: usize, ) { let start = Instant::now(); + assert!(num_accounts > 0, "--num-accounts must be greater than zero"); + assert!(vault_entries > 0, "--vault-entries must be greater than zero"); assert!( vault_entries <= NoteAssets::MAX_NUM_ASSETS, "--vault-entries must be at most {}", NoteAssets::MAX_NUM_ASSETS ); + if account_update_blocks > 0 { + assert!( + storage_map_entries > 0 + && rounded_percentage(num_accounts, public_accounts_percentage) > 0, + "--account-update-blocks requires at least one public account and a non-empty benchmark storage map" + ); + } // Recreate the data directory (it should be empty for store bootstrapping). // @@ -108,29 +164,72 @@ pub async fn seed_store( let _ = fs_err::remove_dir_all(&data_directory); fs_err::create_dir_all(&data_directory).expect("created data directory"); - // generate the faucet account and the genesis state + // Generate the faucet accounts and genesis state. Public account state which cannot fit into a + // transaction account update is inserted at genesis so later partial updates can still exercise + // normal block application. let benchmark_faucets = create_benchmark_faucets(vault_entries); let faucet = benchmark_faucets[0].clone(); let asset_faucet_ids = benchmark_faucets.iter().map(Account::id).collect::>(); + let num_public_accounts = rounded_percentage(num_accounts, public_accounts_percentage); + let seed_public_accounts_at_genesis = + account_update_may_exceed_protocol_limit(storage_map_entries, vault_entries); + let genesis_account_key_pair = if seed_public_accounts_at_genesis { + let coin_seed: [u64; 4] = rand::rng().random(); + let mut rng = RandomCoin::new(coin_seed.map(Felt::new_unchecked).into()); + Some(SecretKey::with_rng(&mut rng)) + } else { + None + }; + let genesis_benchmark_accounts = + genesis_account_key_pair.as_ref().map_or_else(Vec::new, |key| { + create_existing_benchmark_accounts( + num_public_accounts, + key, + &asset_faucet_ids, + storage_map_entries, + vault_entries, + ) + }); + let mut genesis_accounts = benchmark_faucets; + genesis_accounts.extend(genesis_benchmark_accounts); let fee_params = FeeParameters::new(faucet.id(), 0); let signer = EcdsaSecretKey::new(); - let genesis_state = GenesisState::new(benchmark_faucets, fee_params, 1, 1, signer.public_key()); - let genesis_block = genesis_state - .clone() - .into_block(&signer) - .expect("genesis block should be created"); + let genesis_state = GenesisState::new(genesis_accounts, fee_params, 1, 1, signer.public_key()); + let genesis_block = genesis_state.into_block(&signer).expect("genesis block should be created"); + let genesis_header = genesis_block.inner().header().clone(); State::bootstrap(genesis_block, &data_directory).expect("store should bootstrap"); let store_state = load_state(data_directory.clone()).await; + // Recreate the deterministic genesis benchmark accounts after bootstrapping instead of keeping + // another copy of their potentially very large maps alive while the genesis block is built. + let initial_accounts = genesis_account_key_pair.as_ref().map_or_else(Vec::new, |key| { + create_existing_benchmark_accounts( + num_public_accounts, + key, + &asset_faucet_ids, + storage_map_entries, + vault_entries, + ) + }); + let account_batches = if seed_public_accounts_at_genesis { + plan_account_batches(num_accounts - num_public_accounts, 0, ACCOUNT_UPDATES_PER_BLOCK) + } else { + plan_account_batches(num_accounts, public_accounts_percentage, ACCOUNT_UPDATES_PER_BLOCK) + }; + // start generating blocks let accounts_filepath = data_directory.join(ACCOUNTS_FILENAME); let data_directory = miden_node_store::DataDirectory::load(data_directory).expect("data directory should exist"); - let genesis_header = genesis_state.into_block(&signer).unwrap().into_inner(); let metrics = Box::pin(generate_blocks( - num_accounts, - public_accounts_percentage, + account_batches, + initial_accounts, + if seed_public_accounts_at_genesis { + u64::try_from(num_public_accounts).expect("public account count fits into u64") + } else { + 0 + }, faucet, genesis_header, &store_state, @@ -150,47 +249,37 @@ pub async fn seed_store( /// Generates batches of transactions to be inserted into the store. /// -/// The first transaction in each batch sends assets from the faucet to 255 accounts. -/// The rest of the transactions consume the notes created by the faucet in the previous block. +/// Account-creation notes are emitted in batches of at most 255 and consumed in a subsequent +/// block. Optional benchmark updates likewise use separate note-emission and account-update +/// blocks so their insertion times can be reported independently. #[expect(clippy::too_many_arguments)] #[expect(clippy::too_many_lines)] async fn generate_blocks( - num_accounts: usize, - public_accounts_percentage: u8, + account_batches: Vec, + initial_accounts: Vec, + first_account_index: u64, mut faucet: Account, - genesis_block: SignedBlock, + genesis_header: BlockHeader, store_state: &Arc, data_directory: DataDirectory, accounts_filepath: PathBuf, signer: &EcdsaSecretKey, - storage_map_entries: usize, + storage_map_entries: u32, vault_entries: usize, account_update_blocks: usize, asset_faucet_ids: Vec, ) -> SeedingMetrics { - // Each block is composed of [`BATCHES_PER_BLOCK`] batches, and each batch is composed of - // [`TRANSACTIONS_PER_BATCH`] txs. The first note of the block is always a send assets tx from - // the faucet to (BATCHES_PER_BLOCK * TRANSACTIONS_PER_BATCH) - 1 accounts. The rest of the - // notes are consume note txs from the (BATCHES_PER_BLOCK * TRANSACTIONS_PER_BATCH) - 1 accounts - // that were minted in the previous block. let mut metrics = SeedingMetrics::new(data_directory.database_path()); - let mut account_ids = vec![]; - let mut note_nullifiers = vec![]; - let mut account_states: BTreeMap = BTreeMap::new(); + let mut account_ids = initial_accounts.iter().map(Account::id).collect::>(); + let mut account_states = initial_accounts + .into_iter() + .map(|account| (account.id(), account)) + .collect::>(); let mut consume_notes_txs: Vec = vec![]; let mut pending_consumed_accounts: Vec = vec![]; - let consumes_per_block = TRANSACTIONS_PER_BATCH * BATCHES_PER_BLOCK - 1; - #[expect(clippy::cast_sign_loss, clippy::cast_precision_loss)] - let num_public_accounts = (consumes_per_block as f64 - * (f64::from(public_accounts_percentage) / 100.0)) - .round() as usize; - let num_private_accounts = consumes_per_block - num_public_accounts; - // +1 to account for the first block with the send assets tx only - let total_blocks = (num_accounts / consumes_per_block) + 1; - // share random coin seed and key pair for all accounts to avoid key generation overhead let coin_seed: [u64; 4] = rand::rng().random(); let rng = Arc::new(Mutex::new(RandomCoin::new(coin_seed.map(Felt::new_unchecked).into()))); @@ -199,45 +288,49 @@ async fn generate_blocks( SecretKey::with_rng(&mut *rng) }; - let mut prev_block_header = genesis_block.header().clone(); - let mut current_anchor_header = genesis_block.header().clone(); + let mut prev_block_header = genesis_header; + let mut next_account_index = first_account_index; + let mut pending_public_accounts = 0; - for i in 0..total_blocks { + for (batch_index, batch) in account_batches.into_iter().enumerate() { let mut block_txs = Vec::with_capacity(BATCHES_PER_BLOCK * TRANSACTIONS_PER_BATCH); // create public accounts and notes that mint assets for these accounts let (pub_accounts, pub_notes) = create_accounts_and_notes( - num_public_accounts, + batch.public, AccountType::Public, &key_pair, &rng, &asset_faucet_ids, - i, + next_account_index, storage_map_entries, vault_entries, ); + next_account_index += u64::try_from(batch.public).expect("account count fits into u64"); // create private accounts and notes that mint assets for these accounts let (priv_accounts, priv_notes) = create_accounts_and_notes( - num_private_accounts, + batch.private, AccountType::Private, &key_pair, &rng, &asset_faucet_ids, - i, + next_account_index, storage_map_entries, vault_entries, ); + next_account_index += u64::try_from(batch.private).expect("account count fits into u64"); - let notes = [pub_notes, priv_notes].concat(); - let accounts = [pub_accounts, priv_accounts].concat(); + let mut notes = pub_notes; + notes.extend(priv_notes); + let mut accounts = pub_accounts; + accounts.extend(priv_accounts); account_ids.extend(accounts.iter().map(Account::id)); - note_nullifiers.extend(notes.iter().map(|n| n.nullifier().prefix())); - // create the tx that creates the notes let emit_note_tx = create_emit_note_tx(&prev_block_header, &mut faucet, notes.clone()); // collect all the txs + let has_pending_account_creations = !consume_notes_txs.is_empty(); block_txs.push(emit_note_tx); block_txs.extend(consume_notes_txs); @@ -251,13 +344,23 @@ async fn generate_blocks( let block_inputs = get_block_inputs(store_state, &batches, &mut metrics).await; // update blocks - prev_block_header = - apply_block(batches, block_inputs, store_state, &mut metrics, signer).await; + let block_kind = if has_pending_account_creations { + metrics::BlockKind::AccountCreation + } else { + metrics::BlockKind::AccountNoteEmission + }; + prev_block_header = apply_block( + batches, + block_inputs, + store_state, + &mut metrics, + signer, + block_kind, + pending_public_accounts, + ) + .await; account_states .extend(pending_consumed_accounts.into_iter().map(|account| (account.id(), account))); - if current_anchor_header.block_epoch() != prev_block_header.block_epoch() { - current_anchor_header = prev_block_header.clone(); - } // create the consume notes txs to be used in the next block let batch_inputs = @@ -269,25 +372,51 @@ async fn generate_blocks( &batch_inputs.note_proofs, None, ); + pending_public_accounts = batch.public; // track store size every 50 blocks - if i % 50 == 0 { + if batch_index % 50 == 0 { metrics.record_store_size(); } } + // The creation pipeline emits notes in one block and consumes them in the next. Flush the final + // set of consume transactions so `num_accounts` describes persisted accounts rather than merely + // emitted account-creation notes. + if !consume_notes_txs.is_empty() { + let batches: Vec = consume_notes_txs + .par_chunks(TRANSACTIONS_PER_BATCH) + .map(|txs| create_batch(txs, &prev_block_header)) + .collect(); + let block_inputs = get_block_inputs(store_state, &batches, &mut metrics).await; + prev_block_header = apply_block( + batches, + block_inputs, + store_state, + &mut metrics, + signer, + metrics::BlockKind::AccountCreation, + pending_public_accounts, + ) + .await; + account_states + .extend(pending_consumed_accounts.into_iter().map(|account| (account.id(), account))); + metrics.record_store_size(); + } + let update_note_faucet_ids = asset_faucet_ids.iter().take(vault_entries).copied().collect::>(); let mut random = rand::rng(); for update_block_index in 0..account_update_blocks { - let mut block_txs = Vec::with_capacity(BATCHES_PER_BLOCK * TRANSACTIONS_PER_BATCH); - let selected_account_ids = select_random_account_ids_for_update_notes( &account_states, - &pending_consumed_accounts, - consumes_per_block, + ACCOUNT_UPDATES_PER_BLOCK, &mut random, ); + assert!( + !selected_account_ids.is_empty(), + "--account-update-blocks requires at least one public account with a benchmark storage map" + ); let notes = { let mut note_rng = rng.lock().unwrap(); selected_account_ids @@ -297,31 +426,31 @@ async fn generate_blocks( }; let emit_note_tx = create_emit_note_tx(&prev_block_header, &mut faucet, notes.clone()); - block_txs.push(emit_note_tx); - block_txs.extend(consume_notes_txs); - - let batches: Vec = block_txs - .par_chunks(TRANSACTIONS_PER_BATCH) - .map(|txs| create_batch(txs, &prev_block_header)) - .collect(); + let batches = vec![create_batch(std::slice::from_ref(&emit_note_tx), &prev_block_header)]; let block_inputs = get_block_inputs(store_state, &batches, &mut metrics).await; - - prev_block_header = - apply_block(batches, block_inputs, store_state, &mut metrics, signer).await; - account_states - .extend(pending_consumed_accounts.into_iter().map(|account| (account.id(), account))); - if current_anchor_header.block_epoch() != prev_block_header.block_epoch() { - current_anchor_header = prev_block_header.clone(); - } + prev_block_header = apply_block( + batches, + block_inputs, + store_state, + &mut metrics, + signer, + metrics::BlockKind::UpdateNoteEmission, + 0, + ) + .await; let batch_inputs = get_batch_inputs(store_state, &prev_block_header, ¬es, &mut metrics).await; let accounts = selected_account_ids .iter() - .filter_map(|account_id| account_states.get(account_id).cloned()) + .map(|account_id| { + account_states + .remove(account_id) + .expect("selected update account should be present") + }) .collect::>(); - (pending_consumed_accounts, consume_notes_txs) = create_consume_note_txs( + let (updated_accounts, update_txs) = create_consume_note_txs( &prev_block_header, accounts, notes, @@ -331,12 +460,27 @@ async fn generate_blocks( storage_map_entries, }), ); - - if update_block_index % 50 == 0 { - metrics.record_store_size(); - } + let batches: Vec = update_txs + .par_chunks(TRANSACTIONS_PER_BATCH) + .map(|txs| create_batch(txs, &prev_block_header)) + .collect(); + let block_inputs = get_block_inputs(store_state, &batches, &mut metrics).await; + prev_block_header = apply_block( + batches, + block_inputs, + store_state, + &mut metrics, + signer, + metrics::BlockKind::AccountUpdate, + selected_account_ids.len(), + ) + .await; + account_states.extend(updated_accounts.into_iter().map(|account| (account.id(), account))); + metrics.record_store_size(); } + metrics.record_store_size(); + // dump account ids to a file let mut file = fs::File::create(accounts_filepath).await.unwrap(); for id in account_ids { @@ -356,7 +500,10 @@ async fn apply_block( store_state: &Arc, metrics: &mut SeedingMetrics, signer: &EcdsaSecretKey, + block_kind: metrics::BlockKind, + public_account_updates: usize, ) -> BlockHeader { + let transaction_count = batches.iter().map(|batch| batch.transactions().as_slice().len()).sum(); let proposed_block = ProposedBlock::new(block_inputs.clone(), batches).unwrap(); let (header, body) = proposed_block.clone().into_header_and_body().unwrap(); let block_size: usize = header.to_bytes().len() + body.to_bytes().len(); @@ -372,7 +519,13 @@ async fn apply_block( .apply_block_with_proving_inputs(ordered_batches, block_inputs, signed_block) .await .unwrap(); - metrics.track_block_insertion(start.elapsed(), block_size); + metrics.track_block_insertion( + block_kind, + start.elapsed(), + block_size, + transaction_count, + public_account_updates, + ); header } @@ -392,8 +545,8 @@ fn create_accounts_and_notes( key_pair: &SecretKey, rng: &Arc>, asset_faucet_ids: &[AccountId], - block_num: usize, - storage_map_entries: usize, + first_account_index: u64, + storage_map_entries: u32, vault_entries: usize, ) -> (Vec, Vec) { assert!( @@ -410,7 +563,8 @@ fn create_accounts_and_notes( .map(|account_index| { let account = create_account( key_pair.public_key(), - ((block_num * num_accounts) + account_index) as u64, + first_account_index + + u64::try_from(account_index).expect("account index fits into u64"), account_type, storage_map_entries, ); @@ -444,17 +598,16 @@ fn create_note(faucet_ids: &[AccountId], target_id: AccountId, rng: &mut RandomC fn select_random_account_ids_for_update_notes( account_states: &BTreeMap, - pending_accounts: &[Account], max_accounts: usize, rng: &mut R, ) -> Vec { - let mut account_ids = account_states.keys().copied().collect::>(); - for account in pending_accounts { - let account_id = account.id(); - if !account_states.contains_key(&account_id) { - account_ids.push(account_id); - } - } + let mut account_ids = account_states + .values() + .filter(|account| { + account.is_public() && account.storage().get(&benchmark_storage_map_slot()).is_some() + }) + .map(Account::id) + .collect::>(); account_ids.shuffle(rng); account_ids.truncate(max_accounts); @@ -464,14 +617,14 @@ fn select_random_account_ids_for_update_notes( #[derive(Clone, Copy)] struct BenchmarkStorageUpdate { block_index: usize, - storage_map_entries: usize, + storage_map_entries: u32, } -fn benchmark_storage_map_update_value(block_index: usize, tx_index: usize, key_index: u32) -> Word { +fn benchmark_storage_map_update_value(block_index: usize, tx_index: u32, key_index: u32) -> Word { Word::from([ - Felt::ZERO, - Felt::from(u32::try_from(block_index).expect("update block index fits into u32")), - Felt::from(u32::try_from(tx_index).expect("transaction index fits into u32")), + Felt::ONE, + Felt::from(u32::try_from(block_index + 1).expect("update block index fits into u32")), + Felt::from(tx_index.checked_add(1).expect("transaction index fits into u32")), Felt::from(key_index), ]) } @@ -480,14 +633,14 @@ fn update_benchmark_storage_map_entry( account: &mut Account, block_index: usize, tx_index: usize, - storage_map_entries: usize, + storage_map_entries: u32, ) -> bool { if !account.is_public() || storage_map_entries == 0 { return false; } - let key_index = - u32::try_from((tx_index % storage_map_entries) + 1).expect("storage map key fits into u32"); + let tx_index = u32::try_from(tx_index).expect("transaction index fits into u32"); + let key_index = (tx_index % storage_map_entries) + 1; let key = StorageMapKey::from_index(key_index); let value = benchmark_storage_map_update_value(block_index, tx_index, key_index); @@ -507,7 +660,7 @@ fn create_account( public_key: PublicKey, index: u64, account_type: AccountType, - storage_map_entries: usize, + storage_map_entries: u32, ) -> Account { let init_seed: Vec<_> = index.to_be_bytes().into_iter().chain([0u8; 24]).collect(); let mut builder = AccountBuilder::new(init_seed.try_into().unwrap()) @@ -519,15 +672,15 @@ fn create_account( .with_component(BasicWallet); if account_type == AccountType::Public && storage_map_entries > 0 { - let entries = (1..=storage_map_entries) - .map(|i| { - let i = u32::try_from(i).expect("storage map entry index fits into u32"); - ( - StorageMapKey::from_index(i), - Word::from([Felt::ZERO, Felt::ZERO, Felt::ZERO, Felt::from(i)]), - ) - }) - .collect::>(); + let entry_count = + usize::try_from(storage_map_entries).expect("storage map entry count fits into usize"); + let entries = (0..entry_count).map(|index| { + let i = u32::try_from(index + 1).expect("storage map entry index fits into u32"); + ( + StorageMapKey::from_index(i), + Word::from([Felt::ZERO, Felt::ZERO, Felt::ZERO, Felt::from(i)]), + ) + }); let storage_map = StorageMap::with_entries(entries).unwrap(); let component_storage = vec![StorageSlot::with_map(benchmark_storage_map_slot(), storage_map)]; @@ -553,6 +706,34 @@ end", builder.build().unwrap() } +fn create_existing_benchmark_accounts( + num_accounts: usize, + key_pair: &SecretKey, + asset_faucet_ids: &[AccountId], + storage_map_entries: u32, + vault_entries: usize, +) -> Vec { + (0..num_accounts) + .into_par_iter() + .map(|account_index| { + let mut account = create_account( + key_pair.public_key(), + u64::try_from(account_index).expect("account index fits into u64"), + AccountType::Public, + storage_map_entries, + ); + for faucet_id in asset_faucet_ids.iter().take(vault_entries) { + account + .vault_mut() + .add_asset(FungibleAsset::new(*faucet_id, 10).unwrap().into()) + .unwrap(); + } + account.increment_nonce(ONE).unwrap(); + account + }) + .collect() +} + fn create_benchmark_faucets(vault_entries: usize) -> Vec { (0..vault_entries.max(1)) .map(|index| create_faucet_with_seed(index as u64)) @@ -717,8 +898,8 @@ fn create_existing_account_patch( if storage_update.storage_map_entries > 0 && account.storage().get(&benchmark_storage_map_slot()).is_some() => { - let key_index = u32::try_from((tx_index % storage_update.storage_map_entries) + 1) - .expect("storage map key fits into u32"); + let tx_index = u32::try_from(tx_index).expect("transaction index fits into u32"); + let key_index = (tx_index % storage_update.storage_map_entries) + 1; let mut entries = StorageMapPatchEntries::new(); entries.insert( StorageMapKey::from_index(key_index), diff --git a/bin/stress-test/src/seeding/tests.rs b/bin/stress-test/src/seeding/tests.rs index fff424da0..e0be9a7b2 100644 --- a/bin/stress-test/src/seeding/tests.rs +++ b/bin/stress-test/src/seeding/tests.rs @@ -9,6 +9,45 @@ fn benchmark_fungible_faucet_ids(vault_entries: usize) -> Vec { .collect() } +#[test] +fn account_batches_honor_the_exact_requested_count() { + assert_eq!( + plan_account_batches(1, 100, ACCOUNT_UPDATES_PER_BLOCK), + vec![AccountBatch { public: 1, private: 0 }] + ); + assert_eq!( + plan_account_batches(254, 50, ACCOUNT_UPDATES_PER_BLOCK), + vec![AccountBatch { public: 127, private: 127 }] + ); + assert_eq!( + plan_account_batches(255, 50, ACCOUNT_UPDATES_PER_BLOCK), + vec![AccountBatch { public: 128, private: 127 }] + ); + assert_eq!( + plan_account_batches(256, 50, ACCOUNT_UPDATES_PER_BLOCK), + vec![ + AccountBatch { public: 128, private: 127 }, + AccountBatch { public: 0, private: 1 }, + ] + ); +} + +#[test] +fn account_batches_distribute_public_accounts_across_partial_batches() { + let batches = plan_account_batches(1_000, 37, ACCOUNT_UPDATES_PER_BLOCK); + + assert!(batches.iter().all(|batch| batch.public + batch.private <= 255)); + assert_eq!(batches.iter().map(|batch| batch.public).sum::(), 370); + assert_eq!(batches.iter().map(|batch| batch.private).sum::(), 630); +} + +#[test] +fn oversized_account_updates_are_seeded_at_genesis() { + assert!(!account_update_may_exceed_protocol_limit(128, 1)); + assert!(account_update_may_exceed_protocol_limit(4_096, 1)); + assert!(account_update_may_exceed_protocol_limit(250_000, 1)); +} + #[test] fn public_account_can_be_created_with_large_storage_map() { let coin_seed = [1, 2, 3, 4].map(Felt::new_unchecked); @@ -112,3 +151,76 @@ fn private_account_storage_map_update_is_skipped() { assert!(!updated); } + +#[tokio::test(flavor = "multi_thread")] +async fn seed_store_persists_one_public_account_and_applies_one_map_update() { + use miden_node_proto::domain::account::{ + AccountDetailRequest, + AccountRequest, + AccountStorageRequest, + StorageMapEntries, + }; + + let temp_dir = tempfile::tempdir().unwrap(); + let data_directory = temp_dir.path().join("store"); + seed_store(data_directory.clone(), 1, 100, 4, 1, 1).await; + + let account_ids = fs_err::read_to_string(data_directory.join(ACCOUNTS_FILENAME)).unwrap(); + let account_ids = account_ids.lines().collect::>(); + assert_eq!(account_ids.len(), 1); + let account_id = AccountId::from_hex(account_ids[0]).unwrap(); + + let state = load_state(data_directory).await; + let response = state + .get_account(AccountRequest { + account_id, + block_num: None, + details: Some(AccountDetailRequest { + code_commitment: None, + asset_vault_commitment: None, + storage_request: AccountStorageRequest::AllStorageMaps, + }), + }) + .await + .unwrap(); + let details = response.details.expect("public account details should be returned"); + assert_eq!(details.storage_details.map_details.len(), 1); + let map_details = &details.storage_details.map_details[0]; + assert_eq!(map_details.slot_name, benchmark_storage_map_slot()); + let StorageMapEntries::AllEntries(entries) = &map_details.entries else { + panic!("small benchmark map should return all entries"); + }; + assert_eq!(entries.len(), 4); + assert_eq!( + entries + .iter() + .find(|(key, _)| *key == StorageMapKey::from_index(1)) + .map(|(_, value)| *value), + Some(benchmark_storage_map_update_value(0, 0, 1)) + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn seed_store_handles_map_larger_than_transaction_account_update_limit() { + use miden_node_proto::domain::account::AccountRequest; + + let temp_dir = tempfile::tempdir().unwrap(); + let data_directory = temp_dir.path().join("store"); + seed_store(data_directory.clone(), 1, 100, 4_096, 1, 1).await; + + let account_ids = fs_err::read_to_string(data_directory.join(ACCOUNTS_FILENAME)).unwrap(); + let account_ids = account_ids.lines().collect::>(); + assert_eq!(account_ids.len(), 1); + let account_id = AccountId::from_hex(account_ids[0]).unwrap(); + + let state = load_state(data_directory).await; + let response = state + .get_account(AccountRequest { + account_id, + block_num: None, + details: None, + }) + .await + .unwrap(); + assert_ne!(response.witness.state_commitment(), Word::empty()); +} diff --git a/crates/store/Cargo.toml b/crates/store/Cargo.toml index 8cc1a9c05..bd20adc7c 100644 --- a/crates/store/Cargo.toml +++ b/crates/store/Cargo.toml @@ -69,6 +69,10 @@ harness = false name = "account_tree" required-features = ["rocksdb"] +[[bench]] +harness = false +name = "account_state_forest" + [package.metadata.cargo-machete] # This is an indirect dependency for which we need to enable optimisations/features # via feature flags. Because we don't use it directly in code, machete diff --git a/crates/store/benches/account_state_forest.rs b/crates/store/benches/account_state_forest.rs new file mode 100644 index 000000000..3955dfce5 --- /dev/null +++ b/crates/store/benches/account_state_forest.rs @@ -0,0 +1,73 @@ +use std::hint::black_box; + +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use miden_crypto::hash::rpo::Rpo256; +use miden_crypto::merkle::smt::{ + ForestInMemoryBackend, + LargeSmtForest, + LineageId, + SmtForestMutationSet, + SmtForestOperation, + SmtForestUpdateBatch, + SmtUpdateBatch, + TreeId, +}; +use miden_protocol::Word; + +const INITIAL_VERSION: u64 = 1; +const RECREATED_VERSION: u64 = 2; + +fn lineage() -> LineageId { + LineageId::new([1; 32]) +} + +fn word(value: u32) -> Word { + Word::from([value, 0, 0, 0]) +} + +fn key(value: u32) -> Word { + Rpo256::hash(&value.to_le_bytes()) +} + +fn forest_with_entries(entry_count: u32) -> LargeSmtForest { + let mut forest = LargeSmtForest::new(ForestInMemoryBackend::new()).unwrap(); + let updates = + SmtUpdateBatch::new((0..entry_count).map(|index| { + SmtForestOperation::insert(key(index + 1), word(entry_count + index + 1)) + })); + forest.add_lineage(lineage(), INITIAL_VERSION, updates).unwrap(); + forest +} + +/// Mirrors storage-map `Create` preparation when the lineage already contains entries. +fn prepare_recreation( + forest: &LargeSmtForest, +) -> SmtForestMutationSet { + let lineage = lineage(); + let tree = TreeId::new(lineage, INITIAL_VERSION); + let removals = forest + .entries(tree) + .unwrap() + .map(|entry| SmtForestOperation::remove(entry.unwrap().key)); + let mut batch = SmtForestUpdateBatch::empty(); + batch.add_operations(lineage, removals); + batch.operations(lineage).add_insert(key(u32::MAX - 1), word(u32::MAX)); + + forest.compute_forest_mutations(RECREATED_VERSION, batch).unwrap() +} + +fn bench_storage_map_recreation(c: &mut Criterion) { + let mut group = c.benchmark_group("account_state_forest_storage_map_recreation"); + + for entry_count in [16, 256, 4_096, 16_384] { + let forest = forest_with_entries(entry_count); + group.bench_with_input(BenchmarkId::from_parameter(entry_count), &entry_count, |b, _| { + b.iter(|| black_box(prepare_recreation(black_box(&forest)))); + }); + } + + group.finish(); +} + +criterion_group!(benches, bench_storage_map_recreation); +criterion_main!(benches); diff --git a/crates/store/src/account_state_forest/mod.rs b/crates/store/src/account_state_forest/mod.rs index 89a911cd2..e5bcba873 100644 --- a/crates/store/src/account_state_forest/mod.rs +++ b/crates/store/src/account_state_forest/mod.rs @@ -1,3 +1,4 @@ +use std::collections::BTreeMap; use std::num::NonZeroUsize; use miden_crypto::hash::rpo::Rpo256; @@ -15,10 +16,9 @@ use miden_node_utils::tracing::miden_instrument; use miden_protocol::account::{ AccountId, AccountPatch, - AccountStoragePatch, - AccountVaultPatch, StorageMapKey, StorageMapKeyHash, + StoragePatchOperation, StorageSlotName, }; use miden_protocol::asset::{Asset, AssetId, AssetIdHash}; @@ -29,7 +29,9 @@ use miden_protocol::crypto::merkle::smt::{ LineageId, RootInfo, SMT_DEPTH, + SmtForestMutationSet, SmtForestOperation, + SmtForestUpdateBatch, SmtUpdateBatch, TreeId, }; @@ -41,6 +43,8 @@ use thiserror::Error; use crate::COMPONENT; pub use crate::db::models::queries::HISTORICAL_BLOCK_RETENTION; +use crate::db::models::queries::{PrecomputedPublicAccountState, PrecomputedPublicAccountStates}; +use crate::errors::AccountStateForestUpdateError; #[cfg(test)] mod tests; @@ -97,6 +101,31 @@ pub(crate) struct AccountStateForest { pub(crate) vault_key_cache: LruCache, } +/// Account-state forest mutations and roots prepared against one forest snapshot. +/// +/// The update is bound to the `block_num` passed to +/// [`AccountStateForest::compute_block_update_mutations`] and to the forest lineage versions and +/// roots observed by that call. It must be applied exactly once, to the same forest, without an +/// intervening mutation. The type is consumed on application to prevent accidental reuse. +pub(crate) struct PreparedAccountStateForestBlockUpdate { + /// Roots derived from `mutations` for persistence alongside the corresponding block. + pub(crate) account_states: PrecomputedPublicAccountStates, + /// Snapshot-bound forest mutations that produce `account_states`. + mutations: SmtForestMutationSet, + /// Patches retained to update raw-key caches after the forest mutation succeeds. + account_patches: Vec, +} + +/// Forest lineages updated by each account patch in a prepared block update. +/// +/// This mapping associates roots from the forest mutation set with account vaults and named +/// storage-map slots when building the precomputed state passed to SQLite. +#[derive(Default)] +struct AccountUpdateForestLineages { + vault: BTreeMap, + storage: BTreeMap>, +} + #[cfg(test)] impl AccountStateForest { pub(crate) fn new() -> Self { @@ -195,6 +224,258 @@ impl AccountStateForest { .collect() } + fn update_batch_from_operations(operations: Vec) -> SmtUpdateBatch { + if operations.is_empty() { + SmtUpdateBatch::empty() + } else { + SmtUpdateBatch::new(operations.into_iter()) + } + } + + /// Builds the leaf removals needed to replace a storage-map lineage with an empty tree. + /// + /// `LargeSmtForest` does not currently expose a lineage reset or replacement mutation, so this + /// operation enumerates the latest tree and costs O(n) in its number of entries. A removed map + /// remains the latest version of its lineage until it is recreated; pruning cannot discard that + /// latest version. The `account_state_forest` benchmark tracks the resulting recreation cost so + /// this fallback remains visible until the forest API provides a constant-size reset primitive. + fn remove_current_tree_operations( + &self, + lineage: LineageId, + ) -> Result, LargeSmtForestError> { + let Some(version) = self.forest.latest_version(lineage) else { + return Ok(Vec::new()); + }; + + let tree = TreeId::new(lineage, version); + self.forest + .entries(tree)? + .map(|entry| entry.map(|entry| SmtForestOperation::remove(entry.key))) + .collect() + } + + /// Adds the vault operations from `patch` to a prepared forest update batch. + /// + /// Full-state patches always create a vault lineage, including for an empty vault. Partial + /// patches with no vault changes are skipped. Updated lineages are recorded so their computed + /// roots can later be associated with the account. + /// + /// # Errors + /// + /// Returns an error if a full-state patch targets an existing vault lineage. + fn add_vault_updates( + &self, + batch: &mut SmtForestUpdateBatch, + lineages: &mut AccountUpdateForestLineages, + patch: &AccountPatch, + ) -> Result<(), AccountStateForestUpdateError> { + let account_id = patch.id(); + if !patch.is_full_state() && patch.vault().is_empty() { + return Ok(()); + } + + let lineage = Self::vault_lineage_id(account_id); + if patch.is_full_state() && self.forest.latest_version(lineage).is_some() { + return Err(AccountStateForestUpdateError::VaultLineageAlreadyExists { account_id }); + } + + let operations = Self::build_forest_operations( + patch.vault().iter().map(|(key, value)| (key.hash().as_word(), *value)), + ); + let updates = Self::update_batch_from_operations(operations); + batch.operations(lineage).add_operations(updates.into_iter()); + lineages.vault.insert(account_id, lineage); + Ok(()) + } + + /// Adds storage-map lineage creation operations for a full-state account patch. + /// + /// Empty-word entries are omitted from the new map state. Every map lineage is recorded, + /// including empty maps, so its computed root can be passed to SQLite. + /// + /// # Errors + /// + /// Returns an error if any storage-map lineage for the account already exists. + fn add_full_state_storage_updates( + &self, + batch: &mut SmtForestUpdateBatch, + lineages: &mut AccountUpdateForestLineages, + patch: &AccountPatch, + ) -> Result<(), AccountStateForestUpdateError> { + let account_id = patch.id(); + for (slot_name, map_patch) in patch.storage().maps() { + let raw_map_entries = Vec::from_iter( + map_patch + .entries() + .into_iter() + .flat_map(|entries| entries.as_map().iter()) + .filter_map( + |(&key, &value)| { + if value == EMPTY_WORD { None } else { Some((key, value)) } + }, + ), + ); + let operations = Self::build_forest_operations( + raw_map_entries.iter().map(|(raw_key, value)| (raw_key.hash().into(), *value)), + ); + let lineage = Self::storage_lineage_id(account_id, slot_name); + if self.forest.latest_version(lineage).is_some() { + return Err(AccountStateForestUpdateError::StorageLineageAlreadyExists { + account_id, + slot_name: slot_name.clone(), + }); + } + + let updates = Self::update_batch_from_operations(operations); + batch.operations(lineage).add_operations(updates.into_iter()); + lineages + .storage + .entry(account_id) + .or_default() + .insert(slot_name.clone(), lineage); + } + Ok(()) + } + + /// Adds storage-map operations from a partial account patch to a prepared update batch. + /// + /// Updates are applied incrementally. A `Create` replaces an existing lineage by first removing + /// its current leaves, while a `Remove` only removes the slot from the account storage header and + /// leaves its forest lineage available for historical queries. No-op updates are skipped, but an + /// empty `Create` still creates a lineage and records its root. + /// + /// # Errors + /// + /// Returns an error if reading the current lineage while preparing a replacement fails. + fn add_partial_storage_updates( + &self, + batch: &mut SmtForestUpdateBatch, + lineages: &mut AccountUpdateForestLineages, + patch: &AccountPatch, + ) -> Result<(), LargeSmtForestError> { + if patch.storage().is_empty() { + return Ok(()); + } + + let account_id = patch.id(); + for (slot_name, map_patch) in patch.storage().maps() { + let lineage = Self::storage_lineage_id(account_id, slot_name); + let Some(entries) = map_patch.entries() else { + // Removing the slot from the account storage header makes this lineage + // inaccessible. The forest has no lineage reset primitive, so keep it unchanged for + // historical queries. This retains the latest tree indefinitely; a later Create + // replaces it by enumerating and removing its leaves. + continue; + }; + if entries.is_empty() && map_patch.patch_op() != StoragePatchOperation::Create { + continue; + } + + let mut operations = if map_patch.patch_op() == StoragePatchOperation::Create { + self.remove_current_tree_operations(lineage)? + } else { + Vec::new() + }; + operations.extend(Self::build_forest_operations( + entries.as_map().iter().map(|(key, value)| (key.hash().into(), *value)), + )); + + if operations.is_empty() && self.forest.latest_version(lineage).is_none() { + // An empty Create still needs an empty lineage and its root. + if map_patch.patch_op() != StoragePatchOperation::Create { + continue; + } + } + + let updates = Self::update_batch_from_operations(operations); + batch.operations(lineage).add_operations(updates.into_iter()); + lineages + .storage + .entry(account_id) + .or_default() + .insert(slot_name.clone(), lineage); + } + + Ok(()) + } + + /// Builds the forest update batch and lineage lookup for a collection of account patches. + /// + /// This method does not mutate the forest. The returned lineage lookup identifies the vault and + /// storage-map roots that must be extracted after the batch mutations are computed. + /// + /// # Errors + /// + /// Returns an error if a full-state patch targets an existing lineage or a partial storage-map + /// replacement cannot be prepared. + fn prepare_block_update_batch( + &self, + account_patches: &[AccountPatch], + ) -> Result<(SmtForestUpdateBatch, AccountUpdateForestLineages), AccountStateForestUpdateError> + { + let mut batch = SmtForestUpdateBatch::empty(); + let mut lineages = AccountUpdateForestLineages::default(); + + for patch in account_patches { + self.add_vault_updates(&mut batch, &mut lineages, patch)?; + if patch.is_full_state() { + self.add_full_state_storage_updates(&mut batch, &mut lineages, patch)?; + } else { + self.add_partial_storage_updates(&mut batch, &mut lineages, patch)?; + } + } + + Ok((batch, lineages)) + } + + /// Associates roots from a computed forest mutation set with their public accounts. + /// + /// Each result contains the account's final vault root and roots only for storage maps mutated + /// by the corresponding patch. An unchanged vault uses its current forest root. + /// + /// # Errors + /// + /// Returns an error if the mutation set is missing a root for a lineage recorded while building + /// the update batch. + fn precomputed_account_states_from_mutations( + &self, + account_patches: &[AccountPatch], + lineages: &AccountUpdateForestLineages, + mutations: &SmtForestMutationSet, + ) -> Result { + let mut account_states = PrecomputedPublicAccountStates::new(); + let roots = mutations + .roots() + .map(|root| (root.lineage(), root.root())) + .collect::>(); + + for patch in account_patches { + let account_id = patch.id(); + let vault_root = match lineages.vault.get(&account_id) { + Some(lineage) => roots.get(lineage).copied().ok_or( + AccountStateForestUpdateError::MissingComputedRoot { lineage: *lineage }, + )?, + None => self.get_latest_vault_root(account_id), + }; + let mut storage_map_roots = BTreeMap::new(); + if let Some(account_lineages) = lineages.storage.get(&account_id) { + for (slot_name, lineage) in account_lineages { + let root = roots.get(lineage).copied().ok_or( + AccountStateForestUpdateError::MissingComputedRoot { lineage: *lineage }, + )?; + storage_map_roots.insert(slot_name.clone(), root); + } + } + + account_states.insert( + account_id, + PrecomputedPublicAccountState { vault_root, storage_map_roots }, + ); + } + + Ok(account_states) + } + fn cache_hashed_keys_from_patch(&mut self, patch: &AccountPatch) { let raw_keys = patch.storage().maps().flat_map(|(_slot_name, map_patch)| { map_patch.entries().into_iter().flat_map(|e| e.as_map().keys().copied()) @@ -216,30 +497,6 @@ impl AccountStateForest { self.storage_map_key_cache.clear(); } - fn apply_forest_updates( - &mut self, - lineage: LineageId, - block_num: BlockNumber, - operations: Vec, - ) -> Word { - let updates = if operations.is_empty() { - SmtUpdateBatch::empty() - } else { - SmtUpdateBatch::new(operations.into_iter()) - }; - let version = block_num.as_u64(); - let tree = if self.forest.latest_version(lineage).is_some() { - self.forest - .update_tree(lineage, version, updates) - .expect("forest update should succeed") - } else { - self.forest - .add_lineage(lineage, version, updates) - .expect("forest update should succeed") - }; - tree.root() - } - fn map_forest_error(error: LargeSmtForestError) -> MerkleError { match error { LargeSmtForestError::Merkle(merkle) => merkle, @@ -484,239 +741,132 @@ impl AccountStateForest { // PUBLIC INTERFACE // -------------------------------------------------------------------------------------------- - /// Updates the forest with account vault and storage changes from a patch. + /// Prepares an account-state forest update without mutating the forest. /// - /// Iterates through account updates and applies each patch to the forest. - /// Private accounts should be filtered out before calling this method. + /// The returned update is tied to `block_num` and the forest's current lineage versions and + /// roots. It must be passed exactly once to [`Self::apply_precomputed_block_update`] with the + /// same `block_num`, before any intervening forest mutation. Its precomputed account roots must + /// only be persisted as part of committing that same block. /// - /// # Arguments + /// # Errors /// - /// * `block_num` - Block number for which these updates apply - /// * `account_updates` - Iterator of `AccountPatch` for public accounts - #[miden_instrument( - target = COMPONENT, - skip_all, - fields( - block.number = %block_num, - num_pruned = tracing::field::Empty, - ), - )] - pub(crate) fn apply_block_updates( - &mut self, + /// Returns an error when a full-state patch targets an existing lineage, a computed lineage + /// root is missing, or the forest backend cannot prepare the mutation set. + pub(crate) fn compute_block_update_mutations( + &self, block_num: BlockNumber, account_updates: impl IntoIterator, - ) { - for patch in account_updates { - self.update_account(block_num, &patch); - - tracing::trace!( - target: crate::LOG_TARGET, - account_id = %patch.id(), - %block_num, - is_full_state = patch.is_full_state(), - "Updated forest with account patch" - ); - } - - let number_of_pruned_blocks = self.prune(block_num); - tracing::Span::current().record("num_pruned", number_of_pruned_blocks); + ) -> Result, AccountStateForestUpdateError> { + let account_patches = account_updates.into_iter().collect::>(); + let (batch, lineages) = self.prepare_block_update_batch(&account_patches)?; + let mutations = self.forest.compute_forest_mutations(block_num.as_u64(), batch)?; + let account_states = self.precomputed_account_states_from_mutations( + &account_patches, + &lineages, + &mutations, + )?; + + Ok(PreparedAccountStateForestBlockUpdate { + account_states, + mutations, + account_patches, + }) } - /// Updates the forest with account vault and storage changes from a patch. + /// Applies a snapshot-bound forest mutation set and updates the raw-key caches. /// - /// Unified interface for updating all account state in the forest, handling both full-state - /// patches (new accounts or reconstruction from DB) and partial patches (incremental updates - /// during block application). + /// Cache entries are added only after the forest mutation succeeds. This helper does not prune + /// forest history; callers that apply canonical blocks are responsible for pruning afterward. /// - /// Full-state patches (`patch.is_full_state() == true`) populate the forest from scratch using - /// an empty SMT root. Partial patches apply changes on top of the previous block's state. - pub(crate) fn update_account(&mut self, block_num: BlockNumber, patch: &AccountPatch) { - let account_id = patch.id(); - let is_full_state = patch.is_full_state(); - - // Apply vault changes. - if is_full_state { - self.insert_account_vault(block_num, account_id, patch.vault()); - } else if !patch.vault().is_empty() { - self.update_account_vault(block_num, account_id, patch.vault()); - } - - // Apply storage map changes. - if is_full_state { - self.insert_account_storage(block_num, account_id, patch.storage()); - } else if !patch.storage().is_empty() { - self.update_account_storage(block_num, account_id, patch.storage()); - } - - self.cache_hashed_keys_from_patch(patch); - } - - // ASSET VAULT DELTA PROCESSING - // -------------------------------------------------------------------------------------------- - - /// Retrieves the most recent vault SMT root for an account. If no vault root is found for the - /// account, returns an empty SMT root. - pub(crate) fn get_latest_vault_root(&self, account_id: AccountId) -> Word { - let lineage = Self::vault_lineage_id(account_id); - self.forest.latest_root(lineage).unwrap_or_else(empty_smt_root) - } - - /// Inserts asset vault data into the forest for the specified account. Assumes that asset vault - /// for this account does not yet exist in the forest. + /// # Errors /// - /// # Panics - /// Panics if the account's vault is already present in the forest. - fn insert_account_vault( + /// Returns an error if the forest rejects or cannot persist the prepared mutation set. + fn apply_precomputed_update( &mut self, block_num: BlockNumber, - account_id: AccountId, - vault_patch: &AccountVaultPatch, - ) { - let prev_root = self.get_latest_vault_root(account_id); - let lineage = Self::vault_lineage_id(account_id); - assert_eq!(prev_root, empty_smt_root(), "account should not be in the forest"); - assert!( - self.forest.latest_version(lineage).is_none(), - "account should not be in the forest" - ); + update: PreparedAccountStateForestBlockUpdate, + ) -> Result<(), LargeSmtForestError> { + let PreparedAccountStateForestBlockUpdate { + account_states: _, + mutations, + account_patches, + } = update; + + self.forest.apply_mutations(mutations)?; - if vault_patch.is_empty() { - let lineage = Self::vault_lineage_id(account_id); - let new_root = self.apply_forest_updates(lineage, block_num, Vec::new()); + for patch in &account_patches { + self.cache_hashed_keys_from_patch(patch); tracing::trace!( target: crate::LOG_TARGET, - %account_id, + account_id = %patch.id(), %block_num, - %new_root, - vault_entries = 0, - "Inserted vault into forest" + is_full_state = patch.is_full_state(), + "Updated forest with account patch" ); - return; } - let mut entries: Vec<(AssetId, Word)> = Vec::with_capacity(vault_patch.num_assets()); - - for (vault_key, value) in vault_patch.iter() { - entries.push((*vault_key, *value)); - } - - let num_entries = entries.len(); - - let lineage = Self::vault_lineage_id(account_id); - let operations = Self::build_forest_operations( - entries.into_iter().map(|(key, value)| (key.hash().as_word(), value)), - ); - let new_root = self.apply_forest_updates(lineage, block_num, operations); - - tracing::trace!( - target: crate::LOG_TARGET, - %account_id, - %block_num, - %new_root, - vault_entries = num_entries, - "Inserted vault into forest" - ); + Ok(()) } - /// Updates the forest with storage map changes from a patch. + /// Applies a previously prepared block update and prunes expired forest history. + /// + /// `block_num` must match the value used to prepare `update`, and the forest must not have been + /// mutated since preparation. The update is consumed and cannot be applied twice. + /// + /// If the block and [`PreparedAccountStateForestBlockUpdate::account_states`] have already been + /// committed to canonical storage, an error from this method leaves the forest inconsistent + /// with canonical state. The caller must stop normal processing and rebuild the forest before + /// serving forest-backed state; retrying the block is not a recovery mechanism. + /// + /// # Errors /// - /// Assumes that storage maps for the provided account are not in the forest already. - fn insert_account_storage( + /// Returns an error when the backend cannot apply the prepared mutations. + pub(crate) fn apply_precomputed_block_update( &mut self, block_num: BlockNumber, - account_id: AccountId, - storage_patch: &AccountStoragePatch, - ) { - for (slot_name, map_patch) in storage_patch.maps() { - // get the latest root for this map, and make sure the root is for an empty tree - let prev_root = self.get_latest_storage_map_root(account_id, slot_name); - assert_eq!(prev_root, empty_smt_root(), "account should not be in the forest"); - - let raw_map_entries: Vec<(StorageMapKey, Word)> = Vec::from_iter( - map_patch.entries().into_iter().flat_map(|e| e.as_map().iter()).filter_map( - |(&key, &value)| if value == EMPTY_WORD { None } else { Some((key, value)) }, - ), - ); - - if raw_map_entries.is_empty() { - let lineage = Self::storage_lineage_id(account_id, slot_name); - let _new_root = self.apply_forest_updates(lineage, block_num, Vec::new()); - - continue; - } - - let hashed_entries = Vec::from_iter( - raw_map_entries.iter().map(|(raw_key, value)| (raw_key.hash().into(), *value)), - ); + update: PreparedAccountStateForestBlockUpdate, + ) -> Result<(), LargeSmtForestError> { + self.apply_precomputed_update(block_num, update)?; - let lineage = Self::storage_lineage_id(account_id, slot_name); - assert!( - self.forest.latest_version(lineage).is_none(), - "account should not be in the forest" - ); - let operations = Self::build_forest_operations(hashed_entries); - let new_root = self.apply_forest_updates(lineage, block_num, operations); - - let num_entries = raw_map_entries.len(); + let number_of_pruned_blocks = self.prune(block_num); + tracing::Span::current().record("num_pruned", number_of_pruned_blocks); - tracing::trace!( - target: crate::LOG_TARGET, - %account_id, - %block_num, - ?slot_name, - %new_root, - patch_entries = num_entries, - "Inserted storage map into forest" - ); - } + Ok(()) } - // ASSET VAULT PATCH PROCESSING - // -------------------------------------------------------------------------------------------- + fn apply_account_updates_without_pruning( + &mut self, + block_num: BlockNumber, + account_updates: impl IntoIterator, + ) -> Result<(), AccountStateForestUpdateError> { + let update = self.compute_block_update_mutations(block_num, account_updates)?; + self.apply_precomputed_update(block_num, update)?; + Ok(()) + } - /// Updates the forest with the account's vault changes from the patch. + /// Rebuilds fresh account lineages from full-state account patches. /// - /// Writes the patch's absolute vault entries to the vault SMT, where an empty value removes the - /// corresponding asset. + /// Callers must ensure that every patch belongs to an account that is not already present in the + /// forest. Rebuild pages may share a version because their account lineages are disjoint. /// - /// # Panics - /// Panics if the provided vault patch is empty. - fn update_account_vault( + /// Currently used from the loader where we're loading account states from SQLite to the account + /// state forest. + pub(crate) fn apply_rebuild_updates( &mut self, block_num: BlockNumber, - account_id: AccountId, - vault_patch: &AccountVaultPatch, - ) { - assert!(!vault_patch.is_empty(), "expected the patch not to be empty"); - - let mut entries: Vec<(AssetId, Word)> = Vec::new(); - - for (vault_key, value) in vault_patch.iter() { - entries.push((*vault_key, *value)); - } - - let vault_entries = entries.len(); + account_updates: impl IntoIterator, + ) -> Result<(), AccountStateForestUpdateError> { + self.apply_account_updates_without_pruning(block_num, account_updates) + } + /// Retrieves the most recent vault SMT root for an account. If no vault root is found for the + /// account, returns an empty SMT root. + pub(crate) fn get_latest_vault_root(&self, account_id: AccountId) -> Word { let lineage = Self::vault_lineage_id(account_id); - let operations = Self::build_forest_operations( - entries.into_iter().map(|(key, value)| (key.hash().as_word(), value)), - ); - let new_root = self.apply_forest_updates(lineage, block_num, operations); - - tracing::trace!( - target: crate::LOG_TARGET, - %account_id, - %block_num, - %new_root, - %vault_entries, - "Updated vault in forest" - ); + self.forest.latest_root(lineage).unwrap_or_else(empty_smt_root) } - // STORAGE MAP DELTA PROCESSING - // -------------------------------------------------------------------------------------------- - /// Retrieves the most recent storage map SMT root for an account slot. pub(crate) fn get_latest_storage_map_root( &self, @@ -727,50 +877,6 @@ impl AccountStateForest { self.forest.latest_root(lineage).unwrap_or_else(empty_smt_root) } - /// Updates the forest with storage map changes from a patch. - /// - /// # Returns - /// - /// A map from slot name to the new storage map root for that slot. - fn update_account_storage( - &mut self, - block_num: BlockNumber, - account_id: AccountId, - storage_patch: &AccountStoragePatch, - ) { - for (slot_name, map_patch) in storage_patch.maps() { - // map patch shouldn't be empty, but if it is for some reason, there is nothing to do - let Some(entries) = map_patch.entries() else { - continue; - }; - if entries.is_empty() { - continue; - } - - // update the storage map tree in the forest and add an entry to the storage map roots - let lineage = Self::storage_lineage_id(account_id, slot_name); - let patch_entries: Vec<(StorageMapKey, Word)> = - Vec::from_iter(entries.as_map().iter().map(|(key, value)| (*key, *value))); - - let hashed_entries = Vec::from_iter( - patch_entries.iter().map(|(raw_key, value)| (raw_key.hash().into(), *value)), - ); - - let operations = Self::build_forest_operations(hashed_entries); - let new_root = self.apply_forest_updates(lineage, block_num, operations); - - tracing::trace!( - target: crate::LOG_TARGET, - %account_id, - %block_num, - ?slot_name, - %new_root, - patch_entries = patch_entries.len(), - "Updated storage map in forest" - ); - } - } - // PRUNING // -------------------------------------------------------------------------------------------- @@ -791,3 +897,16 @@ impl AccountStateForest { before.saturating_sub(after) } } + +#[cfg(test)] +pub(crate) trait TestAccountStateForestExt { + fn update_account(&mut self, block_num: BlockNumber, patch: &AccountPatch); +} + +#[cfg(test)] +impl TestAccountStateForestExt for AccountStateForest { + fn update_account(&mut self, block_num: BlockNumber, patch: &AccountPatch) { + self.apply_account_updates_without_pruning(block_num, [patch.clone()]) + .expect("test account forest update should succeed"); + } +} diff --git a/crates/store/src/account_state_forest/tests.rs b/crates/store/src/account_state_forest/tests.rs index f7f1f3182..b066b8ab8 100644 --- a/crates/store/src/account_state_forest/tests.rs +++ b/crates/store/src/account_state_forest/tests.rs @@ -1,7 +1,13 @@ use assert_matches::assert_matches; use miden_node_proto::domain::account::{AccountVaultDetails, StorageMapEntries}; use miden_protocol::Felt; -use miden_protocol::account::{AccountCode, AccountType, StorageMapKey}; +use miden_protocol::account::{ + AccountCode, + AccountStoragePatch, + AccountType, + AccountVaultPatch, + StorageMapKey, +}; use miden_protocol::asset::{ Asset, AssetVault, @@ -251,6 +257,339 @@ fn forest_versions_are_continuous_for_sequential_updates() { } } +#[test] +fn compute_block_update_mutations_does_not_mutate_forest() { + use std::collections::BTreeMap; + + use miden_protocol::account::{StorageMapPatch, StorageSlotPatch}; + + let mut forest = AccountStateForest::new(); + let account_id = dummy_account(); + let faucet_id = dummy_faucet(); + let block_num = BlockNumber::GENESIS.child(); + let slot_name = StorageSlotName::mock(11); + let raw_key = StorageMapKey::from_index(11); + let value = Word::from([11u32, 0, 0, 0]); + + let mut vault_patch = AccountVaultPatch::default(); + vault_patch.insert_asset(dummy_fungible_asset(faucet_id, 110)); + let map_patch = StorageMapPatch::from_iters([], [(raw_key, value)]); + let storage_patch = AccountStoragePatch::from_raw(BTreeMap::from_iter([( + slot_name.clone(), + StorageSlotPatch::Map(map_patch), + )])) + .unwrap(); + let patch = dummy_partial_patch(account_id, vault_patch, storage_patch); + + let prepared = forest.compute_block_update_mutations(block_num, [patch]).unwrap(); + let prepared_account_state = prepared.account_states.get(&account_id).unwrap(); + + assert!(forest.get_vault_root(account_id, block_num).is_none()); + assert!(forest.get_storage_map_root(account_id, &slot_name, block_num).is_none()); + assert_eq!(forest.forest.lineage_count(), 0); + assert_ne!(prepared_account_state.vault_root, EMPTY_WORD); + assert_eq!(prepared_account_state.storage_map_roots.len(), 1); + assert_ne!(prepared_account_state.storage_map_roots.get(&slot_name), Some(&EMPTY_WORD)); + + let expected_vault_root = prepared_account_state.vault_root; + let expected_storage_root = prepared_account_state.storage_map_roots[&slot_name]; + + forest.apply_precomputed_block_update(block_num, prepared).unwrap(); + + assert_eq!(forest.get_vault_root(account_id, block_num), Some(expected_vault_root)); + assert_eq!( + forest.get_storage_map_root(account_id, &slot_name, block_num), + Some(expected_storage_root) + ); +} + +#[test] +fn precompute_partial_empty_storage_map_create_records_empty_root() { + use std::collections::BTreeMap; + + use miden_protocol::account::{StorageMapPatch, StorageMapPatchEntries, StorageSlotPatch}; + + let mut forest = AccountStateForest::new(); + let account_id = dummy_account(); + let block_num = BlockNumber::GENESIS.child(); + let slot_name = StorageSlotName::mock(14); + + let map_patch = StorageMapPatch::Create { entries: StorageMapPatchEntries::new() }; + let storage_patch = AccountStoragePatch::from_raw(BTreeMap::from_iter([( + slot_name.clone(), + StorageSlotPatch::Map(map_patch), + )])) + .unwrap(); + let patch = dummy_partial_patch(account_id, AccountVaultPatch::default(), storage_patch); + + let prepared = forest.compute_block_update_mutations(block_num, [patch]).unwrap(); + let prepared_account_state = prepared.account_states.get(&account_id).unwrap(); + let expected_root = AccountStateForest::empty_smt_root(); + + assert_eq!(prepared_account_state.storage_map_roots.get(&slot_name), Some(&expected_root)); + + forest.apply_precomputed_block_update(block_num, prepared).unwrap(); + + assert_eq!( + forest.get_storage_map_root(account_id, &slot_name, block_num), + Some(expected_root) + ); +} + +#[test] +fn storage_map_remove_resets_forest_lineage_for_later_create() { + use std::collections::BTreeMap; + + use miden_protocol::account::{ + StorageMap, + StorageMapPatch, + StorageMapPatchEntries, + StorageSlotPatch, + }; + + let mut forest = AccountStateForest::new(); + let account_id = dummy_account(); + let slot_name = StorageSlotName::mock(15); + let old_key = StorageMapKey::from_index(15); + let new_key = StorageMapKey::from_index(16); + let old_value = Word::from([15u32, 0, 0, 0]); + let new_value = Word::from([16u32, 0, 0, 0]); + + let create_old = StorageMapPatch::Create { + entries: StorageMapPatchEntries::from_iter([(old_key, old_value)]), + }; + let old_patch = dummy_partial_patch( + account_id, + AccountVaultPatch::default(), + AccountStoragePatch::from_raw(BTreeMap::from_iter([( + slot_name.clone(), + StorageSlotPatch::Map(create_old), + )])) + .unwrap(), + ); + forest.update_account(BlockNumber::from(1u32), &old_patch); + + let remove_patch = dummy_partial_patch( + account_id, + AccountVaultPatch::default(), + AccountStoragePatch::from_raw(BTreeMap::from_iter([( + slot_name.clone(), + StorageSlotPatch::Map(StorageMapPatch::Remove), + )])) + .unwrap(), + ); + let prepared_remove = forest + .compute_block_update_mutations(BlockNumber::from(2u32), [remove_patch]) + .unwrap(); + assert!(prepared_remove.account_states[&account_id].storage_map_roots.is_empty()); + let lineage = + AccountStateForest::::storage_lineage_id(account_id, &slot_name); + assert_eq!(forest.forest.latest_version(lineage), Some(1)); + forest + .apply_precomputed_block_update(BlockNumber::from(2u32), prepared_remove) + .unwrap(); + assert_eq!(forest.forest.latest_version(lineage), Some(1)); + + let create_new = StorageMapPatch::Create { + entries: StorageMapPatchEntries::from_iter([(new_key, new_value)]), + }; + let new_patch = dummy_partial_patch( + account_id, + AccountVaultPatch::default(), + AccountStoragePatch::from_raw(BTreeMap::from_iter([( + slot_name.clone(), + StorageSlotPatch::Map(create_new), + )])) + .unwrap(), + ); + let prepared_create = forest + .compute_block_update_mutations(BlockNumber::from(3u32), [new_patch]) + .unwrap(); + let recreated_root = prepared_create.account_states[&account_id].storage_map_roots[&slot_name]; + let expected_root = StorageMap::with_entries([(new_key, new_value)]).unwrap().root(); + let stale_root = StorageMap::with_entries([(old_key, old_value), (new_key, new_value)]) + .unwrap() + .root(); + + assert_eq!(recreated_root, expected_root); + assert_ne!(recreated_root, stale_root); +} + +#[test] +fn precomputed_and_applied_roots_match_protocol_state() { + use std::collections::BTreeMap; + + use miden_protocol::account::{StorageMap, StorageMapPatch, StorageSlotPatch}; + + let account_id = dummy_account(); + let faucet_id = dummy_faucet(); + let slot_name = StorageSlotName::mock(12); + let raw_key = StorageMapKey::from_index(12); + + let mut forest = AccountStateForest::new(); + + let block_1 = BlockNumber::GENESIS.child(); + let value_1 = Word::from([12u32, 0, 0, 0]); + let asset_1 = dummy_fungible_asset(faucet_id, 120); + let expected_vault_root_1 = AssetVault::new(&[asset_1]).unwrap().root(); + let expected_map_root_1 = StorageMap::with_entries([(raw_key, value_1)]).unwrap().root(); + let mut vault_patch_1 = AccountVaultPatch::default(); + vault_patch_1.insert_asset(asset_1); + let map_patch_1 = StorageMapPatch::from_iters([], [(raw_key, value_1)]); + let storage_patch_1 = AccountStoragePatch::from_raw(BTreeMap::from_iter([( + slot_name.clone(), + StorageSlotPatch::Map(map_patch_1), + )])) + .unwrap(); + let patch_1 = dummy_partial_patch(account_id, vault_patch_1, storage_patch_1); + + let prepared_1 = forest.compute_block_update_mutations(block_1, [patch_1.clone()]).unwrap(); + + let account_state_1 = prepared_1.account_states.get(&account_id).unwrap(); + assert_eq!(account_state_1.vault_root, expected_vault_root_1); + assert_eq!(account_state_1.storage_map_roots[&slot_name], expected_map_root_1); + + forest.apply_precomputed_block_update(block_1, prepared_1).unwrap(); + assert_eq!(forest.get_vault_root(account_id, block_1), Some(expected_vault_root_1)); + assert_eq!( + forest.get_storage_map_root(account_id, &slot_name, block_1), + Some(expected_map_root_1) + ); + + let block_2 = block_1.child(); + let value_2 = Word::from([24u32, 0, 0, 0]); + let asset_2 = dummy_fungible_asset(faucet_id, 240); + let expected_vault_root_2 = AssetVault::new(&[asset_2]).unwrap().root(); + let expected_map_root_2 = StorageMap::with_entries([(raw_key, value_2)]).unwrap().root(); + let mut vault_patch_2 = AccountVaultPatch::default(); + vault_patch_2.insert_asset(asset_2); + let map_patch_2 = StorageMapPatch::from_iters([], [(raw_key, value_2)]); + let storage_patch_2 = AccountStoragePatch::from_raw(BTreeMap::from_iter([( + slot_name.clone(), + StorageSlotPatch::Map(map_patch_2), + )])) + .unwrap(); + let patch_2 = dummy_partial_patch(account_id, vault_patch_2, storage_patch_2); + + let prepared_2 = forest.compute_block_update_mutations(block_2, [patch_2.clone()]).unwrap(); + + let account_state_2 = prepared_2.account_states.get(&account_id).unwrap(); + assert_eq!(account_state_2.vault_root, expected_vault_root_2); + assert_eq!(account_state_2.storage_map_roots[&slot_name], expected_map_root_2); + + forest.apply_precomputed_block_update(block_2, prepared_2).unwrap(); + assert_eq!(forest.get_vault_root(account_id, block_2), Some(expected_vault_root_2)); + assert_eq!( + forest.get_storage_map_root(account_id, &slot_name, block_2), + Some(expected_map_root_2) + ); +} + +#[test] +fn rebuild_updates_accept_disjoint_accounts_at_the_same_version() { + let account_1 = dummy_account(); + let account_2 = + AccountId::try_from(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE_2).unwrap(); + let block_num = BlockNumber::from(12u32); + let asset_1 = dummy_fungible_asset(dummy_faucet(), 120); + let asset_2 = dummy_fungible_asset(dummy_faucet(), 240); + let expected_root_1 = AssetVault::new(&[asset_1]).unwrap().root(); + let expected_root_2 = AssetVault::new(&[asset_2]).unwrap().root(); + let patches = [ + dummy_full_state_patch(account_1, &[asset_1]), + dummy_full_state_patch(account_2, &[asset_2]), + ]; + let mut forest = AccountStateForest::new(); + + forest.apply_rebuild_updates(block_num, patches).unwrap(); + + assert_eq!(forest.get_vault_root(account_1, block_num), Some(expected_root_1)); + assert_eq!(forest.get_vault_root(account_2, block_num), Some(expected_root_2)); +} + +#[test] +fn compute_block_update_mutations_rejects_full_state_existing_lineages() { + use std::collections::BTreeMap; + + use miden_protocol::account::{StorageMapPatch, StorageMapPatchEntries, StorageSlotPatch}; + + use crate::errors::AccountStateForestUpdateError; + + let account_id = dummy_account(); + let faucet_id = dummy_faucet(); + let block_1 = BlockNumber::GENESIS.child(); + let block_2 = block_1.child(); + + let mut vault_forest = AccountStateForest::new(); + let mut vault_patch = AccountVaultPatch::default(); + vault_patch.insert_asset(dummy_fungible_asset(faucet_id, 120)); + let initial_vault_patch = + dummy_partial_patch(account_id, vault_patch, AccountStoragePatch::default()); + vault_forest.update_account(block_1, &initial_vault_patch); + + let duplicate_vault_full_state = AccountPatch::new( + account_id, + AccountStoragePatch::default(), + AccountVaultPatch::default(), + Some(AccountCode::mock()), + Some(Felt::ONE), + ) + .unwrap(); + + let Err(err) = + vault_forest.compute_block_update_mutations(block_2, [duplicate_vault_full_state]) + else { + panic!("duplicate full-state vault lineage should fail"); + }; + assert_matches!( + err, + AccountStateForestUpdateError::VaultLineageAlreadyExists { + account_id: duplicate_account_id, + } if duplicate_account_id == account_id + ); + + let mut storage_forest = AccountStateForest::new(); + let slot_name = StorageSlotName::mock(13); + let raw_key = StorageMapKey::from_index(13); + let map_patch = StorageMapPatch::from_iters([], [(raw_key, Word::from([13u32, 0, 0, 0]))]); + let storage_patch = AccountStoragePatch::from_raw(BTreeMap::from_iter([( + slot_name.clone(), + StorageSlotPatch::Map(map_patch), + )])) + .unwrap(); + let initial_storage_patch = + dummy_partial_patch(account_id, AccountVaultPatch::default(), storage_patch); + storage_forest.update_account(block_1, &initial_storage_patch); + + let empty_map_create = StorageMapPatch::Create { entries: StorageMapPatchEntries::new() }; + let duplicate_storage_patch = AccountStoragePatch::from_raw(BTreeMap::from_iter([( + slot_name.clone(), + StorageSlotPatch::Map(empty_map_create), + )])) + .unwrap(); + let duplicate_storage_full_state = AccountPatch::new( + account_id, + duplicate_storage_patch, + AccountVaultPatch::default(), + Some(AccountCode::mock()), + Some(Felt::ONE), + ) + .unwrap(); + + let Err(err) = + storage_forest.compute_block_update_mutations(block_2, [duplicate_storage_full_state]) + else { + panic!("duplicate full-state storage lineage should fail"); + }; + assert_matches!( + err, + AccountStateForestUpdateError::StorageLineageAlreadyExists { + account_id: duplicate_account_id, + slot_name: duplicate_slot_name, + } if duplicate_account_id == account_id && duplicate_slot_name == slot_name + ); +} + #[test] fn vault_state_is_not_available_for_block_gaps() { let mut forest = AccountStateForest::new(); diff --git a/crates/store/src/db/mod.rs b/crates/store/src/db/mod.rs index 745a8855b..90a0995a5 100644 --- a/crates/store/src/db/mod.rs +++ b/crates/store/src/db/mod.rs @@ -39,7 +39,11 @@ pub use crate::db::models::queries::{ PublicAccountIdsPage, PublicAccountStateRootsPage, }; -use crate::db::models::queries::{BlockHeaderCommitment, StorageMapValuesPage}; +use crate::db::models::queries::{ + BlockHeaderCommitment, + PrecomputedPublicAccountStates, + StorageMapValuesPage, +}; use crate::errors::{DatabaseError, NoteSyncError}; use crate::genesis::GenesisBlock; use crate::{COMPONENT, LOG_TARGET}; @@ -211,8 +215,15 @@ impl Db { // Insert genesis block data. let genesis_block = genesis.into_inner(); - conn.transaction(move |conn| models::queries::apply_block(conn, &genesis_block, &[])) - .context("failed to insert genesis block")?; + conn.transaction(move |conn| { + models::queries::apply_block( + conn, + &genesis_block, + &[], + &PrecomputedPublicAccountStates::new(), + ) + }) + .context("failed to insert genesis block")?; Ok(()) } @@ -575,15 +586,16 @@ impl Db { skip_all, err, )] - pub async fn apply_block( + pub(crate) async fn apply_block( &self, allow_acquire: oneshot::Sender<()>, acquire_done: oneshot::Receiver<()>, signed_block: SignedBlock, notes: Vec<(NoteRecord, Option)>, + precomputed_public_states: PrecomputedPublicAccountStates, ) -> Result<()> { self.transact("apply block", move |conn| -> Result<()> { - models::queries::apply_block(conn, &signed_block, ¬es)?; + models::queries::apply_block(conn, &signed_block, ¬es, &precomputed_public_states)?; // XXX FIXME TODO free floating mutex MUST NOT exist it doesn't bind it properly to the // data locked! diff --git a/crates/store/src/db/models/queries/accounts.rs b/crates/store/src/db/models/queries/accounts.rs index 2ab0adfb1..fa658f367 100644 --- a/crates/store/src/db/models/queries/accounts.rs +++ b/crates/store/src/db/models/queries/accounts.rs @@ -46,7 +46,7 @@ use miden_protocol::account::{ use miden_protocol::asset::{Asset, AssetId, AssetVault}; use miden_protocol::block::{BlockAccountUpdate, BlockNumber}; use miden_protocol::utils::serde::{Deserializable, Serializable}; -use miden_standards::account::auth::NetworkAccount; +use miden_standards::account::auth::{NetworkAccount, NetworkAccountNoteAllowlist}; use crate::COMPONENT; use crate::db::models::conv::{SqlTypeConvert, nonce_to_raw_sql, raw_sql_to_nonce}; @@ -61,10 +61,11 @@ pub(crate) use at_block::select_account_header_with_storage_header_at_block; mod delta; use delta::{ AccountStateForInsert, + LatestAccountStateRow, PartialAccountState, - apply_storage_patch, - select_latest_vault_assets, - select_minimal_account_state_headers, + PrecomputedFullAccountState, + apply_storage_patch_with_roots, + select_latest_account_state, }; #[cfg(test)] @@ -323,6 +324,15 @@ pub struct PublicAccountStateRootsPage { pub next_cursor: Option, } +/// Public account state commitments computed by the account state forest before SQLite writes. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct PrecomputedPublicAccountState { + pub(crate) vault_root: Word, + pub(crate) storage_map_roots: BTreeMap, +} + +pub(crate) type PrecomputedPublicAccountStates = BTreeMap; + /// Selects public account IDs with pagination. /// /// Returns up to `page_size` public account IDs, starting after `after_account_id` if provided. @@ -867,56 +877,6 @@ fn select_latest_storage_map_entries_all( group_storage_map_entries(map_values) } -fn select_latest_storage_map_entries_for_slots( - conn: &mut SqliteConnection, - account_id: &AccountId, - slot_names: &[StorageSlotName], -) -> Result>, DatabaseError> { - use schema::account_storage_map_values as t; - - if slot_names.is_empty() { - return Ok(HashMap::new()); - } - - if let [slot_name] = slot_names { - let entries = select_latest_storage_map_entries_for_slot(conn, account_id, slot_name)?; - if entries.is_empty() { - return Ok(HashMap::new()); - } - - let mut map_entries = HashMap::new(); - map_entries.insert(slot_name.clone(), entries); - return Ok(map_entries); - } - - let slot_names = Vec::from_iter(slot_names.iter().cloned().map(StorageSlotName::to_raw_sql)); - let map_values: Vec<(String, Vec, Vec)> = - SelectDsl::select(t::table, (t::slot_name, t::key, t::value)) - .filter(t::account_id.eq(&account_id.to_bytes())) - .filter(t::is_latest.eq(true)) - .filter(t::slot_name.eq_any(slot_names)) - .load(conn)?; - - group_storage_map_entries(map_values) -} - -fn select_latest_storage_map_entries_for_slot( - conn: &mut SqliteConnection, - account_id: &AccountId, - slot_name: &StorageSlotName, -) -> Result, DatabaseError> { - use schema::account_storage_map_values as t; - - let map_values: Vec<(String, Vec, Vec)> = - SelectDsl::select(t::table, (t::slot_name, t::key, t::value)) - .filter(t::account_id.eq(&account_id.to_bytes())) - .filter(t::is_latest.eq(true)) - .filter(t::slot_name.eq(slot_name.clone().to_raw_sql())) - .load(conn)?; - - Ok(group_storage_map_entries(map_values)?.remove(slot_name).unwrap_or_default()) -} - fn group_storage_map_entries( map_values: Vec<(String, Vec, Vec)>, ) -> Result>, DatabaseError> { @@ -1024,14 +984,18 @@ pub(crate) fn insert_account_vault_asset( }) } -/// Insert an account storage map value into the DB using the given [`SqliteConnection`]. +/// Inserts a versioned account storage-map value using the given [`SqliteConnection`]. /// -/// Sets `is_latest=true` for the new row and updates any existing -/// row with the same `(account_id, slot_index, key)` tuple to `is_latest=false`. +/// The new row is marked as latest, and any previous latest row for the same +/// `(account_id, slot_name, key)` tuple is invalidated first. /// /// # Returns /// -/// The number of affected rows. +/// The total number of inserted and invalidated rows. +/// +/// # Errors +/// +/// Returns an error if the previous row cannot be invalidated or the new row cannot be inserted. pub(crate) fn insert_account_storage_map_value( conn: &mut SqliteConnection, account_id: AccountId, @@ -1039,6 +1003,30 @@ pub(crate) fn insert_account_storage_map_value( slot_name: StorageSlotName, key: StorageMapKey, value: Word, +) -> Result { + insert_account_storage_map_value_inner(conn, account_id, block_num, slot_name, key, value, true) +} + +/// Inserts a versioned account storage-map value with optional previous-row invalidation. +/// +/// `invalidate_previous` may be disabled when inserting state for a new account, for which no +/// previous latest row can exist. The inserted row is always marked as latest. +/// +/// # Returns +/// +/// The total number of inserted and invalidated rows. +/// +/// # Errors +/// +/// Returns an error if the requested invalidation or insertion fails. +fn insert_account_storage_map_value_inner( + conn: &mut SqliteConnection, + account_id: AccountId, + block_num: BlockNumber, + slot_name: StorageSlotName, + key: StorageMapKey, + value: Word, + invalidate_previous: bool, ) -> Result { let account_id = account_id.to_bytes(); let key = key.to_bytes(); @@ -1046,16 +1034,20 @@ pub(crate) fn insert_account_storage_map_value( let slot_name = slot_name.to_raw_sql(); let block_num = block_num.to_raw_sql(); - let update_count = diesel::update(schema::account_storage_map_values::table) - .filter( - schema::account_storage_map_values::account_id - .eq(&account_id) - .and(schema::account_storage_map_values::slot_name.eq(&slot_name)) - .and(schema::account_storage_map_values::key.eq(&key)) - .and(schema::account_storage_map_values::is_latest.eq(true)), - ) - .set(schema::account_storage_map_values::is_latest.eq(false)) - .execute(conn)?; + let update_count = if invalidate_previous { + diesel::update(schema::account_storage_map_values::table) + .filter( + schema::account_storage_map_values::account_id + .eq(&account_id) + .and(schema::account_storage_map_values::slot_name.eq(&slot_name)) + .and(schema::account_storage_map_values::key.eq(&key)) + .and(schema::account_storage_map_values::is_latest.eq(true)), + ) + .set(schema::account_storage_map_values::is_latest.eq(false)) + .execute(conn)? + } else { + 0 + }; let record = AccountStorageMapRowInsert { account_id, @@ -1115,17 +1107,116 @@ fn prepare_full_account_update( Ok((AccountStateForInsert::FullAccount(account), storage, assets)) } -/// Prepare partial patch data for account upserts and follow-up storage and vault inserts. +/// Prepares a full public-account insertion using roots computed by the account-state forest. +/// +/// This avoids reconstructing the account's vault and storage maps in SQLite. The returned state +/// contains the account-row fields, while storage-map entries and vault assets are returned +/// separately for insertion after the account row has satisfied their foreign-key dependency. +/// Empty-word map entries and assets are omitted from the pending inserts. +/// +/// # Errors +/// +/// Returns an error if the full-state patch is missing its code or nonce, a required precomputed +/// storage root is absent, an asset is invalid, or the reconstructed account header does not match +/// the update's final state commitment. +fn prepare_precomputed_full_account_update( + update: &BlockAccountUpdate, + patch: &AccountPatch, + precomputed: &PrecomputedPublicAccountState, +) -> Result<(AccountStateForInsert, PendingStorageInserts, PendingAssetInserts), DatabaseError> { + let account_id = patch.id(); + let code = patch.code().cloned().ok_or_else(|| { + DatabaseError::DataCorrupted(format!( + "full-state patch for account {account_id} is missing account code" + )) + })?; + let nonce = patch.final_nonce().ok_or_else(|| { + DatabaseError::DataCorrupted(format!( + "full-state patch for account {account_id} is missing final nonce" + )) + })?; + + let storage_header = apply_storage_patch_with_roots( + &AccountStorageHeader::new(Vec::new())?, + patch.storage(), + &precomputed.storage_map_roots, + )?; + let account_header = miden_protocol::account::AccountHeader::new( + account_id, + nonce, + precomputed.vault_root, + storage_header.to_commitment(), + code.commitment(), + ); + if account_header.to_commitment() != update.final_state_commitment() { + return Err(DatabaseError::AccountCommitmentsMismatch { + calculated: account_header.to_commitment(), + expected: update.final_state_commitment(), + }); + } + + let storage = patch + .storage() + .maps() + .flat_map(|(slot_name, map_patch)| { + map_patch.entries().into_iter().flat_map(move |entries| { + entries + .as_map() + .iter() + .filter(|(_key, value)| **value != Word::empty()) + .map(move |(key, value)| (account_id, slot_name.clone(), *key, *value)) + }) + }) + .collect(); + let assets = patch + .vault() + .iter() + .filter(|(_asset_id, value)| **value != Word::empty()) + .map(|(asset_id, value)| { + Asset::from_id_and_value(*asset_id, *value) + .map(|asset| (account_id, *asset_id, Some(asset))) + }) + .collect::, _>>()?; + + let is_network_account = account_id.is_public() + && patch + .storage() + .maps() + .find(|(slot_name, _)| *slot_name == NetworkAccountNoteAllowlist::slot_name()) + .and_then(|(_slot_name, map_patch)| map_patch.entries()) + .is_some_and(|entries| !entries.is_empty()); + let state = PrecomputedFullAccountState { + nonce, + code, + storage_header, + vault_root: precomputed.vault_root, + is_network_account, + }; + + Ok((AccountStateForInsert::PrecomputedFullState(state), storage, assets)) +} + +/// Prepares a partial public-account update using the latest row and precomputed forest roots. +/// +/// Unchanged header fields are carried forward from `existing`. The returned partial state is used +/// for the next account row, while storage-map values and vault asset updates are returned +/// separately for insertion after that row. Empty vault values are represented as removals. +/// +/// # Errors +/// +/// Returns an error if the existing row is invalid, a required precomputed storage root is absent, +/// a patched asset is invalid, or the reconstructed account header does not match the update's +/// final state commitment. fn prepare_partial_account_update( - conn: &mut SqliteConnection, update: &BlockAccountUpdate, account_id: AccountId, patch: &AccountPatch, + precomputed: &PrecomputedPublicAccountState, + existing: &LatestAccountStateRow, ) -> Result<(AccountStateForInsert, PendingStorageInserts, PendingAssetInserts), DatabaseError> { - // Build the minimal account state needed for partial patch application. Only load the storage - // map entries that will receive updates. The next line fetches the header, which will always - // change unless the patch is empty. - let state_headers = select_minimal_account_state_headers(conn, account_id)?; + // Build the minimal account state needed for partial patch application from the latest row that + // was loaded with the account's creation metadata. + let state_headers = existing.state_headers(account_id)?; // --- Process asset updates. --------------------------------- The patch carries absolute final // values, so encode `Some` as update and `None` (an empty value word) as removal. @@ -1149,28 +1240,14 @@ fn prepare_partial_account_update( } } - // First collect entries that have associated changes. - let slot_names = Vec::from_iter(patch.storage().maps().filter_map(|(slot_name, map_patch)| { - if map_patch.entries().is_none_or(StorageMapPatchEntries::is_empty) { - None - } else { - Some(slot_name.clone()) - } - })); - - let map_entries = select_latest_storage_map_entries_for_slots(conn, &account_id, &slot_names)?; - // Apply the patch storage to the given storage header. - let new_storage_header = - apply_storage_patch(&state_headers.storage_header, patch.storage(), &map_entries)?; - - // --- Update the vault root by constructing the asset vault from DB. - let new_vault_root = { - let assets = select_latest_vault_assets(conn, account_id)?; - let mut vault = AssetVault::new(&assets)?; - vault.apply_patch(patch.vault())?; - vault.root() - }; + let new_storage_header = apply_storage_patch_with_roots( + &state_headers.storage_header, + patch.storage(), + &precomputed.storage_map_roots, + )?; + + let new_vault_root = precomputed.vault_root; // --- Compute updated account state for the accounts row. --- Use the absolute final nonce. let new_nonce = patch.final_nonce().unwrap_or(state_headers.nonce); @@ -1243,28 +1320,20 @@ pub(crate) fn upsert_accounts( conn: &mut SqliteConnection, accounts: &[BlockAccountUpdate], block_num: BlockNumber, + precomputed_public_states: &PrecomputedPublicAccountStates, ) -> Result { let mut count = 0; for update in accounts { let account_id = update.account_id(); let account_id_bytes = account_id.to_bytes(); - // Pull the latest row (if any) so we can carry forward `created_at_block` and the - // `network_account_type` classification, both of which are fixed at account creation. - let existing: Option<(i64, i32)> = QueryDsl::select( - schema::accounts::table.filter( - schema::accounts::account_id - .eq(&account_id_bytes) - .and(schema::accounts::is_latest.eq(true)), - ), - (schema::accounts::created_at_block, schema::accounts::network_account_type), - ) - .first(conn) - .optional() - .map_err(DatabaseError::Diesel)?; + // Pull the latest row once. Partial updates consume the state headers below, while every + // update carries forward creation metadata. + let existing = select_latest_account_state(conn, account_id)?; + let account_is_new = existing.is_none(); - let created_at_block = match existing { - Some((raw, _)) => BlockNumber::from_raw_sql(raw)?, + let created_at_block = match &existing { + Some(row) => row.created_at_block()?, None => block_num, }; @@ -1278,29 +1347,48 @@ pub(crate) fn upsert_accounts( // New account is always a full account, but also comes as an update AccountUpdateDetails::Public(patch) if patch.is_full_state() => { - let account = Account::try_from(patch) - .expect("Patch to full account always works for full state patches"); - debug_assert_eq!(account_id, account.id()); - - prepare_full_account_update(update, account)? + if block_num == BlockNumber::GENESIS { + let account = Account::try_from(patch) + .expect("Patch to full account always works for full state patches"); + debug_assert_eq!(account_id, account.id()); + prepare_full_account_update(update, account)? + } else { + let precomputed = + precomputed_public_states.get(&account_id).ok_or_else(|| { + DatabaseError::DataCorrupted(format!( + "missing precomputed public account state for account {account_id}" + )) + })?; + prepare_precomputed_full_account_update(update, patch, precomputed)? + } }, // Update of an existing account AccountUpdateDetails::Public(patch) => { - prepare_partial_account_update(conn, update, account_id, patch)? + let precomputed = precomputed_public_states.get(&account_id).ok_or_else(|| { + DatabaseError::DataCorrupted(format!( + "missing precomputed public account state for account {account_id}" + )) + })?; + let existing = + existing.as_ref().ok_or(DatabaseError::AccountNotFoundInDb(account_id))?; + prepare_partial_account_update(update, account_id, patch, precomputed, existing)? }, }; // Inherit the classification when the account already exists; otherwise classify it once at // creation based on the new state. - let network_account_type = match existing { - Some((_, raw)) => NetworkAccountType::from_raw_sql(raw)?, + let network_account_type = match &existing { + Some(row) => row.network_account_type()?, None => match &account_state { AccountStateForInsert::FullAccount(account) if NetworkAccount::new(account.clone()).is_ok() => { NetworkAccountType::Network }, + AccountStateForInsert::PrecomputedFullState(state) if state.is_network_account => { + NetworkAccountType::Network + }, _ => NetworkAccountType::None, }, }; @@ -1318,6 +1406,17 @@ pub(crate) fn upsert_accounts( .do_nothing() .execute(conn)?; } + if let AccountStateForInsert::PrecomputedFullState(ref state) = account_state { + let code_value = AccountCodeRowInsert { + code_commitment: state.code.commitment().to_bytes(), + code: state.code.to_bytes(), + }; + diesel::insert_into(schema::account_codes::table) + .values(&code_value) + .on_conflict(schema::account_codes::code_commitment) + .do_nothing() + .execute(conn)?; + } // mark previous rows as non-latest and insert NEW account row diesel::update(schema::accounts::table) @@ -1345,6 +1444,16 @@ pub(crate) fn upsert_accounts( created_at_block, account, ), + AccountStateForInsert::PrecomputedFullState(state) => { + AccountRowInsert::new_from_precomputed_full_state( + account_id, + network_account_type, + update.final_state_commitment(), + block_num, + created_at_block, + state, + ) + }, AccountStateForInsert::PartialState(state) => AccountRowInsert::new_from_partial( account_id, network_account_type, @@ -1364,7 +1473,13 @@ pub(crate) fn upsert_accounts( // insert pending storage map entries TODO consider batching for (acc_id, slot_name, key, value) in pending_storage_inserts { - insert_account_storage_map_value(conn, acc_id, block_num, slot_name, key, value)?; + if account_is_new { + insert_account_storage_map_value_inner( + conn, acc_id, block_num, slot_name, key, value, false, + )?; + } else { + insert_account_storage_map_value(conn, acc_id, block_num, slot_name, key, value)?; + } } for (acc_id, vault_key, update) in pending_asset_inserts { @@ -1445,6 +1560,28 @@ impl AccountRowInsert { } } + fn new_from_precomputed_full_state( + account_id: AccountId, + network_account_type: NetworkAccountType, + account_commitment: Word, + block_num: BlockNumber, + created_at_block: BlockNumber, + state: &PrecomputedFullAccountState, + ) -> Self { + Self { + account_id: account_id.to_bytes(), + network_account_type: network_account_type.to_raw_sql(), + block_num: block_num.to_raw_sql(), + account_commitment: account_commitment.to_bytes(), + code_commitment: Some(state.code.commitment().to_bytes()), + nonce: Some(nonce_to_raw_sql(state.nonce)), + storage_header: Some(state.storage_header.to_bytes()), + vault_root: Some(state.vault_root.to_bytes()), + is_latest: true, + created_at_block: created_at_block.to_raw_sql(), + } + } + /// Creates an insert row from a partial account state (patch update). fn new_from_partial( account_id: AccountId, diff --git a/crates/store/src/db/models/queries/accounts/delta.rs b/crates/store/src/db/models/queries/accounts/delta.rs index 7b9aeac8d..8bdb1c980 100644 --- a/crates/store/src/db/models/queries/accounts/delta.rs +++ b/crates/store/src/db/models/queries/accounts/delta.rs @@ -12,23 +12,27 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use diesel::query_dsl::methods::SelectDsl; use diesel::{ExpressionMethods, OptionalExtension, QueryDsl, RunQueryDsl, SqliteConnection}; +#[cfg(test)] +use miden_protocol::EMPTY_WORD; use miden_protocol::account::{ Account, + AccountCode, AccountId, AccountStorageHeader, AccountStoragePatch, - StorageMap, - StorageMapKey, StoragePatchOperation, StorageSlotHeader, StorageSlotName, StorageSlotType, }; -use miden_protocol::asset::Asset; +#[cfg(test)] +use miden_protocol::account::{StorageMap, StorageMapKey}; +use miden_protocol::block::BlockNumber; use miden_protocol::utils::serde::{Deserializable, Serializable}; -use miden_protocol::{EMPTY_WORD, Felt, Word}; +use miden_protocol::{Felt, Word}; -use crate::db::models::conv::raw_sql_to_nonce; +use super::NetworkAccountType; +use crate::db::models::conv::{SqlTypeConvert, raw_sql_to_nonce}; use crate::db::schema; use crate::errors::DatabaseError; @@ -38,16 +42,53 @@ mod tests; // TYPES // ================================================================================================ -/// Raw row type for account state delta queries. -/// -/// Fields: (`nonce`, `code_commitment`, `storage_header`) +/// Latest account row fields needed by account update preparation. #[derive(diesel::prelude::Queryable)] -struct AccountStateDeltaRow { +pub(super) struct LatestAccountStateRow { + created_at_block: i64, + network_account_type: i32, nonce: Option, code_commitment: Option>, storage_header: Option>, } +impl LatestAccountStateRow { + pub(super) fn created_at_block(&self) -> Result { + Ok(BlockNumber::from_raw_sql(self.created_at_block)?) + } + + pub(super) fn network_account_type(&self) -> Result { + Ok(NetworkAccountType::from_raw_sql(self.network_account_type)?) + } + + pub(super) fn state_headers( + &self, + account_id: AccountId, + ) -> Result { + let nonce = raw_sql_to_nonce(self.nonce.ok_or_else(|| { + DatabaseError::DataCorrupted(format!("No nonce found for account {account_id}")) + })?); + + let code_commitment = self + .code_commitment + .as_deref() + .map(Word::read_from_bytes) + .transpose()? + .ok_or_else(|| { + DatabaseError::DataCorrupted(format!( + "No code_commitment found for account {account_id}" + )) + })?; + + let storage_header = match self.storage_header.as_deref() { + Some(bytes) => AccountStorageHeader::read_from_bytes(bytes)?, + None => AccountStorageHeader::new(Vec::new())?, + }; + + Ok(AccountStateHeadersForDelta { nonce, code_commitment, storage_header }) + } +} + /// 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)] @@ -67,17 +108,25 @@ pub(super) struct PartialAccountState { pub vault_root: Word, } +/// Full account state assembled from a full-state patch and precomputed Merkle roots. +#[derive(Debug, Clone)] +pub(super) struct PrecomputedFullAccountState { + pub nonce: Felt, + pub code: AccountCode, + pub storage_header: AccountStorageHeader, + pub vault_root: Word, + pub is_network_account: bool, +} + /// Represents the account state to be inserted, either from a full account or from a partial delta /// update. -#[expect( - clippy::large_enum_variant, - reason = "built per account update and consumed immediately" -)] pub(super) enum AccountStateForInsert { /// Private account - no public state stored Private, /// Full account state (from full-state delta, i.e., new account) FullAccount(Account), + /// Full account state assembled without reconstructing its vault and storage maps. + PrecomputedFullState(PrecomputedFullAccountState), /// Partial account state (from partial delta, i.e., existing account update) PartialState(PartialAccountState), } @@ -85,27 +134,30 @@ pub(super) enum AccountStateForInsert { // QUERIES // ================================================================================================ -/// Selects the minimal account state needed for applying a delta update. +/// Selects the latest account state needed to prepare any account update. /// -/// Optimized query that only fetches: -/// - `nonce` (to add `nonce_delta`) +/// The query fetches: +/// - `created_at_block` and `network_account_type` for every update +/// - `nonce` (preserved when a partial patch omits its final nonce) /// - `code_commitment` (unchanged in partial deltas) /// - `storage_header` (to apply storage delta) /// /// # Raw SQL /// /// ```sql -/// SELECT nonce, code_commitment, storage_header +/// SELECT created_at_block, network_account_type, nonce, code_commitment, storage_header /// FROM accounts /// WHERE account_id = ?1 AND is_latest = 1 /// ``` -pub(super) fn select_minimal_account_state_headers( +pub(super) fn select_latest_account_state( conn: &mut SqliteConnection, account_id: AccountId, -) -> Result { - let row: AccountStateDeltaRow = SelectDsl::select( +) -> Result, DatabaseError> { + let row = SelectDsl::select( schema::accounts::table, ( + schema::accounts::created_at_block, + schema::accounts::network_account_type, schema::accounts::nonce, schema::accounts::code_commitment, schema::accounts::storage_header, @@ -114,59 +166,9 @@ pub(super) fn select_minimal_account_state_headers( .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())?, - }; + .optional()?; - Ok(AccountStateHeadersForDelta { nonce, code_commitment, storage_header }) -} - -/// 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, DatabaseError> { - use schema::account_vault_assets as vault; - - let entries: Vec<(Vec, Option>)> = - 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::, _>>() - .map_err(Into::into) + Ok(row) } // HELPER FUNCTIONS @@ -178,6 +180,7 @@ pub(super) fn select_latest_vault_assets( /// For map slots, uses the precomputed roots for updated maps. /// Removed slots are dropped from the header and created slots are added to it, mirroring /// [`miden_protocol::account::AccountStorage`]'s patch application. +#[cfg(test)] pub(super) fn apply_storage_patch( header: &AccountStorageHeader, patch: &AccountStoragePatch, @@ -249,3 +252,74 @@ pub(super) fn apply_storage_patch( DatabaseError::DataCorrupted(format!("Failed to create storage header: {e:?}")) }) } + +/// Applies a storage patch to an existing storage header using precomputed map roots. +/// +/// This mirrors the legacy storage patch path for value-slot updates, map-slot removal, no-op map +/// updates, and slot creation. For map slots whose final root is needed, it uses the root supplied +/// by the caller instead of loading the previous map entries and reconstructing the map. +pub(super) fn apply_storage_patch_with_roots( + header: &AccountStorageHeader, + patch: &AccountStoragePatch, + precomputed_map_roots: &BTreeMap, +) -> Result { + let mut value_updates: HashMap<&StorageSlotName, Word> = HashMap::new(); + let mut map_updates: HashMap<&StorageSlotName, Word> = HashMap::new(); + let mut removed: HashSet<&StorageSlotName> = HashSet::new(); + + for (slot_name, value_patch) in patch.values() { + match value_patch.value() { + Some(value) => { + value_updates.insert(slot_name, value); + }, + None => { + removed.insert(slot_name); + }, + } + } + + for (slot_name, map_patch) in patch.maps() { + let Some(map_patch_entries) = map_patch.entries() else { + removed.insert(slot_name); + continue; + }; + // Empty entries are a no-op for updates, but creating a map slot with no entries is + // meaningful and must still produce the slot. + if map_patch_entries.is_empty() && map_patch.patch_op() != StoragePatchOperation::Create { + continue; + } + + let root = precomputed_map_roots.get(slot_name).copied().ok_or_else(|| { + DatabaseError::DataCorrupted(format!( + "missing precomputed storage map root for slot {slot_name}" + )) + })?; + map_updates.insert(slot_name, root); + } + + let mut slots = + Vec::from_iter(header.slots().filter(|slot| !removed.contains(slot.name())).map(|slot| { + let slot_name = slot.name(); + if let Some(new_value) = value_updates.remove(slot_name) { + StorageSlotHeader::new(slot_name.clone(), slot.slot_type(), new_value) + } else if let Some(new_root) = map_updates.remove(slot_name) { + StorageSlotHeader::new(slot_name.clone(), slot.slot_type(), new_root) + } else { + slot.clone() + } + })); + + // Any updates left over belong to slots created by the patch. + for (slot_name, value) in value_updates { + slots.push(StorageSlotHeader::new(slot_name.clone(), StorageSlotType::Value, value)); + } + for (slot_name, root) in map_updates { + slots.push(StorageSlotHeader::new(slot_name.clone(), StorageSlotType::Map, root)); + } + + slots.sort_by_key(StorageSlotHeader::id); + + AccountStorageHeader::new(slots).map_err(|e| { + DatabaseError::DataCorrupted(format!("Failed to create storage header: {e:?}")) + }) +} diff --git a/crates/store/src/db/models/queries/accounts/delta/tests.rs b/crates/store/src/db/models/queries/accounts/delta/tests.rs index 1f7da9b82..54b92b199 100644 --- a/crates/store/src/db/models/queries/accounts/delta/tests.rs +++ b/crates/store/src/db/models/queries/accounts/delta/tests.rs @@ -9,6 +9,7 @@ use miden_node_utils::fee::test_fee_params; use miden_protocol::account::auth::{AuthScheme, PublicKeyCommitment}; use miden_protocol::account::component::AccountComponentMetadata; use miden_protocol::account::{ + Account, AccountBuilder, AccountComponent, AccountId, @@ -40,12 +41,15 @@ use miden_standards::account::auth::{Approver, AuthSingleSig}; use miden_standards::code_builder::CodeBuilder; use crate::db::models::queries::accounts::{ + PrecomputedPublicAccountState, + PrecomputedPublicAccountStates, select_account_header_with_storage_header_at_block, select_account_vault_at_block, select_full_account, upsert_accounts, }; use crate::db::schema::accounts; +use crate::errors::DatabaseError; fn setup_test_db() -> SqliteConnection { crate::db::migrations::test_connection() @@ -82,6 +86,129 @@ fn insert_block_header(conn: &mut SqliteConnection, block_num: BlockNumber) { .expect("Failed to insert block header"); } +fn precomputed_state_from_account(account: &Account) -> PrecomputedPublicAccountState { + PrecomputedPublicAccountState { + vault_root: account.vault().root(), + storage_map_roots: account + .storage() + .slots() + .iter() + .filter_map(|slot| match slot.content() { + miden_protocol::account::StorageSlotContent::Map(map) => { + Some((slot.name().clone(), map.root())) + }, + miden_protocol::account::StorageSlotContent::Value(_) => None, + }) + .collect(), + } +} + +fn precomputed_states_from_account(account: &Account) -> PrecomputedPublicAccountStates { + PrecomputedPublicAccountStates::from_iter([( + account.id(), + precomputed_state_from_account(account), + )]) +} + +fn callback_enabled_faucet_id() -> AccountId { + AccountId::dummy( + [41u8; 15], + AccountIdVersion::Version1, + AccountType::Public, + AssetCallbackFlag::Enabled, + ) +} + +fn callback_delta_test_account(seed: [u8; 32], slot_index: usize) -> Account { + let component_storage = + vec![StorageSlot::with_value(StorageSlotName::mock(slot_index), EMPTY_WORD)]; + let account_component_code = CodeBuilder::default() + .compile_component_code("test::interface", "@account_procedure pub proc vault push.1 end") + .unwrap(); + let component = AccountComponent::new( + account_component_code, + component_storage, + AccountComponentMetadata::new("test"), + ) + .unwrap(); + + AccountBuilder::new(seed) + .account_type(AccountType::Public) + .with_component(component) + .with_auth_component(AuthSingleSig::new(Approver::new( + PublicKeyCommitment::from(EMPTY_WORD), + AuthScheme::Falcon512Poseidon2, + ))) + .build_existing() + .unwrap() +} + +fn insert_public_account(conn: &mut SqliteConnection, block_num: BlockNumber, account: &Account) { + let patch_initial = AccountPatch::try_from(account.clone()).unwrap(); + upsert_accounts( + conn, + &[BlockAccountUpdate::new( + account.id(), + account.to_commitment(), + AccountUpdateDetails::Public(patch_initial), + )], + block_num, + &precomputed_states_from_account(account), + ) + .expect("initial upsert failed"); +} + +fn apply_callback_delta( + conn: &mut SqliteConnection, + account_id: AccountId, + faucet_id: AccountId, + block: BlockNumber, + amount: u64, + nonce_delta: u64, +) -> Account { + let prev = select_full_account(conn, account_id).expect("load account"); + let callback_template = FungibleAsset::new(faucet_id, amount).unwrap(); + let prev_amount = match prev.vault().get(callback_template.id()) { + Some(Asset::Fungible(f)) => f.amount().as_u64(), + _ => 0, + }; + + let absolute_asset = + Asset::Fungible(FungibleAsset::new(faucet_id, prev_amount + amount).unwrap()); + let mut vault_patch = AccountVaultPatch::default(); + vault_patch.insert_asset(absolute_asset); + let final_nonce = Felt::new_unchecked(prev.nonce().as_canonical_u64() + nonce_delta); + let patch = AccountPatch::new( + account_id, + AccountStoragePatch::new(), + vault_patch, + None, + Some(final_nonce), + ) + .unwrap(); + + let mut expected = prev.clone(); + expected.apply_patch(&patch).unwrap(); + let precomputed_public_states = precomputed_states_from_account(&expected); + + upsert_accounts( + conn, + &[BlockAccountUpdate::new( + account_id, + expected.to_commitment(), + AccountUpdateDetails::Public(patch), + )], + block, + &precomputed_public_states, + ) + .expect("partial delta upsert failed"); + + let after = select_full_account(conn, account_id).expect("load account after"); + assert_eq!(after.vault().root(), expected.vault().root(), "vault root mismatch"); + assert_eq!(after.to_commitment(), expected.to_commitment(), "commitment mismatch"); + after +} + /// Tests that the optimized delta update path produces the same results as the old /// method that loads the full account. /// @@ -161,7 +288,13 @@ fn optimized_delta_matches_full_account_method() { account.to_commitment(), AccountUpdateDetails::Public(patch_initial), ); - upsert_accounts(&mut conn, &[account_update_initial], block_1).expect("Initial upsert failed"); + upsert_accounts( + &mut conn, + &[account_update_initial], + block_1, + &precomputed_states_from_account(&account), + ) + .expect("Initial upsert failed"); // Verify initial state let full_account_before = @@ -229,6 +362,7 @@ fn optimized_delta_matches_full_account_method() { let final_commitment = final_account_for_commitment.to_commitment(); let expected_storage_commitment = final_account_for_commitment.storage().to_commitment(); let expected_vault_root = final_account_for_commitment.vault().root(); + let precomputed_public_states = precomputed_states_from_account(&final_account_for_commitment); // ----- Apply the partial patch via upsert_accounts (optimized path) ----- let account_update = BlockAccountUpdate::new( @@ -236,7 +370,8 @@ fn optimized_delta_matches_full_account_method() { final_commitment, AccountUpdateDetails::Public(partial_patch), ); - upsert_accounts(&mut conn, &[account_update], block_2).expect("Partial delta upsert failed"); + upsert_accounts(&mut conn, &[account_update], block_2, &precomputed_public_states) + .expect("Partial delta upsert failed"); // ----- VERIFY: Query the DB and check that optimized path produced correct results ----- @@ -363,7 +498,13 @@ fn optimized_delta_updates_non_empty_vault() { account.to_commitment(), AccountUpdateDetails::Public(patch_initial), ); - upsert_accounts(&mut conn, &[account_update_initial], block_1).expect("Initial upsert failed"); + upsert_accounts( + &mut conn, + &[account_update_initial], + block_1, + &precomputed_states_from_account(&account), + ) + .expect("Initial upsert failed"); let full_account_before = select_full_account(&mut conn, account.id()).expect("Failed to load full account"); @@ -391,13 +532,15 @@ fn optimized_delta_updates_non_empty_vault() { expected_account.apply_patch(&partial_patch).unwrap(); let expected_commitment = expected_account.to_commitment(); let expected_vault_root = expected_account.vault().root(); + let precomputed_public_states = precomputed_states_from_account(&expected_account); let account_update = BlockAccountUpdate::new( account.id(), expected_commitment, AccountUpdateDetails::Public(partial_patch), ); - upsert_accounts(&mut conn, &[account_update], block_2).expect("Partial delta upsert failed"); + upsert_accounts(&mut conn, &[account_update], block_2, &precomputed_public_states) + .expect("Partial delta upsert failed"); let vault_assets_after = select_account_vault_at_block(&mut conn, account.id(), block_2) .expect("Query vault should succeed"); @@ -435,13 +578,15 @@ fn optimized_delta_updates_non_empty_vault() { expected_after_3.apply_patch(&partial_patch_3).unwrap(); let commitment_3 = expected_after_3.to_commitment(); let expected_vault_root_3 = expected_after_3.vault().root(); + let precomputed_public_states_3 = precomputed_states_from_account(&expected_after_3); let account_update_3 = BlockAccountUpdate::new( account.id(), commitment_3, AccountUpdateDetails::Public(partial_patch_3), ); - upsert_accounts(&mut conn, &[account_update_3], block_3).expect("Block 3 upsert failed"); + upsert_accounts(&mut conn, &[account_update_3], block_3, &precomputed_public_states_3) + .expect("Block 3 upsert failed"); let full_account_final = select_full_account(&mut conn, account.id()).expect("Failed to load after block 3"); @@ -470,34 +615,6 @@ fn optimized_delta_updates_preserve_callback_flag() { let mut conn = setup_test_db(); - let faucet_id = AccountId::dummy( - [41u8; 15], - AccountIdVersion::Version1, - AccountType::Public, - AssetCallbackFlag::Enabled, - ); - - let component_storage = - vec![StorageSlot::with_value(StorageSlotName::mock(SLOT_INDEX), EMPTY_WORD)]; - let account_component_code = CodeBuilder::default() - .compile_component_code("test::interface", "@account_procedure pub proc vault push.1 end") - .unwrap(); - let component = AccountComponent::new( - account_component_code, - component_storage, - AccountComponentMetadata::new("test"), - ) - .unwrap(); - let account = AccountBuilder::new(ACCOUNT_SEED) - .account_type(AccountType::Public) - .with_component(component) - .with_auth_component(AuthSingleSig::new(Approver::new( - PublicKeyCommitment::from(EMPTY_WORD), - AuthScheme::Falcon512Poseidon2, - ))) - .build_existing() - .unwrap(); - let block_1 = BlockNumber::from(1u32); let block_2 = BlockNumber::from(2u32); let block_3 = BlockNumber::from(3u32); @@ -505,69 +622,26 @@ fn optimized_delta_updates_preserve_callback_flag() { insert_block_header(&mut conn, block_2); insert_block_header(&mut conn, block_3); - // Block 1: full-state insert of the (asset-less) account. - let patch_initial = AccountPatch::try_from(account.clone()).unwrap(); - upsert_accounts( - &mut conn, - &[BlockAccountUpdate::new( - account.id(), - account.to_commitment(), - AccountUpdateDetails::Public(patch_initial), - )], - block_1, - ) - .expect("Initial upsert failed"); - - // Applies a partial vault delta that adds `amount` of a callback-enabled asset, then asserts - // the store's recomputed state matches the protocol's reference application. - let apply_callback_delta = |conn: &mut SqliteConnection, block, amount| { - let prev = select_full_account(conn, account.id()).expect("load account"); - - let callback_template = FungibleAsset::new(faucet_id, amount).unwrap(); - let vault_key = callback_template.id(); - - // The patch carries the absolute end-state, so accumulate the previous balance for this - // asset and add the new amount on top. - let prev_amount = match prev.vault().get(vault_key) { - Some(Asset::Fungible(f)) => f.amount().as_u64(), - _ => 0, - }; - let absolute_asset = - Asset::Fungible(FungibleAsset::new(faucet_id, prev_amount + amount).unwrap()); - let mut vault_patch = AccountVaultPatch::default(); - vault_patch.insert_asset(absolute_asset); - let final_nonce = Felt::new_unchecked(prev.nonce().as_canonical_u64() + NONCE_DELTA); - let patch = AccountPatch::new( - account.id(), - AccountStoragePatch::new(), - vault_patch, - None, - Some(final_nonce), - ) - .unwrap(); - - let mut expected = prev.clone(); - expected.apply_patch(&patch).unwrap(); - - upsert_accounts( - conn, - &[BlockAccountUpdate::new( - account.id(), - expected.to_commitment(), - AccountUpdateDetails::Public(patch), - )], - block, - ) - .expect("partial delta upsert failed"); - - let after = select_full_account(conn, account.id()).expect("load account after"); - assert_eq!(after.vault().root(), expected.vault().root(), "vault root mismatch"); - assert_eq!(after.to_commitment(), expected.to_commitment(), "commitment mismatch"); - after - }; + let faucet_id = callback_enabled_faucet_id(); + let account = callback_delta_test_account(ACCOUNT_SEED, SLOT_INDEX); + insert_public_account(&mut conn, block_1, &account); - apply_callback_delta(&mut conn, block_2, ADDED_AMOUNT_BLOCK_2); - let final_account = apply_callback_delta(&mut conn, block_3, ADDED_AMOUNT_BLOCK_3); + apply_callback_delta( + &mut conn, + account.id(), + faucet_id, + block_2, + ADDED_AMOUNT_BLOCK_2, + NONCE_DELTA, + ); + let final_account = apply_callback_delta( + &mut conn, + account.id(), + faucet_id, + block_3, + ADDED_AMOUNT_BLOCK_3, + NONCE_DELTA, + ); let final_assets: Vec = final_account.vault().assets().collect(); assert_eq!(final_assets.len(), 1, "Should have exactly 1 vault asset"); @@ -656,7 +730,13 @@ fn optimized_delta_updates_storage_map_header() { account.to_commitment(), AccountUpdateDetails::Public(patch_initial), ); - upsert_accounts(&mut conn, &[account_update_initial], block_1).expect("Initial upsert failed"); + upsert_accounts( + &mut conn, + &[account_update_initial], + block_1, + &precomputed_states_from_account(&account), + ) + .expect("Initial upsert failed"); let full_account_before = select_full_account(&mut conn, account.id()).expect("Failed to load full account"); @@ -683,13 +763,15 @@ fn optimized_delta_updates_storage_map_header() { expected_account.apply_patch(&partial_patch).unwrap(); let expected_commitment = expected_account.to_commitment(); let expected_storage_commitment = expected_account.storage().to_commitment(); + let precomputed_public_states = precomputed_states_from_account(&expected_account); let account_update = BlockAccountUpdate::new( account.id(), expected_commitment, AccountUpdateDetails::Public(partial_patch), ); - upsert_accounts(&mut conn, &[account_update], block_2).expect("Partial delta upsert failed"); + upsert_accounts(&mut conn, &[account_update], block_2, &precomputed_public_states) + .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) @@ -708,6 +790,187 @@ fn optimized_delta_updates_storage_map_header() { ); } +#[test] +fn apply_storage_patch_with_roots_uses_precomputed_map_root() { + use miden_protocol::account::{AccountStorageHeader, StorageSlotHeader, StorageSlotType}; + + use super::apply_storage_patch_with_roots; + + let slot_name = StorageSlotName::mock(70); + let key = StorageMapKey::new(Word::from([Felt::new_unchecked(71); 4])); + let old_value = Word::from([Felt::new_unchecked(72); 4]); + let new_value = Word::from([Felt::new_unchecked(73); 4]); + let old_root = StorageMap::with_entries([(key, old_value)]).unwrap().root(); + let new_root = StorageMap::with_entries([(key, new_value)]).unwrap().root(); + let header = AccountStorageHeader::new(vec![StorageSlotHeader::new( + slot_name.clone(), + StorageSlotType::Map, + old_root, + )]) + .unwrap(); + let patch = AccountStoragePatch::from_raw(BTreeMap::from_iter([( + slot_name.clone(), + StorageSlotPatch::Map(StorageMapPatch::from_iters([], [(key, new_value)])), + )])) + .unwrap(); + let precomputed_roots = BTreeMap::from_iter([(slot_name.clone(), new_root)]); + + let new_header = apply_storage_patch_with_roots(&header, &patch, &precomputed_roots).unwrap(); + + let updated_slot = new_header.find_slot_header_by_name(&slot_name).unwrap(); + assert_eq!(updated_slot.value(), new_root); + assert_eq!(updated_slot.slot_type(), StorageSlotType::Map); +} + +#[test] +fn partial_public_upsert_requires_precomputed_state() { + const ACCOUNT_SEED: [u8; 32] = [80u8; 32]; + const SLOT_INDEX: usize = 0; + + let mut conn = setup_test_db(); + let block_1 = BlockNumber::from(1u32); + let block_2 = BlockNumber::from(2u32); + insert_block_header(&mut conn, block_1); + insert_block_header(&mut conn, block_2); + + let component_storage = + vec![StorageSlot::with_value(StorageSlotName::mock(SLOT_INDEX), EMPTY_WORD)]; + let account_component_code = CodeBuilder::default() + .compile_component_code( + "test::interface", + "@account_procedure pub proc required push.1 end", + ) + .unwrap(); + let component = AccountComponent::new( + account_component_code, + component_storage, + AccountComponentMetadata::new("test"), + ) + .unwrap(); + let account = AccountBuilder::new(ACCOUNT_SEED) + .account_type(AccountType::Public) + .with_component(component) + .with_auth_component(AuthSingleSig::new(Approver::new( + PublicKeyCommitment::from(EMPTY_WORD), + AuthScheme::Falcon512Poseidon2, + ))) + .build_existing() + .unwrap(); + + let patch_initial = AccountPatch::try_from(account.clone()).unwrap(); + upsert_accounts( + &mut conn, + &[BlockAccountUpdate::new( + account.id(), + account.to_commitment(), + AccountUpdateDetails::Public(patch_initial), + )], + block_1, + &precomputed_states_from_account(&account), + ) + .expect("initial full-state upsert failed"); + + let mut current_account = select_full_account(&mut conn, account.id()).unwrap(); + let patch = AccountPatch::new( + account.id(), + AccountStoragePatch::new(), + AccountVaultPatch::default(), + None, + Some(Felt::new_unchecked(current_account.nonce().as_canonical_u64() + 1)), + ) + .unwrap(); + current_account.apply_patch(&patch).unwrap(); + + let err = upsert_accounts( + &mut conn, + &[BlockAccountUpdate::new( + account.id(), + current_account.to_commitment(), + AccountUpdateDetails::Public(patch), + )], + block_2, + &PrecomputedPublicAccountStates::new(), + ) + .expect_err("partial public upsert should require precomputed roots"); + + assert_matches!(err, DatabaseError::DataCorrupted(message) if message.contains("missing precomputed public account state")); +} + +#[test] +fn partial_public_upsert_rejects_bad_precomputed_root() { + const ACCOUNT_SEED: [u8; 32] = [81u8; 32]; + const SLOT_INDEX: usize = 0; + + let mut conn = setup_test_db(); + let block_1 = BlockNumber::from(1u32); + let block_2 = BlockNumber::from(2u32); + insert_block_header(&mut conn, block_1); + insert_block_header(&mut conn, block_2); + + let component_storage = + vec![StorageSlot::with_value(StorageSlotName::mock(SLOT_INDEX), EMPTY_WORD)]; + let account_component_code = CodeBuilder::default() + .compile_component_code("test::interface", "@account_procedure pub proc badroot push.1 end") + .unwrap(); + let component = AccountComponent::new( + account_component_code, + component_storage, + AccountComponentMetadata::new("test"), + ) + .unwrap(); + let account = AccountBuilder::new(ACCOUNT_SEED) + .account_type(AccountType::Public) + .with_component(component) + .with_auth_component(AuthSingleSig::new(Approver::new( + PublicKeyCommitment::from(EMPTY_WORD), + AuthScheme::Falcon512Poseidon2, + ))) + .build_existing() + .unwrap(); + + let patch_initial = AccountPatch::try_from(account.clone()).unwrap(); + upsert_accounts( + &mut conn, + &[BlockAccountUpdate::new( + account.id(), + account.to_commitment(), + AccountUpdateDetails::Public(patch_initial), + )], + block_1, + &precomputed_states_from_account(&account), + ) + .expect("initial full-state upsert failed"); + + let mut expected_account = select_full_account(&mut conn, account.id()).unwrap(); + let patch = AccountPatch::new( + account.id(), + AccountStoragePatch::new(), + AccountVaultPatch::default(), + None, + Some(Felt::new_unchecked(expected_account.nonce().as_canonical_u64() + 1)), + ) + .unwrap(); + expected_account.apply_patch(&patch).unwrap(); + + let mut precomputed = precomputed_states_from_account(&expected_account); + precomputed.get_mut(&account.id()).unwrap().vault_root = + Word::from([Felt::new_unchecked(999); 4]); + + let err = upsert_accounts( + &mut conn, + &[BlockAccountUpdate::new( + account.id(), + expected_account.to_commitment(), + AccountUpdateDetails::Public(patch), + )], + block_2, + &precomputed, + ) + .expect_err("bad precomputed roots should be validated against the final commitment"); + + assert_matches!(err, DatabaseError::AccountCommitmentsMismatch { .. }); +} + /// Tests that a private account update (no public state) is handled correctly. /// /// Private accounts store only the account commitment, not the full state. @@ -746,8 +1009,13 @@ fn upsert_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"); + upsert_accounts( + &mut conn, + &[account_update], + block_num, + &PrecomputedPublicAccountStates::new(), + ) + .expect("Private account upsert failed"); // Verify the account exists and commitment matches @@ -830,8 +1098,13 @@ fn upsert_full_state_delta() { AccountUpdateDetails::Public(patch), ); - upsert_accounts(&mut conn, &[account_update], block_num) - .expect("Full-state delta upsert failed"); + upsert_accounts( + &mut conn, + &[account_update], + block_num, + &precomputed_states_from_account(&account), + ) + .expect("Full-state delta upsert failed"); // Verify the account state was stored correctly let (header, storage_header) = diff --git a/crates/store/src/db/models/queries/accounts/tests.rs b/crates/store/src/db/models/queries/accounts/tests.rs index 796848ec6..44a41e9b6 100644 --- a/crates/store/src/db/models/queries/accounts/tests.rs +++ b/crates/store/src/db/models/queries/accounts/tests.rs @@ -195,6 +195,28 @@ fn insert_block_header(conn: &mut SqliteConnection, block_num: BlockNumber) { .expect("Failed to insert block header"); } +fn precomputed_state_from_account(account: &Account) -> PrecomputedPublicAccountState { + PrecomputedPublicAccountState { + vault_root: account.vault().root(), + storage_map_roots: account + .storage() + .slots() + .iter() + .filter_map(|slot| match slot.content() { + StorageSlotContent::Map(map) => Some((slot.name().clone(), map.root())), + StorageSlotContent::Value(_) => None, + }) + .collect(), + } +} + +fn precomputed_states_from_account(account: &Account) -> PrecomputedPublicAccountStates { + PrecomputedPublicAccountStates::from_iter([( + account.id(), + precomputed_state_from_account(account), + )]) +} + fn create_account_with_map_storage( slot_name: StorageSlotName, entries: Vec<(StorageMapKey, Word)>, @@ -284,7 +306,13 @@ fn test_select_account_header_at_block_returns_correct_header() { AccountUpdateDetails::Public(patch), ); - upsert_accounts(&mut conn, &[account_update], block_num).expect("upsert_accounts failed"); + upsert_accounts( + &mut conn, + &[account_update], + block_num, + &PrecomputedPublicAccountStates::new(), + ) + .expect("upsert_accounts failed"); // Query the account header let (header, _storage_header) = @@ -321,7 +349,13 @@ fn test_select_account_header_at_block_historical_query() { AccountUpdateDetails::Public(patch_1), ); - upsert_accounts(&mut conn, &[account_update_1], block_num_1).expect("First upsert failed"); + upsert_accounts( + &mut conn, + &[account_update_1], + block_num_1, + &PrecomputedPublicAccountStates::new(), + ) + .expect("First upsert failed"); // Query at block 1 - should return the account let (header_1, _) = @@ -360,7 +394,13 @@ fn test_select_account_vault_at_block_empty() { AccountUpdateDetails::Public(patch), ); - upsert_accounts(&mut conn, &[account_update], block_num).expect("upsert_accounts failed"); + upsert_accounts( + &mut conn, + &[account_update], + block_num, + &PrecomputedPublicAccountStates::new(), + ) + .expect("upsert_accounts failed"); // Query vault - should return empty (the test account has no assets) let assets = select_account_vault_at_block(&mut conn, account_id, block_num) @@ -396,7 +436,12 @@ fn test_upsert_accounts_inserts_storage_header() { ); // Upsert account - let result = upsert_accounts(&mut conn, &[account_update], block_num); + let result = upsert_accounts( + &mut conn, + &[account_update], + block_num, + &PrecomputedPublicAccountStates::new(), + ); assert!(result.is_ok(), "upsert_accounts failed: {:?}", result.err()); assert_eq!(result.unwrap(), 1, "Expected 1 account to be inserted"); @@ -443,6 +488,7 @@ fn test_upsert_accounts_updates_is_latest_flag() { let account_commitment_1 = account.to_commitment(); // First update with original account - full state patch + let precomputed_1 = precomputed_states_from_account(&account); let patch_1 = AccountPatch::try_from(account).unwrap(); let account_update_1 = BlockAccountUpdate::new( @@ -451,7 +497,8 @@ fn test_upsert_accounts_updates_is_latest_flag() { AccountUpdateDetails::Public(patch_1), ); - upsert_accounts(&mut conn, &[account_update_1], block_num_1).expect("First upsert failed"); + upsert_accounts(&mut conn, &[account_update_1], block_num_1, &precomputed_1) + .expect("First upsert failed"); // Create modified account with different storage value let storage_value_modified = Word::from([ @@ -488,6 +535,7 @@ fn test_upsert_accounts_updates_is_latest_flag() { let account_commitment_2 = account_2.to_commitment(); // Second update with modified account - full state patch + let precomputed_2 = precomputed_states_from_account(&account_2); let patch_2 = AccountPatch::try_from(account_2).unwrap(); let account_update_2 = BlockAccountUpdate::new( @@ -496,7 +544,8 @@ fn test_upsert_accounts_updates_is_latest_flag() { AccountUpdateDetails::Public(patch_2), ); - upsert_accounts(&mut conn, &[account_update_2], block_num_2).expect("Second upsert failed"); + upsert_accounts(&mut conn, &[account_update_2], block_num_2, &precomputed_2) + .expect("Second upsert failed"); // Verify 2 total account rows exist (both historical records) let total_accounts: i64 = schema::accounts::table @@ -595,8 +644,13 @@ fn test_upsert_accounts_with_multiple_storage_slots() { AccountUpdateDetails::Public(patch), ); - upsert_accounts(&mut conn, &[account_update], block_num) - .expect("Upsert with multiple storage slots failed"); + upsert_accounts( + &mut conn, + &[account_update], + block_num, + &PrecomputedPublicAccountStates::new(), + ) + .expect("Upsert with multiple storage slots failed"); // Query back and verify let queried_storage = @@ -666,8 +720,13 @@ fn test_upsert_accounts_with_empty_storage() { AccountUpdateDetails::Public(patch), ); - upsert_accounts(&mut conn, &[account_update], block_num) - .expect("Upsert with empty storage failed"); + upsert_accounts( + &mut conn, + &[account_update], + block_num, + &PrecomputedPublicAccountStates::new(), + ) + .expect("Upsert with empty storage failed"); // Query back and verify let queried_storage = @@ -741,7 +800,13 @@ fn test_select_latest_account_storage_ordering_semantics() { AccountUpdateDetails::Public(patch), ); - upsert_accounts(&mut conn, &[account_update], block_num).expect("upsert_accounts failed"); + upsert_accounts( + &mut conn, + &[account_update], + block_num, + &PrecomputedPublicAccountStates::new(), + ) + .expect("upsert_accounts failed"); let storage = select_latest_account_storage(&mut conn, account_id).expect("Failed to query storage"); @@ -803,7 +868,13 @@ fn test_select_latest_account_storage_multiple_slots() { AccountUpdateDetails::Public(patch), ); - upsert_accounts(&mut conn, &[account_update], block_num).expect("upsert_accounts failed"); + upsert_accounts( + &mut conn, + &[account_update], + block_num, + &PrecomputedPublicAccountStates::new(), + ) + .expect("upsert_accounts failed"); let storage = select_latest_account_storage(&mut conn, account_id).expect("Failed to query storage"); @@ -842,7 +913,8 @@ fn test_select_latest_account_storage_slot_updates() { AccountUpdateDetails::Public(patch), ); - upsert_accounts(&mut conn, &[account_update], block_1).expect("upsert_accounts failed"); + upsert_accounts(&mut conn, &[account_update], block_1, &PrecomputedPublicAccountStates::new()) + .expect("upsert_accounts failed"); let map_patch = StorageMapPatch::from_iters([], [(key_1, value_2), (key_2, value_3)]); let storage_patch = AccountStoragePatch::from_raw(BTreeMap::from_iter([( @@ -864,6 +936,7 @@ fn test_select_latest_account_storage_slot_updates() { let mut expected_account = account.clone(); expected_account.apply_patch(&partial_patch).unwrap(); let expected_commitment = expected_account.to_commitment(); + let precomputed_public_states = precomputed_states_from_account(&expected_account); let account_update = BlockAccountUpdate::new( account_id, @@ -871,7 +944,8 @@ fn test_select_latest_account_storage_slot_updates() { AccountUpdateDetails::Public(partial_patch), ); - upsert_accounts(&mut conn, &[account_update], block_2).expect("upsert_accounts failed"); + upsert_accounts(&mut conn, &[account_update], block_2, &precomputed_public_states) + .expect("upsert_accounts failed"); let storage = select_latest_account_storage(&mut conn, account_id).expect("Failed to query storage"); @@ -921,8 +995,13 @@ fn test_select_account_vault_at_block_historical_with_updates() { ); for block in [block_1, block_2, block_3] { - upsert_accounts(&mut conn, std::slice::from_ref(&account_update), block) - .expect("upsert_accounts failed"); + upsert_accounts( + &mut conn, + std::slice::from_ref(&account_update), + block, + &precomputed_states_from_account(&account), + ) + .expect("upsert_accounts failed"); } // Insert vault asset at block 1: vault_key_1 = 1000 tokens @@ -1003,8 +1082,13 @@ fn test_select_account_vault_at_block_bounds_read_to_limit() { account.to_commitment(), AccountUpdateDetails::Public(patch), ); - upsert_accounts(&mut conn, std::slice::from_ref(&account_update), block_1) - .expect("upsert_accounts failed"); + upsert_accounts( + &mut conn, + std::slice::from_ref(&account_update), + block_1, + &PrecomputedPublicAccountStates::new(), + ) + .expect("upsert_accounts failed"); // Insert two assets more than the return limit, each with a distinct vault key. let faucet_id = AccountIdBuilder::new() @@ -1054,8 +1138,13 @@ fn test_select_account_vault_at_block_exponential_updates() { ); for block in &blocks { - upsert_accounts(&mut conn, std::slice::from_ref(&account_update), *block) - .expect("upsert_accounts failed"); + upsert_accounts( + &mut conn, + std::slice::from_ref(&account_update), + *block, + &precomputed_states_from_account(&account), + ) + .expect("upsert_accounts failed"); } let vault_key = AssetId::new_fungible(faucet_id); @@ -1112,8 +1201,13 @@ fn test_select_account_vault_at_block_with_deletion() { ); for block in [block_1, block_2, block_3] { - upsert_accounts(&mut conn, std::slice::from_ref(&account_update), block) - .expect("upsert_accounts failed"); + upsert_accounts( + &mut conn, + std::slice::from_ref(&account_update), + block, + &precomputed_states_from_account(&account), + ) + .expect("upsert_accounts failed"); } // Insert vault asset at block 1 @@ -1254,12 +1348,22 @@ fn test_prune_account_code_retains_latest_after_code_change() { let code_commitment_b = account_b.code().commitment(); // Block 0: insert account with code A. - upsert_accounts(&mut conn, &[make_full_state_update(&account_a)], block_0) - .expect("initial upsert failed"); + upsert_accounts( + &mut conn, + &[make_full_state_update(&account_a)], + block_0, + &precomputed_states_from_account(&account_a), + ) + .expect("initial upsert failed"); // Block RETENTION+1: update the same account ID to code B via a full-state delta. - upsert_accounts(&mut conn, &[make_full_state_update(&account_b)], block_code_b) - .expect("code-change upsert failed"); + upsert_accounts( + &mut conn, + &[make_full_state_update(&account_b)], + block_code_b, + &precomputed_states_from_account(&account_b), + ) + .expect("code-change upsert failed"); assert_eq!(count_account_codes(&mut conn), 2, "both codes must exist before pruning"); @@ -1326,14 +1430,29 @@ fn test_prune_account_code_retains_revisited_code() { let code_commitment_b = account_b.code().commitment(); // Block 0: code A. - upsert_accounts(&mut conn, &[make_full_state_update(&account_a)], block_0) - .expect("block 0 upsert failed"); + upsert_accounts( + &mut conn, + &[make_full_state_update(&account_a)], + block_0, + &precomputed_states_from_account(&account_a), + ) + .expect("block 0 upsert failed"); // Block RETENTION+1: code B. - upsert_accounts(&mut conn, &[make_full_state_update(&account_b)], block_code_b) - .expect("block code_b upsert failed"); + upsert_accounts( + &mut conn, + &[make_full_state_update(&account_b)], + block_code_b, + &precomputed_states_from_account(&account_b), + ) + .expect("block code_b upsert failed"); // Block RETENTION+2: back to code A. - upsert_accounts(&mut conn, &[make_full_state_update(&account_a)], block_code_a_again) - .expect("block code_a_again upsert failed"); + upsert_accounts( + &mut conn, + &[make_full_state_update(&account_a)], + block_code_a_again, + &precomputed_states_from_account(&account_a), + ) + .expect("block code_a_again upsert failed"); // Before pruning: both codes must be in account_codes (code A inserted once via ON CONFLICT DO // NOTHING, code B inserted once). diff --git a/crates/store/src/db/models/queries/mod.rs b/crates/store/src/db/models/queries/mod.rs index e4454a4da..b325a8ecb 100644 --- a/crates/store/src/db/models/queries/mod.rs +++ b/crates/store/src/db/models/queries/mod.rs @@ -53,6 +53,7 @@ pub(crate) fn apply_block( conn: &mut SqliteConnection, block: &SignedBlock, notes: &[(NoteRecord, Option)], + precomputed_public_states: &PrecomputedPublicAccountStates, ) -> Result { let mut count = 0; // Note: ordering here is important as the relevant tables have FK dependencies. @@ -66,7 +67,12 @@ pub(crate) fn apply_block( }, }; count += insert_block_header(conn, block.header(), signature)?; - count += upsert_accounts(conn, block.body().updated_accounts(), block.header().block_num())?; + count += upsert_accounts( + conn, + block.body().updated_accounts(), + block.header().block_num(), + precomputed_public_states, + )?; count += insert_scripts(conn, notes.iter().map(|(note, _)| note))?; count += insert_notes(conn, notes)?; count += insert_transactions(conn, block.header().block_num(), block.body().transactions())?; diff --git a/crates/store/src/db/tests.rs b/crates/store/src/db/tests.rs index cf9df2234..4764b6257 100644 --- a/crates/store/src/db/tests.rs +++ b/crates/store/src/db/tests.rs @@ -80,8 +80,17 @@ use rand::RngExt; use tempfile::tempdir; use super::{AccountInfo, NoteRecord, NoteSyncRecord, NullifierInfo, TransactionRecord}; -use crate::account_state_forest::{AccountStorageMapResult, HISTORICAL_BLOCK_RETENTION}; -use crate::db::models::queries::{StorageMapValue, insert_account_storage_map_value}; +use crate::account_state_forest::{ + AccountStorageMapResult, + HISTORICAL_BLOCK_RETENTION, + TestAccountStateForestExt, +}; +use crate::db::models::queries::{ + PrecomputedPublicAccountState, + PrecomputedPublicAccountStates, + StorageMapValue, + insert_account_storage_map_value, +}; use crate::db::models::{queries, utils}; use crate::errors::DatabaseError; @@ -114,6 +123,23 @@ fn create_block(conn: &mut SqliteConnection, block_num: BlockNumber) { .unwrap(); } +fn precomputed_states_from_account(account: &Account) -> PrecomputedPublicAccountStates { + let state = PrecomputedPublicAccountState { + vault_root: account.vault().root(), + storage_map_roots: account + .storage() + .slots() + .iter() + .filter_map(|slot| match slot.content() { + StorageSlotContent::Map(map) => Some((slot.name().clone(), map.root())), + StorageSlotContent::Value(_) => None, + }) + .collect(), + }; + + PrecomputedPublicAccountStates::from_iter([(account.id(), state)]) +} + #[test] #[miden_node_test_macro::enable_logging] fn sql_insert_nullifiers_for_block() { @@ -225,7 +251,13 @@ fn sql_select_note_script_by_root() { let account_id = AccountId::try_from(ACCOUNT_ID_PRIVATE_SENDER).unwrap(); - queries::upsert_accounts(conn, &[mock_block_account_update(account_id, 0)], block_num).unwrap(); + queries::upsert_accounts( + conn, + &[mock_block_account_update(account_id, 0)], + block_num, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); let new_note = create_note(account_id); @@ -270,9 +302,10 @@ fn make_account_and_note( &[BlockAccountUpdate::new( account_id, account.to_commitment(), - AccountUpdateDetails::Public(AccountPatch::try_from(account).unwrap()), + AccountUpdateDetails::Public(AccountPatch::try_from(account.clone()).unwrap()), )], block_num, + &precomputed_states_from_account(&account), ) .unwrap(); @@ -320,6 +353,7 @@ fn sql_select_accounts() { AccountUpdateDetails::Private, )], block_num, + &queries::PrecomputedPublicAccountStates::new(), ); assert_eq!(res.unwrap(), 1, "One element must have been inserted"); @@ -347,8 +381,13 @@ fn sync_account_vault_basic_validation() { create_block(conn, block_to); for block in [block_from, block_mid, block_to] { - queries::upsert_accounts(conn, &[mock_block_account_update(public_account_id, 0)], block) - .unwrap(); + queries::upsert_accounts( + conn, + &[mock_block_account_update(public_account_id, 0)], + block, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); } // Create test vault assets from two different faucets to get different vault keys. @@ -679,7 +718,13 @@ fn notes() { // test insertion - queries::upsert_accounts(conn, &[mock_block_account_update(sender, 0)], block_num_1).unwrap(); + queries::upsert_accounts( + conn, + &[mock_block_account_update(sender, 0)], + block_num_1, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); let new_note = create_note(sender); let note_index = BlockNoteIndex::new(0, 2).unwrap(); @@ -785,6 +830,7 @@ fn note_sync_across_multiple_blocks() { conn, &[mock_block_account_update(sender, block_num_raw.into())], block_num, + &queries::PrecomputedPublicAccountStates::new(), ) .unwrap(); @@ -866,6 +912,7 @@ fn note_sync_multi_respects_payload_limit() { conn, &[mock_block_account_update(sender, block_num_raw.into())], block_num, + &queries::PrecomputedPublicAccountStates::new(), ) .unwrap(); @@ -922,7 +969,13 @@ fn note_sync_no_matching_tags() { let sender = AccountId::try_from(ACCOUNT_ID_PRIVATE_SENDER).unwrap(); let block_num = BlockNumber::from(1); create_block(conn, block_num); - queries::upsert_accounts(conn, &[mock_block_account_update(sender, 0)], block_num).unwrap(); + queries::upsert_accounts( + conn, + &[mock_block_account_update(sender, 0)], + block_num, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); // Insert a note with tag 10. let new_note = create_note(sender); @@ -999,8 +1052,20 @@ fn sql_account_storage_map_values_insertion() { let account_id = AccountId::try_from(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE_2).unwrap(); - queries::upsert_accounts(conn, &[mock_block_account_update(account_id, 0)], block1).unwrap(); - queries::upsert_accounts(conn, &[mock_block_account_update(account_id, 0)], block2).unwrap(); + queries::upsert_accounts( + conn, + &[mock_block_account_update(account_id, 0)], + block1, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); + queries::upsert_accounts( + conn, + &[mock_block_account_update(account_id, 0)], + block2, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); let slot_name = StorageSlotName::mock(3); let key1 = StorageMapKey::new(Word::from([1u32, 2, 3, 4])); @@ -1090,8 +1155,13 @@ fn select_storage_map_sync_values() { let block3 = BlockNumber::from(3); for block in [block1, block2, block3] { - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 0)], block) - .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 0)], + block, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); } // Insert data across multiple blocks using individual inserts Block 1: key1 -> value1, key2 -> @@ -1232,12 +1302,27 @@ fn select_storage_map_sync_values_paginates_until_last_block() { create_block(&mut conn, block2); create_block(&mut conn, block3); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 0)], block1) - .unwrap(); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 1)], block2) - .unwrap(); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 2)], block3) - .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 0)], + block1, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 1)], + block2, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 2)], + block3, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); queries::insert_account_storage_map_value( &mut conn, @@ -1291,8 +1376,13 @@ fn select_storage_map_sync_values_all_entries_in_genesis_block() { let genesis = BlockNumber::GENESIS; create_block(&mut conn, genesis); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 0)], genesis) - .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 0)], + genesis, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); // Insert 3 entries, all in genesis block for i in 0..3 { @@ -1339,8 +1429,13 @@ fn select_storage_map_sync_values_all_entries_in_single_non_genesis_block() { let block5 = BlockNumber::from(5); create_block(&mut conn, block5); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 0)], block5) - .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 0)], + block5, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); for i in 0..3 { queries::insert_account_storage_map_value( @@ -1379,12 +1474,27 @@ fn select_storage_map_sync_values_multi_block_pagination() { create_block(&mut conn, block2); create_block(&mut conn, block3); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 0)], block1) - .unwrap(); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 1)], block2) - .unwrap(); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 2)], block3) - .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 0)], + block1, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 1)], + block2, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 2)], + block3, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); // 1 entry in block 1, 1 in block 2, 1 in block 3 queries::insert_account_storage_map_value( @@ -1450,9 +1560,24 @@ async fn reconstruct_storage_map_from_db_pages_until_latest() { create_block(db_conn, block2); create_block(db_conn, block3); - queries::upsert_accounts(db_conn, &[mock_block_account_update(account_id, 0)], block1)?; - queries::upsert_accounts(db_conn, &[mock_block_account_update(account_id, 1)], block2)?; - queries::upsert_accounts(db_conn, &[mock_block_account_update(account_id, 2)], block3)?; + queries::upsert_accounts( + db_conn, + &[mock_block_account_update(account_id, 0)], + block1, + &queries::PrecomputedPublicAccountStates::new(), + )?; + queries::upsert_accounts( + db_conn, + &[mock_block_account_update(account_id, 1)], + block2, + &queries::PrecomputedPublicAccountStates::new(), + )?; + queries::upsert_accounts( + db_conn, + &[mock_block_account_update(account_id, 2)], + block3, + &queries::PrecomputedPublicAccountStates::new(), + )?; queries::insert_account_storage_map_value( db_conn, @@ -1516,7 +1641,12 @@ async fn reconstruct_storage_map_from_db_returns_limit_exceeded_for_single_block db_conn.transaction(|db_conn| { create_block(db_conn, block5); - queries::upsert_accounts(db_conn, &[mock_block_account_update(account_id, 0)], block5)?; + queries::upsert_accounts( + db_conn, + &[mock_block_account_update(account_id, 0)], + block5, + &queries::PrecomputedPublicAccountStates::new(), + )?; // Insert 3 entries, all in the same block for i in 1..=3 { @@ -1683,7 +1813,13 @@ fn insert_transactions(conn: &mut SqliteConnection) -> usize { mock_block_transaction(AccountId::try_from(ACCOUNT_ID_PRIVATE_SENDER).unwrap(), 2); let ordered_tx_headers = OrderedTransactionHeaders::new_unchecked(vec![mock_tx1, mock_tx2]); - queries::upsert_accounts(conn, &account_updates, block_num).unwrap(); + queries::upsert_accounts( + conn, + &account_updates, + block_num, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); let count = queries::insert_transactions(conn, block_num, &ordered_tx_headers).unwrap(); Ok::<_, DatabaseError>(count) @@ -1760,9 +1896,10 @@ fn test_select_account_code_by_commitment() { &[BlockAccountUpdate::new( account.id(), account.to_commitment(), - AccountUpdateDetails::Public(AccountPatch::try_from(account).unwrap()), + AccountUpdateDetails::Public(AccountPatch::try_from(account.clone()).unwrap()), )], block_num_1, + &precomputed_states_from_account(&account), ) .unwrap(); @@ -1809,9 +1946,10 @@ fn test_select_account_code_by_commitment_multiple_codes() { &[BlockAccountUpdate::new( account_v1.id(), account_v1.to_commitment(), - AccountUpdateDetails::Public(AccountPatch::try_from(account_v1).unwrap()), + AccountUpdateDetails::Public(AccountPatch::try_from(account_v1.clone()).unwrap()), )], block_num_1, + &precomputed_states_from_account(&account_v1), ) .unwrap(); @@ -1843,9 +1981,10 @@ fn test_select_account_code_by_commitment_multiple_codes() { &[BlockAccountUpdate::new( account_v2.id(), account_v2.to_commitment(), - AccountUpdateDetails::Public(AccountPatch::try_from(account_v2).unwrap()), + AccountUpdateDetails::Public(AccountPatch::try_from(account_v2.clone()).unwrap()), )], block_num_2, + &precomputed_states_from_account(&account_v2), ) .unwrap(); @@ -2162,7 +2301,13 @@ fn regression_1461_full_state_delta_inserts_vault_assets() { AccountUpdateDetails::Public(account_patch), ); - queries::upsert_accounts(&mut conn, &[block_update], block_num).unwrap(); + queries::upsert_accounts( + &mut conn, + &[block_update], + block_num, + &precomputed_states_from_account(&account), + ) + .unwrap(); let (_, vault_assets) = queries::select_account_vault_assets( &mut conn, @@ -2376,7 +2521,13 @@ fn db_roundtrip_account() { account_commitment, AccountUpdateDetails::Public(account_patch), ); - queries::upsert_accounts(&mut conn, &[block_update], block_num).unwrap(); + queries::upsert_accounts( + &mut conn, + &[block_update], + block_num, + &precomputed_states_from_account(&account), + ) + .unwrap(); // Retrieve let retrieved = queries::select_all_accounts(&mut conn).unwrap(); @@ -2402,8 +2553,13 @@ fn db_roundtrip_notes() { create_block(&mut conn, block_num); let sender = AccountId::try_from(ACCOUNT_ID_PRIVATE_SENDER).unwrap(); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(sender, 0)], block_num) - .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(sender, 0)], + block_num, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); let new_note = create_note(sender); let note_index = BlockNoteIndex::new(0, 0).unwrap(); @@ -2455,8 +2611,13 @@ fn db_roundtrip_vault_assets() { let account_id = AccountId::try_from(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE).unwrap(); // Create account first - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 0)], block_num) - .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 0)], + block_num, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); let fungible_asset = FungibleAsset::new(faucet_id, 5000).unwrap(); let asset: Asset = fungible_asset.into(); @@ -2490,14 +2651,24 @@ fn db_roundtrip_storage_map_values() { create_block(&mut conn, block_num); let account_id = AccountId::try_from(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE).unwrap(); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 0)], block_num) - .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 0)], + block_num, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); let slot_name = StorageSlotName::mock(5); let key = StorageMapKey::from_index(12345u32); let value = num_to_word(67890); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 1)], block_num) - .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 1)], + block_num, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); // Insert queries::insert_account_storage_map_value( @@ -2598,7 +2769,13 @@ fn db_roundtrip_account_storage_with_maps() { account.to_commitment(), AccountUpdateDetails::Public(account_patch), ); - queries::upsert_accounts(&mut conn, &[block_update], block_num).unwrap(); + queries::upsert_accounts( + &mut conn, + &[block_update], + block_num, + &precomputed_states_from_account(&account), + ) + .unwrap(); // Retrieve the storage using select_latest_account_storage (reconstructs from header + map // values) @@ -2731,8 +2908,13 @@ fn test_prune_history() { // Create account for block in [block_0, block_old, block_cutoff, block_update, block_tip] { - queries::upsert_accounts(conn, &[mock_block_account_update(public_account_id, 0)], block) - .unwrap(); + queries::upsert_accounts( + conn, + &[mock_block_account_update(public_account_id, 0)], + block, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); } // Insert vault assets at different blocks - use different faucets for different vault keys. @@ -3035,12 +3217,27 @@ fn account_state_forest_matches_db_storage_map_roots_across_updates() { create_block(&mut conn, block2); create_block(&mut conn, block3); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 0)], block1) - .unwrap(); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 1)], block2) - .unwrap(); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(account_id, 2)], block3) - .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 0)], + block1, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 1)], + block2, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(account_id, 2)], + block3, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); let slot_map = StorageSlotName::mock(1); let slot_value = StorageSlotName::mock(2); @@ -3470,7 +3667,13 @@ fn db_roundtrip_transactions() { create_block(&mut conn, block_num); let bob = AccountId::try_from(ACCOUNT_ID_PRIVATE_SENDER).unwrap(); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(bob, 0)], block_num).unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(bob, 0)], + block_num, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); let tx = mock_block_transaction(bob, 1); let ordered = OrderedTransactionHeaders::new_unchecked(vec![tx.clone()]); @@ -3534,7 +3737,13 @@ fn db_roundtrip_transactions_filters_missing_output_note_sync_records() { create_block(&mut conn, block_num); let bob = AccountId::try_from(ACCOUNT_ID_PRIVATE_SENDER).unwrap(); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(bob, 0)], block_num).unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(bob, 0)], + block_num, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); let tx = mock_block_transaction(bob, 1); let ordered = OrderedTransactionHeaders::new_unchecked(vec![tx.clone()]); @@ -3570,7 +3779,13 @@ fn select_transactions_records_resolves_consumed_public_note_refs() { create_block(&mut conn, block_num); let bob = AccountId::try_from(ACCOUNT_ID_PRIVATE_SENDER).unwrap(); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(bob, 0)], block_num).unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(bob, 0)], + block_num, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); // `mock_block_transaction` consumes a single authenticated input whose nullifier is // `num_to_nullifier(num)`. Record a public note carrying that same nullifier so the node can @@ -3619,8 +3834,20 @@ fn select_transactions_records_reports_truncation_below_payload_cap() { let block2 = BlockNumber::from(2); create_block(&mut conn, block1); create_block(&mut conn, block2); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(bob, 0)], block1).unwrap(); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(bob, 1)], block2).unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(bob, 0)], + block1, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(bob, 1)], + block2, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); let cap = miden_node_utils::limiter::MAX_RESPONSE_PAYLOAD_BYTES; // `block1` nearly fills the cap on its own, then `block2` pushes the running total past it. The @@ -3663,7 +3890,13 @@ fn select_transactions_records_errors_when_single_block_exceeds_payload_cap() { let block1 = BlockNumber::from(1); create_block(&mut conn, block1); - queries::upsert_accounts(&mut conn, &[mock_block_account_update(bob, 0)], block1).unwrap(); + queries::upsert_accounts( + &mut conn, + &[mock_block_account_update(bob, 0)], + block1, + &queries::PrecomputedPublicAccountStates::new(), + ) + .unwrap(); let cap = miden_node_utils::limiter::MAX_RESPONSE_PAYLOAD_BYTES; let oversized_notes = cap / OUTPUT_NOTE_SIZE_BYTES + 100; diff --git a/crates/store/src/errors.rs b/crates/store/src/errors.rs index c74c1298f..0969a58ad 100644 --- a/crates/store/src/errors.rs +++ b/crates/store/src/errors.rs @@ -8,6 +8,7 @@ use miden_protocol::account::AccountId; use miden_protocol::block::BlockNumber; use miden_protocol::crypto::merkle::MerkleError; use miden_protocol::crypto::merkle::mmr::MmrError; +use miden_protocol::crypto::merkle::smt::{LargeSmtForestError, LineageId}; use miden_protocol::crypto::utils::DeserializationError; use miden_protocol::errors::{ AccountDeltaError, @@ -26,6 +27,29 @@ use tokio::sync::oneshot::error::RecvError; use crate::db::models::conv::DatabaseTypeConversionError; +/// Errors produced while preparing or rebuilding account-state forest updates. +/// +/// The underlying [`LargeSmtForestError`] is preserved so callers can distinguish fatal backend +/// failures from invalid update preparation. +#[derive(Debug, Error)] +pub enum AccountStateForestUpdateError { + /// A full-state patch attempted to create a vault lineage that is already present. + #[error("account {account_id} vault lineage already exists")] + VaultLineageAlreadyExists { account_id: AccountId }, + /// A full-state patch attempted to create a storage-map lineage that is already present. + #[error("account {account_id} storage map lineage for slot {slot_name} already exists")] + StorageLineageAlreadyExists { + account_id: AccountId, + slot_name: miden_protocol::account::StorageSlotName, + }, + /// The forest mutation set did not contain a root for a lineage it was expected to update. + #[error("computed forest mutations are missing lineage {lineage}")] + MissingComputedRoot { lineage: LineageId }, + /// The forest backend failed while computing, reading, or applying mutations. + #[error(transparent)] + Forest(#[from] LargeSmtForestError), +} + // DATABASE ERRORS // ================================================================================================= @@ -98,6 +122,8 @@ pub enum StateInitializationError { NullifierTreeIoError(String), #[error("account state forest IO error: {0}")] AccountStateForestIoError(String), + #[error("failed to rebuild account state forest")] + AccountStateForestRebuild(#[source] AccountStateForestUpdateError), #[error("database error")] DatabaseError(#[from] DatabaseError), #[error("failed to create nullifier tree")] @@ -209,6 +235,8 @@ pub enum ApplyBlockError { ClosedChannel(#[from] RecvError), #[error("concurrent write detected")] ConcurrentWrite, + #[error("account state forest update preparation failed")] + AccountStateForestPreparation(#[source] AccountStateForestUpdateError), #[error("database doesn't have any block header data")] DbBlockHeaderEmpty, #[error("database update failed: {0}")] diff --git a/crates/store/src/lib.rs b/crates/store/src/lib.rs index 657509886..416e08736 100644 --- a/crates/store/src/lib.rs +++ b/crates/store/src/lib.rs @@ -26,6 +26,7 @@ pub use db::{ TransactionRecord, }; pub use errors::{ + AccountStateForestUpdateError, ApplyBlockError, ApplyBlockWithProvingInputsError, DatabaseError, diff --git a/crates/store/src/state/apply_block.rs b/crates/store/src/state/apply_block.rs index e2adf4e70..9ee27b4e8 100644 --- a/crates/store/src/state/apply_block.rs +++ b/crates/store/src/state/apply_block.rs @@ -1,3 +1,4 @@ +use std::fmt::Display; use std::sync::Arc; use miden_node_proto::domain::proof_request::BlockProofRequest; @@ -17,7 +18,7 @@ use tracing::{Instrument, info_span}; use crate::db::NoteRecord; use crate::errors::{ApplyBlockError, ApplyBlockWithProvingInputsError, InvalidBlockError}; -use crate::state::{BlockNotification, State}; +use crate::state::{BlockNotification, InnerState, State}; use crate::{COMPONENT, HistoricalError, LOG_TARGET}; impl State { @@ -69,13 +70,16 @@ impl State { /// - a transaction is open in the DB and the writes are started. /// - while the transaction is not committed, concurrent reads are allowed, both the DB and the /// in-memory representations, which are consistent at this stage. - /// - prior to committing the changes to the DB, an exclusive lock to the in-memory data is - /// acquired, preventing concurrent reads to the in-memory data, since that will be - /// out-of-sync w.r.t. the DB. + /// - prior to committing the changes to the DB, exclusive locks to the canonical in-memory + /// state and account-state forest are acquired, preventing readers from observing them at + /// different block heights. /// - the DB transaction is committed, and requests that read only from the DB can proceed to /// use the fresh data. - /// - the in-memory structures are updated, including the latest block pointer and the lock is - /// released. + /// - the account-state forest and canonical in-memory structures are updated, including the + /// latest block pointer, and both locks are released. + /// - if any persistent tree update fails after the DB commit, the process aborts. The durable + /// database remains authoritative, and divergent tree storage must be rebuilt before the node + /// resumes normal processing. // TODO: This span is logged in a root span, we should connect it to the parent span. #[miden_instrument( target = COMPONENT, @@ -134,14 +138,29 @@ impl State { AccountUpdateDetails::Private => None, }, )); + let account_forest_update = self.with_forest_read_blocking(|forest| { + forest + .compute_block_update_mutations(block_num, account_patches) + .map_err(ApplyBlockError::AccountStateForestPreparation) + })?; + let precomputed_public_states = account_forest_update.account_states.clone(); // The DB and in-memory state updates need to be synchronized and are partially 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 = Arc::clone(&self.db); let db_update_task = tokio::spawn( - async move { db.apply_block(allow_acquire, acquire_done, signed_block, notes).await } - .in_current_span(), + async move { + db.apply_block( + allow_acquire, + acquire_done, + signed_block, + notes, + precomputed_public_states, + ) + .await + } + .in_current_span(), ); // Wait for the message from the DB update task, that we ready to commit the DB transaction. @@ -153,7 +172,7 @@ impl State { // Awaiting the block saving task to complete without errors. block_save_task.await??; - self.with_inner_write_blocking(|inner| { + self.with_inner_and_forest_write_blocking(|inner, forest| { // We need to check that neither the nullifier tree nor the account tree have changed // while we were waiting for the DB preparation task to complete. If either of them did // change, we do not proceed with in-memory and database updates, since it may lead to @@ -177,25 +196,25 @@ impl State { .block_on(db_update_task)? .map_err(|err| ApplyBlockError::DbUpdateTaskFailed(err.as_report()))?; - // Update the in-memory data structures after successful commit of the DB transaction - 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.blockchain.push(block_commitment); + // The DB is now committed. Keep both write locks held while advancing the forest and + // canonical in-memory state so readers cannot observe different block heights. + let InnerState { nullifier_tree, blockchain, account_tree } = inner; + forest + .apply_precomputed_block_update(block_num, account_forest_update) + .unwrap_or_else(|error| { + Self::abort_after_post_commit_failure("account-state forest", &error) + }); + nullifier_tree.apply_mutations(nullifier_tree_update).unwrap_or_else(|error| { + Self::abort_after_post_commit_failure("nullifier tree", &error) + }); + account_tree.apply_mutations(account_tree_update).unwrap_or_else(|error| { + Self::abort_after_post_commit_failure("account tree", &error) + }); + blockchain.push(block_commitment); Ok(()) })?; - self.with_forest_write_blocking(|forest| { - forest.apply_block_updates(block_num, account_patches); - }); - // Push to cache and notify replica subscribers. self.block_cache .push(block_num, BlockNotification::new(block_num, cache_bytes)) @@ -207,6 +226,25 @@ impl State { Ok(()) } + /// Terminates after a persistent state failure that occurred after the canonical DB commit. + /// + /// Returning would expose components at different block heights. Tests panic so the fatal path + /// can be asserted without terminating the test process; production aborts immediately. + fn abort_after_post_commit_failure(component: &str, error: &impl Display) -> ! { + tracing::error!( + target: LOG_TARGET, + component, + error = %error, + "persistent state update failed after database commit; aborting" + ); + + #[cfg(test)] + panic!("{component} update failed after database commit: {error}"); + + #[cfg(not(test))] + std::process::abort(); + } + /// Saves the proving inputs for the given block to the block store. pub async fn save_proving_inputs( &self, diff --git a/crates/store/src/state/loader.rs b/crates/store/src/state/loader.rs index d573c9b4a..49df6c420 100644 --- a/crates/store/src/state/loader.rs +++ b/crates/store/src/state/loader.rs @@ -558,7 +558,7 @@ pub async fn rebuild_account_state_forest( break; } - // Process each account in this page + let mut patches = Vec::with_capacity(page.account_ids.len()); for account_id in page.account_ids { // TODO: Loading the full account from the database is inefficient and will need to go // away. @@ -572,9 +572,13 @@ pub async fn rebuild_account_state_forest( StateInitializationError::AccountToDeltaConversionFailed(e.to_string()) })?; - forest.update_account(block_num, &patch); + patches.push(patch); } + forest + .apply_rebuild_updates(block_num, patches) + .map_err(StateInitializationError::AccountStateForestRebuild)?; + cursor = page.next_cursor; if cursor.is_none() { break; diff --git a/crates/store/src/state/mod.rs b/crates/store/src/state/mod.rs index b7f39fde8..185341a59 100644 --- a/crates/store/src/state/mod.rs +++ b/crates/store/src/state/mod.rs @@ -311,15 +311,24 @@ impl State { }) } - /// Runs a synchronous mutable operation over the inner state on Tokio's blocking path. + /// Runs a synchronous mutable operation while holding both in-memory state write locks. /// - /// See [`Self::with_inner_read_blocking`] for why this uses `block_in_place`. - fn with_inner_write_blocking(&self, f: impl FnOnce(&mut InnerState) -> R) -> R { + /// Locks are always acquired in inner-state then forest order. Holding both across the database + /// commit window prevents readers from observing canonical tree state from one block and + /// account-state forest data from another. + fn with_inner_and_forest_write_blocking( + &self, + f: impl FnOnce( + &mut InnerState, + &mut AccountStateForest, + ) -> R, + ) -> R { let span = Span::current(); tokio::task::block_in_place(|| { span.in_scope(|| { let mut inner = self.inner.blocking_write(); - f(&mut inner) + let mut forest = self.forest.blocking_write(); + f(&mut inner, &mut forest) }) }) } @@ -342,22 +351,6 @@ impl State { }) } - /// Runs a synchronous mutable operation over the account state forest on Tokio's blocking path. - /// - /// See [`Self::with_forest_read_blocking`] for why this uses `block_in_place`. - fn with_forest_write_blocking( - &self, - f: impl FnOnce(&mut AccountStateForest) -> R, - ) -> R { - let span = Span::current(); - tokio::task::block_in_place(|| { - span.in_scope(|| { - let mut forest = self.forest.blocking_write(); - f(&mut forest) - }) - }) - } - // STATE ACCESSORS // --------------------------------------------------------------------------------------------