1use std::collections::{BTreeMap, BTreeSet, HashSet};
7use std::num::NonZeroUsize;
8use std::path::{Path, PathBuf};
9use std::sync::Arc;
10
11use miden_node_proto::domain::batch::BatchInputs;
12use miden_node_utils::clap::StorageOptions;
13use miden_node_utils::formatting::format_array;
14use miden_node_utils::tracing::miden_instrument;
15use miden_protocol::Word;
16use miden_protocol::account::AccountId;
17use miden_protocol::block::account_tree::AccountWitness;
18use miden_protocol::block::nullifier_tree::{NullifierTree, NullifierWitness};
19use miden_protocol::block::{BlockHeader, BlockInputs, BlockNumber, Blockchain};
20use miden_protocol::crypto::merkle::mmr::{MmrProof, PartialMmr};
21use miden_protocol::crypto::merkle::smt::{LargeSmt, SmtStorage};
22use miden_protocol::note::{NoteId, NoteScript, Nullifier};
23use miden_protocol::transaction::PartialBlockchain;
24use tokio::sync::{Mutex, RwLock, watch};
25use tracing::{Instrument, Span};
26
27use crate::account_state_forest::{AccountStateForest, AccountStateForestBackend};
28use crate::accounts::AccountTreeWithHistory;
29use crate::blocks::BlockStore;
30use crate::db::{Db, NoteRecord, NullifierInfo};
31use crate::errors::{
32 DatabaseError,
33 GetBatchInputsError,
34 GetBlockHeaderError,
35 GetBlockInputsError,
36 StateInitializationError,
37};
38use crate::proven_tip::ProvenTipWriter;
39use crate::{COMPONENT, DataDirectory, DatabaseOptions};
40
41const BLOCK_CACHE_CAPACITY: NonZeroUsize = NonZeroUsize::new(512).unwrap();
43
44const PROOF_CACHE_CAPACITY: NonZeroUsize = NonZeroUsize::new(512).unwrap();
46
47mod loader;
48use loader::{
49 ACCOUNT_STATE_FOREST_STORAGE_DIR,
50 ACCOUNT_TREE_STORAGE_DIR,
51 AccountForestLoader,
52 NULLIFIER_TREE_STORAGE_DIR,
53 TreeStorage,
54 TreeStorageLoader,
55 load_mmr,
56 verify_account_state_forest_consistency,
57 verify_tree_consistency,
58};
59
60mod replica;
61pub use replica::{BlockCache, BlockNotification, ProofCache, ProofNotification};
62
63mod account;
64
65mod apply_block;
66mod apply_proof;
67mod block_lifecycle;
68mod bootstrap;
69mod disk_monitor;
70mod sync_state;
71
72#[derive(Debug, Clone, Copy)]
77pub enum Finality {
78 Committed,
80 Proven,
82}
83
84#[derive(Debug, Default)]
88pub struct TransactionInputs {
89 pub account_commitment: Word,
90 pub nullifiers: Vec<NullifierInfo>,
91 pub found_unauthenticated_notes: HashSet<Word>,
92 pub new_account_id_prefix_is_unique: Option<bool>,
93}
94
95type BlockInputWitnesses = (
96 BlockNumber,
97 BTreeMap<AccountId, AccountWitness>,
98 BTreeMap<Nullifier, NullifierWitness>,
99 PartialMmr,
100);
101
102struct InnerState<S>
104where
105 S: SmtStorage,
106{
107 nullifier_tree: NullifierTree<LargeSmt<S>>,
108 blockchain: Blockchain,
109 account_tree: AccountTreeWithHistory<S>,
110}
111
112impl<S: SmtStorage> InnerState<S> {
113 fn latest_block_num(&self) -> BlockNumber {
115 self.blockchain
116 .chain_tip()
117 .expect("chain should always have at least the genesis block")
118 }
119}
120
121pub struct State {
126 data_directory: PathBuf,
128
129 db: Arc<Db>,
132
133 block_store: Arc<BlockStore>,
135
136 inner: RwLock<InnerState<TreeStorage>>,
140
141 forest: RwLock<AccountStateForest<AccountStateForestBackend>>,
143
144 writer: Mutex<()>,
147
148 proven_tip: ProvenTipWriter,
150
151 committed_tip_tx: watch::Sender<BlockNumber>,
154
155 pub(crate) block_cache: BlockCache,
158
159 pub(crate) proof_cache: ProofCache,
162}
163
164impl State {
165 #[miden_instrument(
173 target = COMPONENT,
174 skip_all,
175 )]
176 pub async fn load(
177 data_path: &Path,
178 storage_options: StorageOptions,
179 ) -> Result<Self, StateInitializationError> {
180 Self::load_with_database_options(data_path, storage_options, DatabaseOptions::default())
181 .await
182 }
183
184 #[miden_instrument(
189 target = COMPONENT,
190 skip_all,
191 )]
192 pub async fn load_with_database_options(
193 data_path: &Path,
194 storage_options: StorageOptions,
195 database_options: DatabaseOptions,
196 ) -> Result<Self, StateInitializationError> {
197 let data_directory = DataDirectory::load(data_path.to_path_buf())
198 .map_err(StateInitializationError::DataDirectoryLoadError)?;
199
200 let block_store = Arc::new(
201 BlockStore::load(data_directory.block_store_dir())
202 .map_err(StateInitializationError::BlockStoreLoadError)?,
203 );
204
205 let database_filepath = data_directory.database_path();
206 let mut db = Db::load_with_pool_size(
207 database_filepath.clone(),
208 database_options.connection_pool_size,
209 )
210 .await
211 .map_err(StateInitializationError::DatabaseLoadError)?;
212
213 let blockchain = load_mmr(&mut db).await?;
214 let latest_block_num = blockchain.chain_tip().unwrap_or(BlockNumber::GENESIS);
215
216 #[cfg(feature = "rocksdb")]
217 let (account_storage_config, nullifier_storage_config, forest_storage_config) = (
218 storage_options.account_tree.into(),
219 storage_options.nullifier_tree.into(),
220 storage_options.account_state_forest.into(),
221 );
222 #[cfg(not(feature = "rocksdb"))]
223 let (account_storage_config, nullifier_storage_config, forest_storage_config) = {
224 let _ = &storage_options;
225 ((), (), ())
226 };
227 let account_storage =
228 TreeStorage::create(data_path, &account_storage_config, ACCOUNT_TREE_STORAGE_DIR)?;
229 let account_tree = account_storage.load_account_tree(&mut db).await?;
230
231 let nullifier_storage =
232 TreeStorage::create(data_path, &nullifier_storage_config, NULLIFIER_TREE_STORAGE_DIR)?;
233 let nullifier_tree = nullifier_storage.load_nullifier_tree(&mut db).await?;
234
235 verify_tree_consistency(account_tree.root(), nullifier_tree.root(), &mut db).await?;
239
240 let account_tree = AccountTreeWithHistory::new(account_tree, latest_block_num);
241
242 let forest_backend = AccountStateForestBackend::create(
243 data_path,
244 &forest_storage_config,
245 ACCOUNT_STATE_FOREST_STORAGE_DIR,
246 )?;
247 let forest = forest_backend.load_account_state_forest(&mut db, latest_block_num).await?;
248 verify_account_state_forest_consistency(&forest, &mut db).await?;
249
250 let inner = RwLock::new(InnerState { nullifier_tree, blockchain, account_tree });
251
252 let forest = RwLock::new(forest);
253 let writer = Mutex::new(());
254 let db = Arc::new(db);
255
256 let proven_tip_init = block_store
258 .load_proven_tip()
259 .map_err(StateInitializationError::ProvenTipLoadError)?;
260 let (proven_tip, _rx) = ProvenTipWriter::new(proven_tip_init);
261
262 let (committed_tip_tx, _rx) = watch::channel(latest_block_num);
264
265 Ok(Self {
266 data_directory: data_path.to_path_buf(),
267 db,
268 block_store,
269 inner,
270 forest,
271 writer,
272 proven_tip,
273 committed_tip_tx,
274 block_cache: BlockCache::new(BLOCK_CACHE_CAPACITY),
275 proof_cache: ProofCache::new(PROOF_CACHE_CAPACITY),
276 })
277 }
278
279 pub fn subscribe_committed_tip(&self) -> watch::Receiver<BlockNumber> {
281 self.committed_tip_tx.subscribe()
282 }
283
284 pub async fn load_proving_inputs(
286 &self,
287 block_num: BlockNumber,
288 ) -> std::io::Result<Option<Vec<u8>>> {
289 self.block_store.load_proving_inputs(block_num).await
290 }
291
292 pub fn subscribe_proven_tip(&self) -> watch::Receiver<BlockNumber> {
294 self.proven_tip.subscribe()
295 }
296
297 fn with_inner_read_blocking<R>(&self, f: impl FnOnce(&InnerState<TreeStorage>) -> R) -> R {
306 let span = Span::current();
307 tokio::task::block_in_place(|| {
308 span.in_scope(|| {
309 let inner = self.inner.blocking_read();
310 f(&inner)
311 })
312 })
313 }
314
315 fn with_inner_and_forest_write_blocking<R>(
321 &self,
322 f: impl FnOnce(
323 &mut InnerState<TreeStorage>,
324 &mut AccountStateForest<AccountStateForestBackend>,
325 ) -> R,
326 ) -> R {
327 let span = Span::current();
328 tokio::task::block_in_place(|| {
329 span.in_scope(|| {
330 let mut inner = self.inner.blocking_write();
331 let mut forest = self.forest.blocking_write();
332 f(&mut inner, &mut forest)
333 })
334 })
335 }
336
337 fn with_forest_read_blocking<R>(
343 &self,
344 f: impl FnOnce(&AccountStateForest<AccountStateForestBackend>) -> R,
345 ) -> R {
346 let span = Span::current();
347 tokio::task::block_in_place(|| {
348 span.in_scope(|| {
349 let forest = self.forest.blocking_read();
350 f(&forest)
351 })
352 })
353 }
354
355 #[miden_instrument(
363 level = "debug",
364 target = COMPONENT,
365 skip_all,
366 err,
367 )]
368 pub async fn get_block_header(
369 &self,
370 block_num: Option<BlockNumber>,
371 include_mmr_proof: bool,
372 ) -> Result<(Option<BlockHeader>, Option<MmrProof>), GetBlockHeaderError> {
373 let block_header = self.db.select_block_header_by_block_num(block_num).await?;
374 if let Some(header) = block_header {
375 let mmr_proof = if include_mmr_proof {
376 let inner = self.inner.read().await;
377 let mmr_proof = inner.blockchain.open(header.block_num())?;
378 Some(mmr_proof)
379 } else {
380 None
381 };
382 Ok((Some(header), mmr_proof))
383 } else {
384 Ok((None, None))
385 }
386 }
387
388 pub async fn get_notes_by_id(
393 &self,
394 note_ids: Vec<NoteId>,
395 ) -> Result<Vec<NoteRecord>, DatabaseError> {
396 self.db.select_notes_by_id(note_ids).await
397 }
398
399 pub async fn get_batch_inputs(
418 &self,
419 tx_reference_blocks: BTreeSet<BlockNumber>,
420 unauthenticated_note_commitments: BTreeSet<Word>,
421 ) -> Result<BatchInputs, GetBatchInputsError> {
422 if tx_reference_blocks.is_empty() {
423 return Err(GetBatchInputsError::TransactionBlockReferencesEmpty);
424 }
425
426 let note_proofs = self
430 .db
431 .select_note_inclusion_proofs(unauthenticated_note_commitments)
432 .await
433 .map_err(GetBatchInputsError::SelectNoteInclusionProofError)?;
434
435 let note_blocks = note_proofs.values().map(|proof| proof.location().block_num());
437
438 let mut blocks: BTreeSet<BlockNumber> = tx_reference_blocks;
442 blocks.extend(note_blocks);
443
444 let (batch_reference_block, partial_mmr) = {
447 let inner_state = self.inner.read().await;
448
449 let latest_block_num = inner_state.latest_block_num();
450
451 let highest_block_num =
452 *blocks.last().expect("we should have checked for empty block references");
453 if highest_block_num > latest_block_num {
454 return Err(GetBatchInputsError::UnknownTransactionBlockReference {
455 highest_block_num,
456 latest_block_num,
457 });
458 }
459
460 blocks.remove(&latest_block_num);
464
465 let partial_mmr = inner_state
473 .blockchain
474 .partial_mmr_from_blocks(&blocks, latest_block_num)
475 .expect("latest block num should exist and all blocks in set should be < than latest block");
476
477 (latest_block_num, partial_mmr)
478 };
479
480 let mut headers = self
483 .db
484 .select_block_headers(blocks.into_iter().chain(std::iter::once(batch_reference_block)))
485 .await
486 .map_err(GetBatchInputsError::SelectBlockHeaderError)?;
487
488 let header_index = headers
490 .iter()
491 .enumerate()
492 .find_map(|(index, header)| {
493 (header.block_num() == batch_reference_block).then_some(index)
494 })
495 .expect("DB should have returned the header of the batch reference block");
496
497 let batch_reference_block_header = headers.swap_remove(header_index);
499
500 let partial_block_chain = PartialBlockchain::new_unchecked(partial_mmr, headers)
509 .expect("partial mmr and block headers should be consistent");
510
511 Ok(BatchInputs {
512 batch_reference_block_header,
513 note_proofs,
514 partial_block_chain,
515 })
516 }
517
518 pub async fn get_block_inputs(
520 &self,
521 account_ids: Vec<AccountId>,
522 nullifiers: Vec<Nullifier>,
523 unauthenticated_note_commitments: BTreeSet<Word>,
524 reference_blocks: BTreeSet<BlockNumber>,
525 ) -> Result<BlockInputs, GetBlockInputsError> {
526 let unauthenticated_note_proofs = self
530 .db
531 .select_note_inclusion_proofs(unauthenticated_note_commitments)
532 .await
533 .map_err(GetBlockInputsError::SelectNoteInclusionProofError)?;
534
535 let note_proof_reference_blocks =
537 unauthenticated_note_proofs.values().map(|proof| proof.location().block_num());
538
539 let mut blocks = reference_blocks;
541 blocks.extend(note_proof_reference_blocks);
542
543 let (latest_block_number, account_witnesses, nullifier_witnesses, partial_mmr) =
544 self.get_block_inputs_witnesses(&mut blocks, &account_ids, &nullifiers)?;
545
546 let mut headers = self
549 .db
550 .select_block_headers(blocks.into_iter().chain(std::iter::once(latest_block_number)))
551 .await
552 .map_err(GetBlockInputsError::SelectBlockHeaderError)?;
553
554 let latest_block_header_index = headers
557 .iter()
558 .enumerate()
559 .find_map(|(index, header)| {
560 (header.block_num() == latest_block_number).then_some(index)
561 })
562 .expect("DB should have returned the header of the latest block header");
563
564 let latest_block_header = headers.swap_remove(latest_block_header_index);
566
567 let partial_block_chain = PartialBlockchain::new_unchecked(partial_mmr, headers)
576 .expect("partial mmr and block headers should be consistent");
577
578 Ok(BlockInputs::new(
579 latest_block_header,
580 partial_block_chain,
581 account_witnesses,
582 nullifier_witnesses,
583 unauthenticated_note_proofs,
584 ))
585 }
586
587 fn get_block_inputs_witnesses(
594 &self,
595 blocks: &mut BTreeSet<BlockNumber>,
596 account_ids: &[AccountId],
597 nullifiers: &[Nullifier],
598 ) -> Result<BlockInputWitnesses, GetBlockInputsError> {
599 self.with_inner_read_blocking(|inner| {
600 let latest_block_number = inner.latest_block_num();
601
602 let highest_block_number = blocks.last().copied().unwrap_or(latest_block_number);
604 if highest_block_number > latest_block_number {
605 return Err(GetBlockInputsError::UnknownBatchBlockReference {
606 highest_block_number,
607 latest_block_number,
608 });
609 }
610
611 blocks.remove(&latest_block_number);
614
615 let partial_mmr =
626 inner.blockchain.partial_mmr_from_blocks(blocks, latest_block_number).expect(
627 "latest block num should exist and all blocks in set should be < than latest block",
628 );
629
630 let account_witnesses = account_ids
632 .iter()
633 .copied()
634 .map(|account_id| (account_id, inner.account_tree.open_latest(account_id)))
635 .collect::<BTreeMap<AccountId, AccountWitness>>();
636
637 let nullifier_witnesses: BTreeMap<Nullifier, NullifierWitness> = nullifiers
640 .iter()
641 .copied()
642 .map(|nullifier| (nullifier, inner.nullifier_tree.open(&nullifier)))
643 .collect();
644
645 Ok((latest_block_number, account_witnesses, nullifier_witnesses, partial_mmr))
646 })
647 }
648
649 #[miden_instrument(
651 target = COMPONENT,
652 skip_all,
653 fields(
654 account.id=%account_id,
655 nullifiers = %format_array(nullifiers),
656 ),
657 )]
658 pub async fn get_transaction_inputs(
659 &self,
660 account_id: AccountId,
661 nullifiers: &[Nullifier],
662 unauthenticated_note_commitments: Vec<Word>,
663 ) -> Result<TransactionInputs, DatabaseError> {
664 let tree_inputs = self.with_inner_read_blocking(|inner| {
665 let account_commitment = inner.account_tree.get_latest_commitment(account_id);
666
667 let new_account_id_prefix_is_unique = if account_commitment.is_empty() {
668 Some(!inner.account_tree.contains_account_id_prefix_in_latest(account_id.prefix()))
669 } else {
670 None
671 };
672
673 if let Some(false) = new_account_id_prefix_is_unique {
675 return Err(TransactionInputs {
676 new_account_id_prefix_is_unique,
677 ..Default::default()
678 });
679 }
680
681 let nullifiers = nullifiers
682 .iter()
683 .map(|nullifier| NullifierInfo {
684 nullifier: *nullifier,
685 block_num: inner.nullifier_tree.get_block_num(nullifier).unwrap_or_default(),
686 })
687 .collect();
688
689 Ok((account_commitment, nullifiers, new_account_id_prefix_is_unique))
690 });
691 let (account_commitment, nullifiers, new_account_id_prefix_is_unique) = match tree_inputs {
692 Ok(inputs) => inputs,
693 Err(inputs) => return Ok(inputs),
694 };
695
696 let found_unauthenticated_notes = self
697 .db
698 .select_existing_note_commitments(unauthenticated_note_commitments)
699 .await?;
700
701 Ok(TransactionInputs {
702 account_commitment,
703 nullifiers,
704 found_unauthenticated_notes,
705 new_account_id_prefix_is_unique,
706 })
707 }
708
709 pub async fn filter_network_accounts(
711 &self,
712 account_ids: &[AccountId],
713 ) -> Result<HashSet<AccountId>, DatabaseError> {
714 self.db.select_network_accounts_subset(account_ids.to_vec()).await
715 }
716
717 pub async fn chain_tip(&self, finality: Finality) -> BlockNumber {
723 match finality {
724 Finality::Committed => self
725 .inner
726 .read()
727 .instrument(tracing::info_span!("acquire_inner"))
728 .await
729 .latest_block_num(),
730 Finality::Proven => self.proven_tip.read(),
731 }
732 }
733
734 pub async fn load_block(
737 &self,
738 block_num: BlockNumber,
739 ) -> Result<Option<Vec<u8>>, DatabaseError> {
740 if block_num > self.chain_tip(Finality::Committed).await {
741 return Ok(None);
742 }
743 if let Some(block) = self.block_cache.get(block_num) {
744 return Ok(Some(block.block_bytes().to_vec()));
745 }
746 self.block_store.load_block(block_num).await.map_err(Into::into)
747 }
748
749 pub async fn load_proof(
752 &self,
753 block_num: BlockNumber,
754 ) -> Result<Option<Vec<u8>>, DatabaseError> {
755 if block_num > self.chain_tip(Finality::Proven).await {
756 return Ok(None);
757 }
758 if let Some(proof) = self.proof_cache.get(block_num) {
759 return Ok(Some(proof.proof_bytes().to_vec()));
760 }
761 self.block_store.load_proof(block_num).await.map_err(Into::into)
762 }
763
764 pub async fn get_note_script_by_root(
766 &self,
767 root: Word,
768 ) -> Result<Option<NoteScript>, DatabaseError> {
769 self.db.select_note_script_by_root(root).await
770 }
771}