miden_testing/mock_chain/
chain.rs

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// MOCK CHAIN
47// ================================================================================================
48
49/// The [`MockChain`] simulates a simplified blockchain environment for testing purposes.
50/// It allows creating and managing accounts, minting assets, executing transactions, and applying
51/// state updates.
52///
53/// This struct is designed to mock transaction workflows, asset transfers, and
54/// note creation in a test setting. Once entities are set up, [`TransactionContextBuilder`] objects
55/// can be obtained in order to execute transactions accordingly.
56///
57/// On a high-level, there are two ways to interact with the mock chain:
58/// - Generating transactions yourself and adding them to the mock chain "mempool" using
59///   [`MockChain::add_pending_executed_transaction`] or
60///   [`MockChain::add_pending_proven_transaction`]. Once some transactions have been added, they
61///   can be proven into a block using [`MockChain::prove_next_block`], which commits them to the
62///   chain state.
63/// - Using any of the other pending APIs to _magically_ add new notes, accounts or nullifiers in
64///   the next block. For example, [`MockChain::add_pending_p2id_note`] will create a new P2ID note
65///   in the next proven block, without actually containing a transaction that creates that note.
66///
67/// Both approaches can be mixed in the same block, within limits. In particular, avoid modification
68/// of the _same_ entities using both regular transactions and the magic pending APIs.
69///
70/// The mock chain uses the batch and block provers underneath to process pending transactions, so
71/// the generated blocks are realistic and indistinguishable from a real node. The only caveat is
72/// that no real ZK proofs are generated or validated as part of transaction, batch or block
73/// building. If realistic data is important for your use case, avoid using any pending APIs except
74/// for [`MockChain::add_pending_executed_transaction`] and
75/// [`MockChain::add_pending_proven_transaction`].
76///
77/// # Examples
78///
79/// ## Create mock objects and build a transaction context
80/// ```no_run
81/// # use miden_testing::{Auth, MockChain, TransactionContextBuilder};
82/// # use miden_objects::{asset::FungibleAsset, Felt, note::NoteType};
83/// let mut mock_chain = MockChain::new();
84/// let faucet = mock_chain.add_pending_new_faucet(Auth::BasicAuth, "USDT", 100_000);  // Create a USDT faucet
85/// let asset = faucet.mint(1000);
86/// let sender = mock_chain.add_pending_new_wallet(Auth::BasicAuth);
87/// let target = mock_chain.add_pending_new_wallet(Auth::BasicAuth);
88/// let note = mock_chain
89///     .add_pending_p2id_note(
90///         faucet.id(),
91///         target.id(),
92///         &[FungibleAsset::mock(10)],
93///         NoteType::Public,
94///       None,
95///     )
96///   .unwrap();
97/// mock_chain.prove_next_block();
98/// let tx_context = mock_chain.build_tx_context(sender.id(), &[note.id()], &[]).build();
99/// let result = tx_context.execute();
100/// ```
101///
102/// ## Executing a Simple Transaction
103///
104/// ```
105/// # use miden_objects::{
106/// #   asset::{Asset, FungibleAsset},
107/// #   note::NoteType,
108/// # };
109/// # use miden_testing::{Auth, MockChain};
110/// let mut mock_chain = MockChain::new();
111/// // Add a recipient wallet.
112/// let receiver = mock_chain.add_pending_new_wallet(Auth::BasicAuth);
113/// // Add a wallet with assets.
114/// let sender = mock_chain.add_pending_existing_wallet(Auth::NoAuth, vec![]);
115/// let fungible_asset = FungibleAsset::mock(10).unwrap_fungible();
116///
117/// // Add a pending P2ID note to the chain.
118/// let note = mock_chain
119///     .add_pending_p2id_note(
120///         sender.id(),
121///         receiver.id(),
122///         &[Asset::Fungible(fungible_asset)],
123///         NoteType::Public,
124///         None,
125///     )
126///     .unwrap();
127/// // Prove the next block to add the pending note to the chain state, making it available for
128/// // consumption.
129/// mock_chain.prove_next_block();
130///
131/// // Create a transaction context that consumes the note and execute it.
132/// let transaction = mock_chain
133///     .build_tx_context(receiver.id(), &[note.id()], &[])
134///     .build()
135///     .execute()
136///     .unwrap();
137///
138/// // Add the transaction to the mock chain's "mempool" of pending transactions.
139/// mock_chain.add_pending_executed_transaction(&transaction);
140///
141/// // Prove the next block to include the transaction in the chain state.
142/// mock_chain.prove_next_block();
143///
144/// // Check that the receiver's balance has increased.
145/// assert_eq!(
146///     mock_chain
147///         .committed_account(receiver.id())
148///         .vault()
149///         .get_balance(fungible_asset.faucet_id())
150///         .unwrap(),
151///     fungible_asset.amount()
152/// );
153/// ```
154#[derive(Debug, Clone)]
155pub struct MockChain {
156    /// An append-only structure used to represent the history of blocks produced for this chain.
157    chain: Blockchain,
158
159    /// History of produced blocks.
160    blocks: Vec<ProvenBlock>,
161
162    /// Tree containing all nullifiers.
163    nullifier_tree: NullifierTree,
164
165    /// Tree containing the state commitments of all accounts.
166    account_tree: AccountTree,
167
168    /// Objects that have not yet been finalized.
169    ///
170    /// These will become available once the block is proven.
171    ///
172    /// Note:
173    /// - The [Note]s in this container do not have the `proof` set.
174    pending_objects: PendingObjects,
175
176    /// Transactions that have been submitted to the chain but have not yet been included in a
177    /// block.
178    pending_transactions: Vec<ProvenTransaction>,
179
180    /// NoteID |-> MockChainNote mapping to simplify note retrieval.
181    committed_notes: BTreeMap<NoteId, MockChainNote>,
182
183    /// AccountId |-> MockAccount mapping to simplify transaction creation.
184    committed_accounts: BTreeMap<AccountId, MockAccount>,
185
186    // The RNG used to generate note serial numbers, account seeds or cryptographic keys.
187    rng: ChaCha20Rng,
188}
189
190impl MockChain {
191    // CONSTANTS
192    // ----------------------------------------------------------------------------------------
193
194    /// The timestamp of the genesis block of the chain. Chosen as an easily readable number.
195    pub const TIMESTAMP_START_SECS: u32 = 1700000000;
196
197    /// The number of seconds by which a block's timestamp increases over the previous block's
198    /// timestamp, unless overwritten when calling [`Self::prove_next_block_at`].
199    pub const TIMESTAMP_STEP_SECS: u32 = 10;
200
201    // CONSTRUCTORS
202    // ----------------------------------------------------------------------------------------
203
204    /// Creates a new `MockChain` with an empty genesis block.
205    pub fn new() -> Self {
206        Self::with_accounts(&[])
207    }
208
209    /// Creates a new `MockChain` with a genesis block containing the provided accounts.
210    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            // Initialize RNG with default seed.
224            rng: ChaCha20Rng::from_seed(Default::default()),
225        };
226
227        // We do not have to apply the tree changes, because the account tree is already initialized
228        // and the nullifier tree is empty at genesis.
229        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    // PUBLIC ACCESSORS
252    // ----------------------------------------------------------------------------------------
253
254    /// Returns a reference to the current [`Blockchain`].
255    pub fn blockchain(&self) -> &Blockchain {
256        &self.chain
257    }
258
259    /// Returns a [`PartialBlockchain`] instantiated from the current [`Blockchain`] and with
260    /// authentication paths for all all blocks in the chain.
261    pub fn latest_partial_blockchain(&self) -> PartialBlockchain {
262        // We have to exclude the latest block because we need to fetch the state of the chain at
263        // that latest block, which does not include itself.
264        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    /// Creates a new [`PartialBlockchain`] with all reference blocks in the given iterator except
271    /// for the latest block header in the chain and returns that latest block header.
272    ///
273    /// The intended use for the latest block header is to become the reference block of a new
274    /// transaction batch or block.
275    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        // Deduplicate block numbers so each header will be included just once. This is required so
281        // PartialBlockchain::from_blockchain does not panic.
282        let reference_blocks: BTreeSet<_> = reference_blocks.into_iter().collect();
283
284        // Include all block headers of the reference blocks except the latest block.
285        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    /// Returns a map of [`AccountWitness`]es for the requested account IDs from the current
298    /// [`AccountTree`] in the chain.
299    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    /// Returns a map of [`NullifierWitness`]es for the requested nullifiers from the current
314    /// [`NullifierTree`] in the chain.
315    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    /// Returns all note inclusion proofs for the requested note IDs, **if they are available for
330    /// consumption**. Therefore, not all of the requested notes will be guaranteed to have an entry
331    /// in the returned map.
332    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(&note) {
339                proofs.insert(note, input_note.inclusion_proof().clone());
340            }
341        }
342
343        proofs
344    }
345
346    /// Returns a reference to the latest [`BlockHeader`] in the chain.
347    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    /// Returns the [`BlockHeader`] with the specified `block_number`.
354    pub fn block_header(&self, block_number: usize) -> BlockHeader {
355        self.blocks[block_number].header().clone()
356    }
357
358    /// Returns a reference to slice of all created proven blocks.
359    pub fn proven_blocks(&self) -> &[ProvenBlock] {
360        &self.blocks
361    }
362
363    /// Returns a reference to the nullifier tree.
364    pub fn nullifier_tree(&self) -> &NullifierTree {
365        &self.nullifier_tree
366    }
367
368    /// Returns the map of note IDs to committed notes.
369    ///
370    /// These notes are committed for authenticated consumption.
371    pub fn committed_notes(&self) -> &BTreeMap<NoteId, MockChainNote> {
372        &self.committed_notes
373    }
374
375    /// Returns an [`InputNote`] for the given note ID. If the note does not exist or is not
376    /// public, `None` is returned.
377    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    /// Returns a reference to the account identified by the given account ID and panics if it does
383    /// not exist.
384    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    /// Returns a reference to the [`AccountTree`] of the chain.
392    pub fn account_tree(&self) -> &AccountTree {
393        &self.account_tree
394    }
395
396    // BATCH APIS
397    // ----------------------------------------------------------------------------------------
398
399    /// Proposes a new transaction batch from the provided transactions and returns it.
400    ///
401    /// This method does not modify the chain state.
402    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    /// Mock-proves a proposed transaction batch from the provided [`ProposedBatch`] and returns it.
428    ///
429    /// This method does not modify the chain state.
430    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        // SAFETY: This satisfies the requirements of the ordered tx headers.
444        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    // BLOCK APIS
466    // ----------------------------------------------------------------------------------------
467
468    /// Proposes a new block from the provided batches with the given timestamp and returns it.
469    ///
470    /// This method does not modify the chain state.
471    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    /// Proposes a new block from the provided batches and returns it.
488    ///
489    /// This method does not modify the chain state.
490    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        // We can't access system time because we are in a no-std environment, so we use the
498        // minimally correct next timestamp.
499        let timestamp = self.latest_block_header().timestamp() + 1;
500
501        self.propose_block_at(batches, timestamp)
502    }
503
504    /// Mock-proves a proposed block into a proven block and returns it.
505    ///
506    /// This method does not modify the chain state.
507    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    // TRANSACTION APIS
515    // ----------------------------------------------------------------------------------------
516
517    /// Initializes a [`TransactionContextBuilder`].
518    ///
519    /// Depending on the provided `input`, the builder is initialized differently:
520    /// - [`TxContextInput::AccountId`]: Initialize the builder with [`TransactionInputs`] fetched
521    ///   from the chain for the account identified by the ID.
522    /// - [`TxContextInput::Account`]: Initialize the builder with [`TransactionInputs`] where the
523    ///   account is passed as-is to the inputs.
524    /// - [`TxContextInput::ExecutedTransaction`]: Initialize the builder with [`TransactionInputs`]
525    ///   where the account passed to the inputs is the final account of the executed transaction.
526    ///   This is the initial account of the transaction with the account delta applied.
527    ///
528    /// In all cases, if the chain contains a seed or authenticator for the account, they are added
529    /// to the builder. Additionally, if the account is set to authenticate with
530    /// [`Auth::BasicAuth`], the executed transaction script is defaulted to
531    /// [`DEFAULT_AUTH_SCRIPT`].
532    ///
533    /// [`TxContextInput::Account`] and [`TxContextInput::ExecutedTransaction`] can be used to build
534    /// a chain of transactions against the same account that build on top of each other. For
535    /// example, transaction A modifies an account from state 0 to 1, and transaction B modifies
536    /// it from state 1 to 2.
537    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    // INPUTS APIS
591    // ----------------------------------------------------------------------------------------
592
593    /// Returns a valid [`TransactionInputs`] for the specified entities.
594    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 the account is new, add the anchor block's header from which the account ID is derived
625        // to the MMR.
626        if account.is_new() {
627            let epoch_block_num = BlockNumber::from_epoch(account.id().anchor_epoch());
628            // The reference block of the transaction is added to the MMR in
629            // prologue::process_chain_data so we can skip adding it to the block headers here.
630            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    /// Returns inputs for a transaction batch for all the reference blocks of the provided
656    /// transactions.
657    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        // Fetch note proofs for notes that exist in the chain.
663        let unauthenticated_note_proofs = self.unauthenticated_note_proofs(unauthenticated_notes);
664
665        // We also need to fetch block inclusion proofs for any of the blocks that contain
666        // unauthenticated notes for which we want to prove inclusion.
667        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    /// Gets foreign account inputs to execute FPI transactions.
680    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 there are storage maps, we populate the merkle store and advice map
689            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    /// Gets the inputs for a block for the provided batches.
699    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    // PUBLIC MUTATORS
735    // ----------------------------------------------------------------------------------------
736
737    /// Creates the next block in the mock chain.
738    ///
739    /// This will make all the objects currently pending available for use.
740    pub fn prove_next_block(&mut self) -> ProvenBlock {
741        self.prove_block_inner(None).unwrap()
742    }
743
744    /// Proves the next block in the mock chain at the given timestamp.
745    pub fn prove_next_block_at(&mut self, timestamp: u32) -> anyhow::Result<ProvenBlock> {
746        self.prove_block_inner(Some(timestamp))
747    }
748
749    /// Proves new blocks until the block with the given target block number has been created.
750    ///
751    /// For example, if the latest block is `5` and this function is called with `10`, then blocks
752    /// `6..=10` will be created and block 10 will be returned.
753    ///
754    /// # Panics
755    ///
756    /// Panics if:
757    /// - the given block number is smaller or equal to the number of the latest block in the chain.
758    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    /// Sets the seed for the internal RNG.
778    pub fn set_rng_seed(&mut self, seed: [u8; 32]) {
779        self.rng = ChaCha20Rng::from_seed(seed);
780    }
781
782    // PUBLIC MUTATORS (PENDING APIS)
783    // ----------------------------------------------------------------------------------------
784
785    /// Adds the given [`ExecutedTransaction`] to the list of pending transactions.
786    ///
787    /// A block has to be created to apply the transaction effects to the chain state, e.g. using
788    /// [`MockChain::prove_next_block`].
789    ///
790    /// Returns the resulting state of the executing account after executing the transaction.
791    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        // This essentially transforms an executed tx into a proven tx with a dummy proof.
799        let proven_tx = ProvenTransaction::from_executed_transaction_mocked(transaction.clone());
800
801        self.pending_transactions.push(proven_tx);
802
803        account
804    }
805
806    /// Adds the given [`ProvenTransaction`] to the list of pending transactions.
807    ///
808    /// A block has to be created to apply the transaction effects to the chain state, e.g. using
809    /// [`MockChain::prove_next_block`].
810    pub fn add_pending_proven_transaction(&mut self, transaction: ProvenTransaction) {
811        self.pending_transactions.push(transaction);
812    }
813
814    /// Adds the given [`OutputNote`] to the list of pending notes.
815    ///
816    /// A block has to be created to add the note to that block and make it available in the chain
817    /// state, e.g. using [`MockChain::prove_next_block`].
818    pub fn add_pending_note(&mut self, note: OutputNote) {
819        self.pending_objects.output_notes.push(note);
820    }
821
822    /// Adds a P2ID [`OutputNote`] to the list of pending notes.
823    ///
824    /// A block has to be created to add the note to that block and make it available in the chain
825    /// state, e.g. using [`MockChain::prove_next_block`].
826    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    /// Adds the [`Nullifier`] to the list of pending nullifiers.
863    ///
864    /// A block has to be created to add the nullifier to the nullifier tree as part of that block,
865    /// e.g. using [`MockChain::prove_next_block`].
866    pub fn add_pending_nullifier(&mut self, nullifier: Nullifier) {
867        self.pending_objects.created_nullifiers.push(nullifier);
868    }
869
870    /// Adds a new [`BasicWallet`] account to the list of pending accounts.
871    ///
872    /// A block has to be created to add the account to the chain state as part of that block,
873    /// e.g. using [`MockChain::prove_next_block`].
874    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    /// Adds an existing [`BasicWallet`] account with nonce `1` to the list of pending accounts.
883    ///
884    /// A block has to be created to add the account to the chain state as part of that block,
885    /// e.g. using [`MockChain::prove_next_block`].
886    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    /// Adds a new [`BasicFungibleFaucet`] account with the specified authentication method and the
900    /// given token metadata to the list of pending accounts.
901    ///
902    /// A block has to be created to add the account to the chain state as part of that block,
903    /// e.g. using [`MockChain::prove_next_block`].
904    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    /// Adds an existing [`BasicFungibleFaucet`] account with the specified authentication method
930    /// and the given token metadata to the list of pending accounts.
931    ///
932    /// A block has to be created to add the account to the chain state as part of that block,
933    /// e.g. using [`MockChain::prove_next_block`].
934    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        // The faucet's reserved slot is initialized to an empty word by default.
963        // If total_issuance is set, overwrite it.
964        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        // We have to insert these into the committed accounts so the authenticator is available.
972        // Without this, the account couldn't be authenticated.
973        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    /// Adds the [`AccountComponent`](miden_objects::account::AccountComponent) corresponding to
981    /// `auth_method` to the account in the builder and builds a new or existing account
982    /// depending on `account_state`.
983    ///
984    /// The account is added to the list of committed accounts _and_, if [`AccountState::Exists`] is
985    /// passed, is also added to the list of pending accounts. Adding it to committed accounts
986    /// makes the account seed and authenticator available for account creation and
987    /// authentication, respectively. If the account exists, then the next block that is created
988    /// will add the pending accounts to the chain state.
989    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        // Add account to the committed accounts so transaction inputs can be retrieved via the mock
1016        // chain APIs.
1017        // We also have to insert these into the committed accounts so the account seed and
1018        // authenticator are available. Without this, the account couldn't be created or
1019        // authenticated.
1020        self.committed_accounts
1021            .insert(account.id(), MockAccount::new(account.clone(), seed, authenticator));
1022
1023        // Do not add new accounts to the pending accounts. Usually, new accounts are added in tests
1024        // to create transactions that create these new accounts in the chain. If we add them to the
1025        // pending accounts and the test calls prove_next_block (which typically happens to add all
1026        // pending objects to the chain state as part of the test setup), the account will be added
1027        // to the chain which means the account-creating transaction fails because the
1028        // account already exists. So new accounts are added only to the committed accounts
1029        // so the transaction context APIs work as expected.
1030        if let AccountState::Exists = account_state {
1031            self.add_pending_account(account.clone());
1032        }
1033
1034        account
1035    }
1036
1037    /// Adds a new `Account` to the list of pending objects.
1038    ///
1039    /// A block has to be created to finalize the new entity.
1040    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    // PRIVATE HELPERS
1052    // ----------------------------------------------------------------------------------------
1053
1054    /// Inserts the given block's account updates and created nullifiers into the account tree and
1055    /// nullifier tree, respectively.
1056    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            // TODO: Remove from self.committed_notes. This is not critical to have for now. It is
1069            // not straightforward, because committed_notes are indexed by note IDs rather than
1070            // nullifiers, so we'll have to create a second index to do this.
1071        }
1072
1073        Ok(())
1074    }
1075
1076    /// Applies the given block to the chain state, which means:
1077    ///
1078    /// - Updated accounts from the block are updated in the committed accounts.
1079    /// - Created notes are inserted into the committed notes.
1080    /// - Consumed notes are removed from the committed notes.
1081    /// - The block is appended to the [`BlockChain`] and the list of proven blocks.
1082    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        // Batches must contain at least one transaction, so if there are no pending transactions,
1156        // return early.
1157        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        // TODO: Distribute the transactions into multiple batches if the transactions would not fit
1164        // into a single batch (according to max input notes, max output notes and max accounts).
1165        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        // Add pending accounts to block.
1177        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        // Add pending nullifiers to block.
1200        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        // Add pending output notes to block.
1221        let output_notes_block: BTreeSet<NoteId> =
1222            proven_block.output_notes().map(|(_, output_note)| output_note.id()).collect();
1223
1224        // We could distribute notes over multiple batches (if space is available), but most likely
1225        // one is sufficient.
1226        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        // Update account tree and nullifier tree root in the block.
1256        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    /// Creates a new block in the mock chain.
1276    ///
1277    /// This will make all the objects currently pending available for use.
1278    ///
1279    /// If a `timestamp` is provided, it will be set on the block.
1280    ///
1281    /// Block building is divided into a few steps:
1282    ///
1283    /// 1. Build batches from pending transactions and a block from those batches. This results in a
1284    ///    block.
1285    /// 2. Take that block and apply only its account/nullifier tree updates to the chain.
1286    /// 3. Then take the pending objects and insert them directly into the proven block. This means
1287    ///    we have to update the header of the block as well, with the newly inserted pending
1288    ///    accounts/nullifiers/notes. This is why we already did step 2, so that we can insert the
1289    ///    pending objects directly into the account/nullifier tree to get the latest correct state
1290    ///    of those trees. Then take the root of the trees and update them in the header of the
1291    ///    block. This should be pretty efficient because we don't have to do any tree insertions
1292    ///    multiple times (which would be slow).
1293    /// 4. Finally, now the block contains both the updates from the regular transactions/batches as
1294    ///    well as the pending objects. Now insert all the remaining updates into the chain state.
1295    fn prove_block_inner(&mut self, timestamp: Option<u32>) -> anyhow::Result<ProvenBlock> {
1296        // Create batches from pending transactions.
1297        // ----------------------------------------------------------------------------------------
1298
1299        let batches = self
1300            .pending_transactions_to_batches()
1301            .context("failed to convert pending transactions to batch")?;
1302
1303        // Create block.
1304        // ----------------------------------------------------------------------------------------
1305
1306        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        // We apply the block tree updates here, so that apply_pending_objects_to_block can easily
1318        // update the block header of this block with the pending accounts and nullifiers.
1319        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
1334/// Creates the genesis state, consisting of a block containing the provided account updates and an
1335/// account tree with those accounts.
1336fn 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// PENDING OBJECTS
1406// ================================================================================================
1407
1408/// Aggregates all entities that were added using the _pending_ APIs of the [`MockChain`].
1409#[derive(Default, Debug, Clone)]
1410struct PendingObjects {
1411    /// Account updates for the block.
1412    updated_accounts: BTreeMap<AccountId, BlockAccountUpdate>,
1413
1414    /// Note batches created in transactions in the block.
1415    output_notes: Vec<OutputNote>,
1416
1417    /// Nullifiers produced in transactions in the block.
1418    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    /// Returns `true` if there are no pending objects, `false` otherwise.
1431    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
1438// ACCOUNT STATE
1439// ================================================================================================
1440
1441/// Helper type for increased readability at call-sites. Indicates whether to build a new (nonce =
1442/// ZERO) or existing account (nonce = ONE).
1443pub enum AccountState {
1444    New,
1445    Exists,
1446}
1447
1448// TX CONTEXT INPUT
1449// ================================================================================================
1450
1451/// Helper type to abstract over the inputs to [`MockChain::build_tx_context`]. See that method's
1452/// docs for details.
1453#[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// TESTS
1479// ================================================================================================
1480
1481#[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        // Check that transaction inputs retrieved from the chain are against the block header with
1509        // the current account tree root.
1510        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}