1use alloc::{
2 boxed::Box,
3 collections::{BTreeMap, BTreeSet},
4 vec::Vec,
5};
6
7use anyhow::Context;
8use miden_block_prover::{LocalBlockProver, ProvenBlockError};
9use miden_lib::{
10 account::{faucets::BasicFungibleFaucet, wallets::BasicWallet},
11 note::{create_p2id_note, create_p2idr_note},
12 transaction::{TransactionKernel, memory},
13};
14use miden_objects::{
15 MAX_BATCHES_PER_BLOCK, MAX_OUTPUT_NOTES_PER_BATCH, NoteError, ProposedBatchError,
16 ProposedBlockError,
17 account::{
18 Account, AccountBuilder, AccountId, AccountIdAnchor, AccountStorageMode, AccountType,
19 StorageSlot, delta::AccountUpdateDetails,
20 },
21 asset::{Asset, TokenSymbol},
22 batch::{ProposedBatch, ProvenBatch},
23 block::{
24 AccountTree, AccountWitness, BlockAccountUpdate, BlockHeader, BlockInputs, BlockNoteTree,
25 BlockNumber, Blockchain, NullifierTree, NullifierWitness, ProposedBlock, ProvenBlock,
26 },
27 crypto::merkle::SmtProof,
28 note::{Note, NoteHeader, NoteId, NoteInclusionProof, NoteType, Nullifier},
29 testing::account_code::DEFAULT_AUTH_SCRIPT,
30 transaction::{
31 AccountInputs, ExecutedTransaction, InputNote, InputNotes, OrderedTransactionHeaders,
32 OutputNote, PartialBlockchain, ProvenTransaction, TransactionHeader, TransactionInputs,
33 TransactionScript,
34 },
35};
36use rand::{Rng, SeedableRng};
37use rand_chacha::ChaCha20Rng;
38use vm_processor::{Digest, Felt, Word, ZERO, crypto::RpoRandomCoin};
39
40use super::note::MockChainNote;
41use crate::{
42 Auth, MockFungibleFaucet, ProvenTransactionExt, TransactionContextBuilder,
43 mock_chain::account::MockAccount,
44};
45
46#[derive(Debug, Clone)]
155pub struct MockChain {
156 chain: Blockchain,
158
159 blocks: Vec<ProvenBlock>,
161
162 nullifier_tree: NullifierTree,
164
165 account_tree: AccountTree,
167
168 pending_objects: PendingObjects,
175
176 pending_transactions: Vec<ProvenTransaction>,
179
180 committed_notes: BTreeMap<NoteId, MockChainNote>,
182
183 committed_accounts: BTreeMap<AccountId, MockAccount>,
185
186 rng: ChaCha20Rng,
188}
189
190impl MockChain {
191 pub const TIMESTAMP_START_SECS: u32 = 1700000000;
196
197 pub const TIMESTAMP_STEP_SECS: u32 = 10;
200
201 pub fn new() -> Self {
206 Self::with_accounts(&[])
207 }
208
209 pub fn with_accounts(accounts: &[Account]) -> Self {
211 let (genesis_block, account_tree) =
212 create_genesis_state(accounts.iter().cloned()).expect("TODO: turn into error");
213
214 let mut chain = MockChain {
215 chain: Blockchain::default(),
216 blocks: vec![],
217 nullifier_tree: NullifierTree::default(),
218 account_tree,
219 pending_objects: PendingObjects::new(),
220 pending_transactions: Vec::new(),
221 committed_notes: BTreeMap::new(),
222 committed_accounts: BTreeMap::new(),
223 rng: ChaCha20Rng::from_seed(Default::default()),
225 };
226
227 chain
230 .apply_block(genesis_block)
231 .context("failed to apply genesis block")
232 .unwrap();
233
234 debug_assert_eq!(chain.blocks.len(), 1);
235 debug_assert_eq!(chain.account_tree.num_accounts(), accounts.len());
236 debug_assert_eq!(chain.committed_accounts.len(), accounts.len());
237 for added_account in accounts {
238 debug_assert_eq!(
239 chain.account_tree.get(added_account.id()),
240 added_account.commitment()
241 );
242 debug_assert_eq!(
243 chain.committed_account(added_account.id()).commitment(),
244 added_account.commitment(),
245 );
246 }
247
248 chain
249 }
250
251 pub fn blockchain(&self) -> &Blockchain {
256 &self.chain
257 }
258
259 pub fn latest_partial_blockchain(&self) -> PartialBlockchain {
262 let block_headers =
265 self.blocks.iter().map(|b| b.header()).take(self.blocks.len() - 1).cloned();
266
267 PartialBlockchain::from_blockchain(&self.chain, block_headers).unwrap()
268 }
269
270 pub fn latest_selective_partial_blockchain(
276 &self,
277 reference_blocks: impl IntoIterator<Item = BlockNumber>,
278 ) -> (BlockHeader, PartialBlockchain) {
279 let latest_block_header = self.latest_block_header().clone();
280 let reference_blocks: BTreeSet<_> = reference_blocks.into_iter().collect();
283
284 let block_headers: Vec<_> = reference_blocks
286 .into_iter()
287 .map(|block_ref_num| self.block_header(block_ref_num.as_usize()))
288 .filter(|block_header| block_header.commitment() != latest_block_header.commitment())
289 .collect();
290
291 let partial_blockchain =
292 PartialBlockchain::from_blockchain(&self.chain, block_headers).unwrap();
293
294 (latest_block_header, partial_blockchain)
295 }
296
297 pub fn account_witnesses(
300 &self,
301 account_ids: impl IntoIterator<Item = AccountId>,
302 ) -> BTreeMap<AccountId, AccountWitness> {
303 let mut account_witnesses = BTreeMap::new();
304
305 for account_id in account_ids {
306 let witness = self.account_tree.open(account_id);
307 account_witnesses.insert(account_id, witness);
308 }
309
310 account_witnesses
311 }
312
313 pub fn nullifier_witnesses(
316 &self,
317 nullifiers: impl IntoIterator<Item = Nullifier>,
318 ) -> BTreeMap<Nullifier, NullifierWitness> {
319 let mut nullifier_proofs = BTreeMap::new();
320
321 for nullifier in nullifiers {
322 let witness = self.nullifier_tree.open(&nullifier);
323 nullifier_proofs.insert(nullifier, witness);
324 }
325
326 nullifier_proofs
327 }
328
329 pub fn unauthenticated_note_proofs(
333 &self,
334 notes: impl IntoIterator<Item = NoteId>,
335 ) -> BTreeMap<NoteId, NoteInclusionProof> {
336 let mut proofs = BTreeMap::default();
337 for note in notes {
338 if let Some(input_note) = self.committed_notes.get(¬e) {
339 proofs.insert(note, input_note.inclusion_proof().clone());
340 }
341 }
342
343 proofs
344 }
345
346 pub fn latest_block_header(&self) -> BlockHeader {
348 let chain_tip =
349 self.chain.chain_tip().expect("chain should contain at least the genesis block");
350 self.blocks[chain_tip.as_usize()].header().clone()
351 }
352
353 pub fn block_header(&self, block_number: usize) -> BlockHeader {
355 self.blocks[block_number].header().clone()
356 }
357
358 pub fn proven_blocks(&self) -> &[ProvenBlock] {
360 &self.blocks
361 }
362
363 pub fn nullifier_tree(&self) -> &NullifierTree {
365 &self.nullifier_tree
366 }
367
368 pub fn committed_notes(&self) -> &BTreeMap<NoteId, MockChainNote> {
372 &self.committed_notes
373 }
374
375 pub fn get_public_note(&self, note_id: &NoteId) -> Option<InputNote> {
378 let note = self.committed_notes.get(note_id)?;
379 note.clone().try_into().ok()
380 }
381
382 pub fn committed_account(&self, account_id: AccountId) -> &Account {
385 self.committed_accounts
386 .get(&account_id)
387 .expect("account should be available")
388 .account()
389 }
390
391 pub fn account_tree(&self) -> &AccountTree {
393 &self.account_tree
394 }
395
396 pub fn propose_transaction_batch<I>(
403 &self,
404 txs: impl IntoIterator<Item = ProvenTransaction, IntoIter = I>,
405 ) -> Result<ProposedBatch, ProposedBatchError>
406 where
407 I: Iterator<Item = ProvenTransaction> + Clone,
408 {
409 let transactions: Vec<_> = txs.into_iter().map(alloc::sync::Arc::new).collect();
410
411 let (batch_reference_block, partial_blockchain, unauthenticated_note_proofs) = self
412 .get_batch_inputs(
413 transactions.iter().map(|tx| tx.ref_block_num()),
414 transactions
415 .iter()
416 .flat_map(|tx| tx.unauthenticated_notes().map(NoteHeader::id)),
417 );
418
419 ProposedBatch::new(
420 transactions,
421 batch_reference_block,
422 partial_blockchain,
423 unauthenticated_note_proofs,
424 )
425 }
426
427 pub fn prove_transaction_batch(&self, proposed_batch: ProposedBatch) -> ProvenBatch {
431 let (
432 transactions,
433 block_header,
434 _partial_blockchain,
435 _unauthenticated_note_proofs,
436 id,
437 account_updates,
438 input_notes,
439 output_notes,
440 batch_expiration_block_num,
441 ) = proposed_batch.into_parts();
442
443 let tx_headers = OrderedTransactionHeaders::new_unchecked(
445 transactions
446 .iter()
447 .map(AsRef::as_ref)
448 .map(TransactionHeader::from)
449 .collect::<Vec<_>>(),
450 );
451
452 ProvenBatch::new(
453 id,
454 block_header.commitment(),
455 block_header.block_num(),
456 account_updates,
457 input_notes,
458 output_notes,
459 batch_expiration_block_num,
460 tx_headers,
461 )
462 .expect("failed to create ProvenBatch")
463 }
464
465 pub fn propose_block_at<I>(
472 &self,
473 batches: impl IntoIterator<Item = ProvenBatch, IntoIter = I>,
474 timestamp: u32,
475 ) -> Result<ProposedBlock, ProposedBlockError>
476 where
477 I: Iterator<Item = ProvenBatch> + Clone,
478 {
479 let batches: Vec<_> = batches.into_iter().collect();
480 let block_inputs = self.get_block_inputs(batches.iter());
481
482 let proposed_block = ProposedBlock::new_at(block_inputs, batches, timestamp)?;
483
484 Ok(proposed_block)
485 }
486
487 pub fn propose_block<I>(
491 &self,
492 batches: impl IntoIterator<Item = ProvenBatch, IntoIter = I>,
493 ) -> Result<ProposedBlock, ProposedBlockError>
494 where
495 I: Iterator<Item = ProvenBatch> + Clone,
496 {
497 let timestamp = self.latest_block_header().timestamp() + 1;
500
501 self.propose_block_at(batches, timestamp)
502 }
503
504 pub fn prove_block(
508 &self,
509 proposed_block: ProposedBlock,
510 ) -> Result<ProvenBlock, ProvenBlockError> {
511 LocalBlockProver::new(0).prove_without_batch_verification(proposed_block)
512 }
513
514 pub fn build_tx_context(
538 &self,
539 input: impl Into<TxContextInput>,
540 note_ids: &[NoteId],
541 unauthenticated_notes: &[Note],
542 ) -> TransactionContextBuilder {
543 let mock_account = match input.into() {
544 TxContextInput::AccountId(account_id) => {
545 self.committed_accounts.get(&account_id).unwrap().clone()
546 },
547 TxContextInput::Account(account) => {
548 let committed_account = self.committed_accounts.get(&account.id());
549 let authenticator = committed_account.and_then(|account| account.authenticator());
550 let seed = committed_account.and_then(|account| account.seed());
551 MockAccount::new(account, seed.cloned(), authenticator.cloned())
552 },
553 TxContextInput::ExecutedTransaction(executed_transaction) => {
554 let mut initial_account = executed_transaction.initial_account().clone();
555 initial_account
556 .apply_delta(executed_transaction.account_delta())
557 .expect("delta from tx should be valid for initial account from tx");
558 let committed_account = self.committed_accounts.get(&initial_account.id());
559 let authenticator = committed_account.and_then(|account| account.authenticator());
560 let seed = committed_account.and_then(|account| account.seed());
561 MockAccount::new(initial_account, seed.cloned(), authenticator.cloned())
562 },
563 };
564
565 let tx_inputs = self.get_transaction_inputs(
566 mock_account.account().clone(),
567 mock_account.seed().cloned(),
568 note_ids,
569 unauthenticated_notes,
570 );
571
572 let mut tx_context_builder = TransactionContextBuilder::new(mock_account.account().clone())
573 .authenticator(mock_account.authenticator().cloned())
574 .account_seed(mock_account.seed().cloned())
575 .tx_inputs(tx_inputs);
576
577 if mock_account.authenticator().is_some() {
578 let tx_script = TransactionScript::compile(
579 DEFAULT_AUTH_SCRIPT,
580 vec![],
581 TransactionKernel::testing_assembler_with_mock_account(),
582 )
583 .unwrap();
584 tx_context_builder = tx_context_builder.tx_script(tx_script);
585 }
586
587 tx_context_builder
588 }
589
590 pub fn get_transaction_inputs(
595 &self,
596 account: Account,
597 account_seed: Option<Word>,
598 notes: &[NoteId],
599 unauthenticated_notes: &[Note],
600 ) -> TransactionInputs {
601 let block = self.blocks.last().unwrap();
602
603 let mut input_notes = vec![];
604 let mut block_headers_map: BTreeMap<BlockNumber, BlockHeader> = BTreeMap::new();
605 for note in notes {
606 let input_note: InputNote = self
607 .committed_notes
608 .get(note)
609 .expect("Note not found")
610 .clone()
611 .try_into()
612 .expect("Note should be public");
613 let note_block_num = input_note.location().unwrap().block_num();
614 if note_block_num != block.header().block_num() {
615 block_headers_map.insert(
616 note_block_num,
617 self.blocks.get(note_block_num.as_usize()).unwrap().header().clone(),
618 );
619 }
620
621 input_notes.push(input_note);
622 }
623
624 if account.is_new() {
627 let epoch_block_num = BlockNumber::from_epoch(account.id().anchor_epoch());
628 if epoch_block_num != block.header().block_num() {
631 block_headers_map.insert(
632 epoch_block_num,
633 self.blocks.get(epoch_block_num.as_usize()).unwrap().header().clone(),
634 );
635 }
636 }
637
638 for note in unauthenticated_notes {
639 input_notes.push(InputNote::Unauthenticated { note: note.clone() })
640 }
641
642 let block_headers = block_headers_map.values().cloned();
643 let mmr = PartialBlockchain::from_blockchain(&self.chain, block_headers).unwrap();
644
645 TransactionInputs::new(
646 account,
647 account_seed,
648 block.header().clone(),
649 mmr,
650 InputNotes::new(input_notes).unwrap(),
651 )
652 .unwrap()
653 }
654
655 pub fn get_batch_inputs(
658 &self,
659 tx_reference_blocks: impl IntoIterator<Item = BlockNumber>,
660 unauthenticated_notes: impl Iterator<Item = NoteId>,
661 ) -> (BlockHeader, PartialBlockchain, BTreeMap<NoteId, NoteInclusionProof>) {
662 let unauthenticated_note_proofs = self.unauthenticated_note_proofs(unauthenticated_notes);
664
665 let required_blocks = tx_reference_blocks.into_iter().chain(
668 unauthenticated_note_proofs
669 .values()
670 .map(|note_proof| note_proof.location().block_num()),
671 );
672
673 let (batch_reference_block, partial_block_chain) =
674 self.latest_selective_partial_blockchain(required_blocks);
675
676 (batch_reference_block, partial_block_chain, unauthenticated_note_proofs)
677 }
678
679 pub fn get_foreign_account_inputs(&self, account_id: AccountId) -> AccountInputs {
681 let account = self.committed_account(account_id);
682
683 let account_witness = self.account_tree().open(account_id);
684 assert_eq!(account_witness.state_commitment(), account.commitment());
685
686 let mut storage_map_proofs = vec![];
687 for slot in account.storage().slots() {
688 if let StorageSlot::Map(map) = slot {
690 let proofs: Vec<SmtProof> = map.entries().map(|(key, _)| map.open(key)).collect();
691 storage_map_proofs.extend(proofs);
692 }
693 }
694
695 AccountInputs::new(account.into(), account_witness)
696 }
697
698 pub fn get_block_inputs<'batch, I>(
700 &self,
701 batch_iter: impl IntoIterator<Item = &'batch ProvenBatch, IntoIter = I>,
702 ) -> BlockInputs
703 where
704 I: Iterator<Item = &'batch ProvenBatch> + Clone,
705 {
706 let batch_iterator = batch_iter.into_iter();
707
708 let unauthenticated_note_proofs =
709 self.unauthenticated_note_proofs(batch_iterator.clone().flat_map(|batch| {
710 batch.input_notes().iter().filter_map(|note| note.header().map(NoteHeader::id))
711 }));
712
713 let (block_reference_block, partial_blockchain) = self.latest_selective_partial_blockchain(
714 batch_iterator.clone().map(ProvenBatch::reference_block_num).chain(
715 unauthenticated_note_proofs.values().map(|proof| proof.location().block_num()),
716 ),
717 );
718
719 let account_witnesses =
720 self.account_witnesses(batch_iterator.clone().flat_map(ProvenBatch::updated_accounts));
721
722 let nullifier_proofs =
723 self.nullifier_witnesses(batch_iterator.flat_map(ProvenBatch::created_nullifiers));
724
725 BlockInputs::new(
726 block_reference_block,
727 partial_blockchain,
728 account_witnesses,
729 nullifier_proofs,
730 unauthenticated_note_proofs,
731 )
732 }
733
734 pub fn prove_next_block(&mut self) -> ProvenBlock {
741 self.prove_block_inner(None).unwrap()
742 }
743
744 pub fn prove_next_block_at(&mut self, timestamp: u32) -> anyhow::Result<ProvenBlock> {
746 self.prove_block_inner(Some(timestamp))
747 }
748
749 pub fn prove_until_block(
759 &mut self,
760 target_block_num: impl Into<BlockNumber>,
761 ) -> anyhow::Result<ProvenBlock> {
762 let target_block_num = target_block_num.into();
763 let latest_block_num = self.latest_block_header().block_num();
764 assert!(
765 target_block_num > latest_block_num,
766 "target block number must be greater than the number of the latest block in the chain"
767 );
768
769 let mut last_block = None;
770 for _ in latest_block_num.as_usize()..target_block_num.as_usize() {
771 last_block = Some(self.prove_next_block());
772 }
773
774 Ok(last_block.expect("at least one block should have been created"))
775 }
776
777 pub fn set_rng_seed(&mut self, seed: [u8; 32]) {
779 self.rng = ChaCha20Rng::from_seed(seed);
780 }
781
782 pub fn add_pending_executed_transaction(
792 &mut self,
793 transaction: &ExecutedTransaction,
794 ) -> Account {
795 let mut account = transaction.initial_account().clone();
796 account.apply_delta(transaction.account_delta()).unwrap();
797
798 let proven_tx = ProvenTransaction::from_executed_transaction_mocked(transaction.clone());
800
801 self.pending_transactions.push(proven_tx);
802
803 account
804 }
805
806 pub fn add_pending_proven_transaction(&mut self, transaction: ProvenTransaction) {
811 self.pending_transactions.push(transaction);
812 }
813
814 pub fn add_pending_note(&mut self, note: OutputNote) {
819 self.pending_objects.output_notes.push(note);
820 }
821
822 pub fn add_pending_p2id_note(
827 &mut self,
828 sender_account_id: AccountId,
829 target_account_id: AccountId,
830 asset: &[Asset],
831 note_type: NoteType,
832 reclaim_height: Option<BlockNumber>,
833 ) -> Result<Note, NoteError> {
834 let mut rng = RpoRandomCoin::new(Word::default());
835
836 let note = if let Some(height) = reclaim_height {
837 create_p2idr_note(
838 sender_account_id,
839 target_account_id,
840 asset.to_vec(),
841 note_type,
842 Default::default(),
843 height,
844 &mut rng,
845 )?
846 } else {
847 create_p2id_note(
848 sender_account_id,
849 target_account_id,
850 asset.to_vec(),
851 note_type,
852 Default::default(),
853 &mut rng,
854 )?
855 };
856
857 self.add_pending_note(OutputNote::Full(note.clone()));
858
859 Ok(note)
860 }
861
862 pub fn add_pending_nullifier(&mut self, nullifier: Nullifier) {
867 self.pending_objects.created_nullifiers.push(nullifier);
868 }
869
870 pub fn add_pending_new_wallet(&mut self, auth_method: Auth) -> Account {
875 let account_builder = AccountBuilder::new(self.rng.random())
876 .storage_mode(AccountStorageMode::Public)
877 .with_component(BasicWallet);
878
879 self.add_pending_account_from_builder(auth_method, account_builder, AccountState::New)
880 }
881
882 pub fn add_pending_existing_wallet(
887 &mut self,
888 auth_method: Auth,
889 assets: Vec<Asset>,
890 ) -> Account {
891 let account_builder = Account::builder(self.rng.random())
892 .storage_mode(AccountStorageMode::Public)
893 .with_component(BasicWallet)
894 .with_assets(assets);
895
896 self.add_pending_account_from_builder(auth_method, account_builder, AccountState::Exists)
897 }
898
899 pub fn add_pending_new_faucet(
905 &mut self,
906 auth_method: Auth,
907 token_symbol: &str,
908 max_supply: u64,
909 ) -> MockFungibleFaucet {
910 let account_builder = AccountBuilder::new(self.rng.random())
911 .storage_mode(AccountStorageMode::Public)
912 .account_type(AccountType::FungibleFaucet)
913 .with_component(
914 BasicFungibleFaucet::new(
915 TokenSymbol::new(token_symbol).unwrap(),
916 10,
917 max_supply.try_into().unwrap(),
918 )
919 .unwrap(),
920 );
921
922 MockFungibleFaucet::new(self.add_pending_account_from_builder(
923 auth_method,
924 account_builder,
925 AccountState::New,
926 ))
927 }
928
929 pub fn add_pending_existing_faucet(
935 &mut self,
936 auth_method: Auth,
937 token_symbol: &str,
938 max_supply: u64,
939 total_issuance: Option<u64>,
940 ) -> MockFungibleFaucet {
941 let mut account_builder = AccountBuilder::new(self.rng.random())
942 .storage_mode(AccountStorageMode::Public)
943 .with_component(
944 BasicFungibleFaucet::new(
945 TokenSymbol::new(token_symbol).unwrap(),
946 10u8,
947 Felt::new(max_supply),
948 )
949 .unwrap(),
950 )
951 .account_type(AccountType::FungibleFaucet);
952
953 let authenticator = match auth_method.build_component() {
954 Some((auth_component, authenticator)) => {
955 account_builder = account_builder.with_component(auth_component);
956 Some(authenticator)
957 },
958 None => None,
959 };
960 let mut account = account_builder.build_existing().unwrap();
961
962 if let Some(issuance) = total_issuance {
965 account
966 .storage_mut()
967 .set_item(memory::FAUCET_STORAGE_DATA_SLOT, [ZERO, ZERO, ZERO, Felt::new(issuance)])
968 .unwrap();
969 }
970
971 self.committed_accounts
974 .insert(account.id(), MockAccount::new(account.clone(), None, authenticator));
975 self.add_pending_account(account.clone());
976
977 MockFungibleFaucet::new(account)
978 }
979
980 pub fn add_pending_account_from_builder(
990 &mut self,
991 auth_method: Auth,
992 mut account_builder: AccountBuilder,
993 account_state: AccountState,
994 ) -> Account {
995 let authenticator = match auth_method.build_component() {
996 Some((auth_component, authenticator)) => {
997 account_builder = account_builder.with_component(auth_component);
998 Some(authenticator)
999 },
1000 None => None,
1001 };
1002
1003 let (account, seed) = if let AccountState::New = account_state {
1004 let epoch_block_number = self.latest_block_header().epoch_block_num();
1005 let account_id_anchor =
1006 self.blocks.get(epoch_block_number.as_usize()).unwrap().header();
1007 account_builder =
1008 account_builder.anchor(AccountIdAnchor::try_from(account_id_anchor).unwrap());
1009
1010 account_builder.build().map(|(account, seed)| (account, Some(seed))).unwrap()
1011 } else {
1012 account_builder.build_existing().map(|account| (account, None)).unwrap()
1013 };
1014
1015 self.committed_accounts
1021 .insert(account.id(), MockAccount::new(account.clone(), seed, authenticator));
1022
1023 if let AccountState::Exists = account_state {
1031 self.add_pending_account(account.clone());
1032 }
1033
1034 account
1035 }
1036
1037 pub fn add_pending_account(&mut self, account: Account) {
1041 self.pending_objects.updated_accounts.insert(
1042 account.id(),
1043 BlockAccountUpdate::new(
1044 account.id(),
1045 account.commitment(),
1046 AccountUpdateDetails::New(account),
1047 ),
1048 );
1049 }
1050
1051 fn apply_block_tree_updates(&mut self, proven_block: &ProvenBlock) -> anyhow::Result<()> {
1057 for account_update in proven_block.updated_accounts() {
1058 self.account_tree
1059 .insert(account_update.account_id(), account_update.final_state_commitment())
1060 .context("failed to insert account update into account tree")?;
1061 }
1062
1063 for nullifier in proven_block.created_nullifiers() {
1064 self.nullifier_tree
1065 .mark_spent(*nullifier, proven_block.header().block_num())
1066 .context("failed to mark block nullifier as spent")?;
1067
1068 }
1072
1073 Ok(())
1074 }
1075
1076 fn apply_block(&mut self, proven_block: ProvenBlock) -> anyhow::Result<()> {
1083 for account_update in proven_block.updated_accounts() {
1084 match account_update.details() {
1085 AccountUpdateDetails::New(account) => {
1086 let committed_account =
1087 self.committed_accounts.get(&account_update.account_id());
1088 let authenticator =
1089 committed_account.and_then(|account| account.authenticator());
1090 let seed = committed_account.and_then(|account| account.seed());
1091
1092 self.committed_accounts.insert(
1093 account.id(),
1094 MockAccount::new(account.clone(), seed.cloned(), authenticator.cloned()),
1095 );
1096 },
1097 AccountUpdateDetails::Delta(account_delta) => {
1098 let committed_account =
1099 self.committed_accounts.get_mut(&account_update.account_id()).ok_or_else(
1100 || anyhow::anyhow!("account delta in block for non-existent account"),
1101 )?;
1102 committed_account
1103 .apply_delta(account_delta)
1104 .context("failed to apply account delta to committed account")?;
1105 },
1106 AccountUpdateDetails::Private => {
1107 todo!("private accounts are not yet supported")
1108 },
1109 }
1110 }
1111
1112 let notes_tree = proven_block.build_output_note_tree();
1113 for (block_note_index, created_note) in proven_block.output_notes() {
1114 let note_path = notes_tree.get_note_path(block_note_index);
1115 let note_inclusion_proof = NoteInclusionProof::new(
1116 proven_block.header().block_num(),
1117 block_note_index.leaf_index_value(),
1118 note_path,
1119 )
1120 .context("failed to construct note inclusion proof")?;
1121
1122 if let OutputNote::Full(note) = created_note {
1123 self.committed_notes
1124 .insert(note.id(), MockChainNote::Public(note.clone(), note_inclusion_proof));
1125 } else {
1126 self.committed_notes.insert(
1127 created_note.id(),
1128 MockChainNote::Private(
1129 created_note.id(),
1130 *created_note.metadata(),
1131 note_inclusion_proof,
1132 ),
1133 );
1134 }
1135 }
1136
1137 debug_assert_eq!(
1138 self.chain.commitment(),
1139 proven_block.header().chain_commitment(),
1140 "current mock chain commitment and new block's chain commitment should match"
1141 );
1142 debug_assert_eq!(
1143 BlockNumber::from(self.chain.as_mmr().forest() as u32),
1144 proven_block.header().block_num(),
1145 "current mock chain length and new block's number should match"
1146 );
1147
1148 self.chain.push(proven_block.header().commitment());
1149 self.blocks.push(proven_block);
1150
1151 Ok(())
1152 }
1153
1154 fn pending_transactions_to_batches(&mut self) -> anyhow::Result<Vec<ProvenBatch>> {
1155 if self.pending_transactions.is_empty() {
1158 return Ok(vec![]);
1159 }
1160
1161 let pending_transactions = core::mem::take(&mut self.pending_transactions);
1162
1163 let proven_batch = self
1166 .propose_transaction_batch(pending_transactions)
1167 .map(|proposed_batch| self.prove_transaction_batch(proposed_batch))?;
1168
1169 Ok(vec![proven_batch])
1170 }
1171
1172 fn apply_pending_objects_to_block(
1173 &mut self,
1174 proven_block: &mut ProvenBlock,
1175 ) -> anyhow::Result<()> {
1176 let pending_account_updates = core::mem::take(&mut self.pending_objects.updated_accounts);
1178
1179 let updated_accounts_block: BTreeSet<AccountId> = proven_block
1180 .updated_accounts()
1181 .iter()
1182 .map(|update| update.account_id())
1183 .collect();
1184
1185 for (id, account_update) in pending_account_updates {
1186 if updated_accounts_block.contains(&id) {
1187 anyhow::bail!(
1188 "account {id} is already modified through a transaction in the block so it cannot also be modified through pending objects"
1189 );
1190 }
1191
1192 self.account_tree
1193 .insert(id, account_update.final_state_commitment())
1194 .context("failed to insert pending account into tree")?;
1195
1196 proven_block.updated_accounts_mut().push(account_update);
1197 }
1198
1199 let pending_created_nullifiers =
1201 core::mem::take(&mut self.pending_objects.created_nullifiers);
1202
1203 let created_nullifiers_block: BTreeSet<Nullifier> =
1204 proven_block.created_nullifiers().iter().copied().collect();
1205
1206 for nullifier in pending_created_nullifiers {
1207 if created_nullifiers_block.contains(&nullifier) {
1208 anyhow::bail!(
1209 "nullifier {nullifier} is already created by a transaction in the block so it cannot also be added through pending objects"
1210 );
1211 }
1212
1213 self.nullifier_tree
1214 .mark_spent(nullifier, proven_block.header().block_num())
1215 .context("failed to insert pending nullifier into tree")?;
1216
1217 proven_block.created_nullifiers_mut().push(nullifier);
1218 }
1219
1220 let output_notes_block: BTreeSet<NoteId> =
1222 proven_block.output_notes().map(|(_, output_note)| output_note.id()).collect();
1223
1224 if self.pending_objects.output_notes.len() > MAX_OUTPUT_NOTES_PER_BATCH {
1227 anyhow::bail!(
1228 "cannot create more than {MAX_OUTPUT_NOTES_PER_BATCH} notes through pending objects"
1229 );
1230 }
1231
1232 let mut pending_note_batch = Vec::with_capacity(self.pending_objects.output_notes.len());
1233 let pending_output_notes = core::mem::take(&mut self.pending_objects.output_notes);
1234 for (note_idx, output_note) in pending_output_notes.into_iter().enumerate() {
1235 if output_notes_block.contains(&output_note.id()) {
1236 anyhow::bail!(
1237 "output note {} is already created by a transaction in the block so it cannot also be created through pending objects",
1238 output_note.id()
1239 );
1240 }
1241
1242 pending_note_batch.push((note_idx, output_note));
1243 }
1244
1245 if (proven_block.output_note_batches().len() + 1) > MAX_BATCHES_PER_BLOCK {
1246 anyhow::bail!(
1247 "failed to add pending notes to block because max number of batches is already reached"
1248 )
1249 }
1250
1251 proven_block.output_note_batches_mut().push(pending_note_batch);
1252
1253 let updated_block_note_tree = proven_block.build_output_note_tree().root();
1254
1255 let block_header = proven_block.header();
1257 let updated_header = BlockHeader::new(
1258 block_header.version(),
1259 block_header.prev_block_commitment(),
1260 block_header.block_num(),
1261 block_header.chain_commitment(),
1262 self.account_tree.root(),
1263 self.nullifier_tree.root(),
1264 updated_block_note_tree,
1265 block_header.tx_commitment(),
1266 block_header.tx_kernel_commitment(),
1267 block_header.proof_commitment(),
1268 block_header.timestamp(),
1269 );
1270 proven_block.set_block_header(updated_header);
1271
1272 Ok(())
1273 }
1274
1275 fn prove_block_inner(&mut self, timestamp: Option<u32>) -> anyhow::Result<ProvenBlock> {
1296 let batches = self
1300 .pending_transactions_to_batches()
1301 .context("failed to convert pending transactions to batch")?;
1302
1303 let block_timestamp =
1307 timestamp.unwrap_or(self.latest_block_header().timestamp() + Self::TIMESTAMP_STEP_SECS);
1308
1309 let mut proven_block = self
1310 .propose_block_at(batches, block_timestamp)
1311 .context("failed to propose block")
1312 .and_then(|proposed_block| {
1313 self.prove_block(proposed_block)
1314 .context("failed to prove proposed block into proven block")
1315 })?;
1316
1317 self.apply_block_tree_updates(&proven_block)
1320 .context("failed to apply account and nullifier tree changes from block")?;
1321
1322 if !self.pending_objects.is_empty() {
1323 self.apply_pending_objects_to_block(&mut proven_block)
1324 .context("failed to add pending objects to block")?;
1325 }
1326
1327 self.apply_block(proven_block.clone())
1328 .context("failed to apply proven block to chain state")?;
1329
1330 Ok(proven_block)
1331 }
1332}
1333
1334fn create_genesis_state(
1337 accounts: impl IntoIterator<Item = Account>,
1338) -> anyhow::Result<(ProvenBlock, AccountTree)> {
1339 let block_account_updates: Vec<BlockAccountUpdate> = accounts
1340 .into_iter()
1341 .map(|account| {
1342 BlockAccountUpdate::new(
1343 account.id(),
1344 account.commitment(),
1345 AccountUpdateDetails::New(account),
1346 )
1347 })
1348 .collect();
1349
1350 let account_tree = AccountTree::with_entries(
1351 block_account_updates
1352 .iter()
1353 .map(|account| (account.account_id(), account.final_state_commitment())),
1354 )
1355 .context("failed to create genesis account tree")?;
1356
1357 let output_note_batches = Vec::new();
1358 let created_nullifiers = Vec::new();
1359 let transactions = OrderedTransactionHeaders::new_unchecked(Vec::new());
1360
1361 let version = 0;
1362 let prev_block_commitment = Digest::default();
1363 let block_num = BlockNumber::from(0u32);
1364 let chain_commitment = Blockchain::new().commitment();
1365 let account_root = account_tree.root();
1366 let nullifier_root = NullifierTree::new().root();
1367 let note_root = BlockNoteTree::empty().root();
1368 let tx_commitment = transactions.commitment();
1369 let tx_kernel_commitment = TransactionKernel::kernel_commitment();
1370 let proof_commitment = Digest::default();
1371 let timestamp = MockChain::TIMESTAMP_START_SECS;
1372
1373 let header = BlockHeader::new(
1374 version,
1375 prev_block_commitment,
1376 block_num,
1377 chain_commitment,
1378 account_root,
1379 nullifier_root,
1380 note_root,
1381 tx_commitment,
1382 tx_kernel_commitment,
1383 proof_commitment,
1384 timestamp,
1385 );
1386
1387 Ok((
1388 ProvenBlock::new_unchecked(
1389 header,
1390 block_account_updates,
1391 output_note_batches,
1392 created_nullifiers,
1393 transactions,
1394 ),
1395 account_tree,
1396 ))
1397}
1398
1399impl Default for MockChain {
1400 fn default() -> Self {
1401 MockChain::new()
1402 }
1403}
1404
1405#[derive(Default, Debug, Clone)]
1410struct PendingObjects {
1411 updated_accounts: BTreeMap<AccountId, BlockAccountUpdate>,
1413
1414 output_notes: Vec<OutputNote>,
1416
1417 created_nullifiers: Vec<Nullifier>,
1419}
1420
1421impl PendingObjects {
1422 pub fn new() -> PendingObjects {
1423 PendingObjects {
1424 updated_accounts: BTreeMap::new(),
1425 output_notes: vec![],
1426 created_nullifiers: vec![],
1427 }
1428 }
1429
1430 pub fn is_empty(&self) -> bool {
1432 self.updated_accounts.is_empty()
1433 && self.output_notes.is_empty()
1434 && self.created_nullifiers.is_empty()
1435 }
1436}
1437
1438pub enum AccountState {
1444 New,
1445 Exists,
1446}
1447
1448#[derive(Debug, Clone)]
1454pub enum TxContextInput {
1455 AccountId(AccountId),
1456 Account(Account),
1457 ExecutedTransaction(Box<ExecutedTransaction>),
1458}
1459
1460impl From<AccountId> for TxContextInput {
1461 fn from(account: AccountId) -> Self {
1462 Self::AccountId(account)
1463 }
1464}
1465
1466impl From<Account> for TxContextInput {
1467 fn from(account: Account) -> Self {
1468 Self::Account(account)
1469 }
1470}
1471
1472impl From<ExecutedTransaction> for TxContextInput {
1473 fn from(tx: ExecutedTransaction) -> Self {
1474 Self::ExecutedTransaction(Box::new(tx))
1475 }
1476}
1477
1478#[cfg(test)]
1482mod tests {
1483 use miden_objects::{
1484 account::{AccountStorage, AccountStorageMode},
1485 testing::account_component::AccountMockComponent,
1486 };
1487
1488 use super::*;
1489
1490 #[test]
1491 fn with_accounts() {
1492 let account = AccountBuilder::new([4; 32])
1493 .storage_mode(AccountStorageMode::Public)
1494 .with_component(
1495 AccountMockComponent::new_with_slots(
1496 TransactionKernel::testing_assembler(),
1497 vec![AccountStorage::mock_item_2().slot],
1498 )
1499 .unwrap(),
1500 )
1501 .build_existing()
1502 .unwrap();
1503
1504 let mock_chain = MockChain::with_accounts(&[account.clone()]);
1505
1506 assert_eq!(mock_chain.committed_account(account.id()), &account);
1507
1508 let tx_context = mock_chain.build_tx_context(account.id(), &[], &[]).build();
1511 assert_eq!(tx_context.tx_inputs().block_header().block_num(), BlockNumber::from(0u32));
1512 assert_eq!(
1513 tx_context.tx_inputs().block_header().account_root(),
1514 mock_chain.account_tree.root()
1515 );
1516 }
1517
1518 #[test]
1519 fn prove_until_block() -> anyhow::Result<()> {
1520 let mut chain = MockChain::new();
1521 let block = chain.prove_until_block(5)?;
1522 assert_eq!(block.header().block_num(), 5u32.into());
1523 assert_eq!(chain.proven_blocks().len(), 6);
1524
1525 Ok(())
1526 }
1527}