miden_testing/mock_chain/
chain.rs

1use alloc::collections::{BTreeMap, BTreeSet};
2use alloc::vec::Vec;
3
4use anyhow::Context;
5use miden_block_prover::{LocalBlockProver, ProvenBlockError};
6use miden_objects::account::auth::AuthSecretKey;
7use miden_objects::account::delta::AccountUpdateDetails;
8use miden_objects::account::{Account, AccountId, PartialAccount};
9use miden_objects::batch::{ProposedBatch, ProvenBatch};
10use miden_objects::block::account_tree::AccountTree;
11use miden_objects::block::{
12    AccountWitness,
13    BlockHeader,
14    BlockInputs,
15    BlockNumber,
16    Blockchain,
17    NullifierTree,
18    NullifierWitness,
19    ProposedBlock,
20    ProvenBlock,
21};
22use miden_objects::note::{Note, NoteHeader, NoteId, NoteInclusionProof, Nullifier};
23use miden_objects::transaction::{
24    ExecutedTransaction,
25    InputNote,
26    InputNotes,
27    OutputNote,
28    PartialBlockchain,
29    ProvenTransaction,
30    TransactionInputs,
31};
32use miden_processor::DeserializationError;
33use miden_tx::LocalTransactionProver;
34use miden_tx::auth::BasicAuthenticator;
35use miden_tx::utils::{ByteReader, Deserializable, Serializable};
36use miden_tx_batch_prover::LocalBatchProver;
37use winterfell::ByteWriter;
38
39use super::note::MockChainNote;
40use crate::{MockChainBuilder, TransactionContextBuilder};
41
42// MOCK CHAIN
43// ================================================================================================
44
45/// The [`MockChain`] simulates a simplified blockchain environment for testing purposes.
46///
47/// The typical usage of a mock chain is:
48/// - Creating it using a [`MockChainBuilder`], which allows adding accounts and notes to the
49///   genesis state.
50/// - Creating transactions against the chain state and executing them.
51/// - Adding executed or proven transactions to the set of pending transactions (the "mempool"),
52///   e.g. using [`MockChain::add_pending_executed_transaction`].
53/// - Proving a block, which adds all pending transactions to the chain state, e.g. using
54///   [`MockChain::prove_next_block`].
55///
56/// The mock chain uses the batch and block provers underneath to process pending transactions, so
57/// the generated blocks are realistic and indistinguishable from a real node. The only caveat is
58/// that no real ZK proofs are generated or validated as part of transaction, batch or block
59/// building.
60///
61/// # Examples
62///
63/// ## Executing a simple transaction
64/// ```
65/// # use anyhow::Result;
66/// # use miden_objects::{
67/// #    asset::{Asset, FungibleAsset},
68/// #    note::NoteType,
69/// # };
70/// # use miden_testing::{Auth, MockChain};
71/// #
72/// # #[tokio::main(flavor = "current_thread")]
73/// # async fn main() -> Result<()> {
74/// // Build a genesis state for a mock chain using a MockChainBuilder.
75/// // --------------------------------------------------------------------------------------------
76///
77/// let mut builder = MockChain::builder();
78///
79/// // Add a recipient wallet.
80/// let receiver = builder.add_existing_wallet(Auth::BasicAuth)?;
81///
82/// // Add a wallet with assets.
83/// let sender = builder.add_existing_wallet(Auth::IncrNonce)?;
84///
85/// let fungible_asset = FungibleAsset::mock(10).unwrap_fungible();
86/// // Add a P2ID note with a fungible asset to the chain.
87/// let note = builder.add_p2id_note(
88///     sender.id(),
89///     receiver.id(),
90///     &[Asset::Fungible(fungible_asset)],
91///     NoteType::Public,
92/// )?;
93///
94/// let mut mock_chain: MockChain = builder.build()?;
95///
96/// // Create a transaction against the receiver account consuming the note.
97/// // --------------------------------------------------------------------------------------------
98///
99/// let transaction = mock_chain
100///     .build_tx_context(receiver.id(), &[note.id()], &[])?
101///     .build()?
102///     .execute()
103///     .await?;
104///
105/// // Add the transaction to the chain state.
106/// // --------------------------------------------------------------------------------------------
107///
108/// // Add the transaction to the mock chain's "mempool" of pending transactions.
109/// mock_chain.add_pending_executed_transaction(&transaction)?;
110///
111/// // Prove the next block to include the transaction in the chain state.
112/// mock_chain.prove_next_block()?;
113///
114/// // The receiver account should now have the asset in its account vault.
115/// assert_eq!(
116///     mock_chain
117///         .committed_account(receiver.id())?
118///         .vault()
119///         .get_balance(fungible_asset.faucet_id())?,
120///     fungible_asset.amount()
121/// );
122/// # Ok(())
123/// # }
124/// ```
125///
126/// ## Create mock objects and build a transaction context
127///
128/// ```
129/// # use anyhow::Result;
130/// # use miden_objects::{Felt, asset::{Asset, FungibleAsset}, note::NoteType};
131/// # use miden_testing::{Auth, MockChain, TransactionContextBuilder};
132/// #
133/// # #[tokio::main(flavor = "current_thread")]
134/// # async fn main() -> Result<()> {
135/// let mut builder = MockChain::builder();
136///
137/// let faucet = builder.create_new_faucet(Auth::BasicAuth, "USDT", 100_000)?;
138/// let asset = Asset::from(FungibleAsset::new(faucet.id(), 10)?);
139///
140/// let sender = builder.create_new_wallet(Auth::BasicAuth)?;
141/// let target = builder.create_new_wallet(Auth::BasicAuth)?;
142///
143/// let note = builder.add_p2id_note(faucet.id(), target.id(), &[asset], NoteType::Public)?;
144///
145/// let mock_chain = builder.build()?;
146///
147/// // The target account is a new account so we move it into the build_tx_context, since the
148/// // chain's committed accounts do not yet contain it.
149/// let tx_context = mock_chain.build_tx_context(target, &[note.id()], &[])?.build()?;
150/// let executed_transaction = tx_context.execute().await?;
151/// # Ok(())
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    /// Transactions that have been submitted to the chain but have not yet been included in a
169    /// block.
170    pending_transactions: Vec<ProvenTransaction>,
171
172    /// NoteID |-> MockChainNote mapping to simplify note retrieval.
173    committed_notes: BTreeMap<NoteId, MockChainNote>,
174
175    /// AccountId |-> Account mapping to simplify transaction creation. Latest known account
176    /// state is maintained for each account here.
177    ///
178    /// The map always holds the most recent *public* state known for every account. For private
179    /// accounts, however, transactions do not emit the post-transaction state, so their entries
180    /// remain at the last observed state.
181    committed_accounts: BTreeMap<AccountId, Account>,
182
183    /// AccountId |-> AccountAuthenticator mapping to store the authenticator for accounts to
184    /// simplify transaction creation.
185    account_authenticators: BTreeMap<AccountId, AccountAuthenticator>,
186}
187
188impl MockChain {
189    // CONSTANTS
190    // ----------------------------------------------------------------------------------------
191
192    /// The timestamp of the genesis block of the chain. Chosen as an easily readable number.
193    pub const TIMESTAMP_START_SECS: u32 = 1700000000;
194
195    /// The number of seconds by which a block's timestamp increases over the previous block's
196    /// timestamp, unless overwritten when calling [`Self::prove_next_block_at`].
197    pub const TIMESTAMP_STEP_SECS: u32 = 10;
198
199    // CONSTRUCTORS
200    // ----------------------------------------------------------------------------------------
201
202    /// Creates a new `MockChain` with an empty genesis block.
203    pub fn new() -> Self {
204        Self::builder().build().expect("empty chain should be valid")
205    }
206
207    /// Returns a new, empty [`MockChainBuilder`].
208    pub fn builder() -> MockChainBuilder {
209        MockChainBuilder::new()
210    }
211
212    /// Creates a new `MockChain` with the provided genesis block and account tree.
213    pub(super) fn from_genesis_block(
214        genesis_block: ProvenBlock,
215        account_tree: AccountTree,
216        account_authenticators: BTreeMap<AccountId, AccountAuthenticator>,
217    ) -> anyhow::Result<Self> {
218        let mut chain = MockChain {
219            chain: Blockchain::default(),
220            blocks: vec![],
221            nullifier_tree: NullifierTree::default(),
222            account_tree,
223            pending_transactions: Vec::new(),
224            committed_notes: BTreeMap::new(),
225            committed_accounts: BTreeMap::new(),
226            account_authenticators,
227        };
228
229        // We do not have to apply the tree changes, because the account tree is already initialized
230        // and the nullifier tree is empty at genesis.
231        chain
232            .apply_block(genesis_block)
233            .context("failed to build account from builder")?;
234
235        debug_assert_eq!(chain.blocks.len(), 1);
236        debug_assert_eq!(chain.committed_accounts.len(), chain.account_tree.num_accounts());
237
238        Ok(chain)
239    }
240
241    // PUBLIC ACCESSORS
242    // ----------------------------------------------------------------------------------------
243
244    /// Returns a reference to the current [`Blockchain`].
245    pub fn blockchain(&self) -> &Blockchain {
246        &self.chain
247    }
248
249    /// Returns a [`PartialBlockchain`] instantiated from the current [`Blockchain`] and with
250    /// authentication paths for all all blocks in the chain.
251    pub fn latest_partial_blockchain(&self) -> PartialBlockchain {
252        // We have to exclude the latest block because we need to fetch the state of the chain at
253        // that latest block, which does not include itself.
254        let block_headers =
255            self.blocks.iter().map(|b| b.header()).take(self.blocks.len() - 1).cloned();
256
257        PartialBlockchain::from_blockchain(&self.chain, block_headers)
258            .expect("blockchain should be valid by construction")
259    }
260
261    /// Creates a new [`PartialBlockchain`] with all reference blocks in the given iterator except
262    /// for the latest block header in the chain and returns that latest block header.
263    ///
264    /// The intended use for the latest block header is to become the reference block of a new
265    /// transaction batch or block.
266    pub fn latest_selective_partial_blockchain(
267        &self,
268        reference_blocks: impl IntoIterator<Item = BlockNumber>,
269    ) -> anyhow::Result<(BlockHeader, PartialBlockchain)> {
270        let latest_block_header = self.latest_block_header();
271
272        self.selective_partial_blockchain(latest_block_header.block_num(), reference_blocks)
273    }
274
275    /// Creates a new [`PartialBlockchain`] with all reference blocks in the given iterator except
276    /// for the reference block header in the chain and returns that reference block header.
277    ///
278    /// The intended use for the reference block header is to become the reference block of a new
279    /// transaction batch or block.
280    pub fn selective_partial_blockchain(
281        &self,
282        reference_block: BlockNumber,
283        reference_blocks: impl IntoIterator<Item = BlockNumber>,
284    ) -> anyhow::Result<(BlockHeader, PartialBlockchain)> {
285        let reference_block_header = self.block_header(reference_block.as_usize());
286        // Deduplicate block numbers so each header will be included just once. This is required so
287        // PartialBlockchain::from_blockchain does not panic.
288        let reference_blocks: BTreeSet<_> = reference_blocks.into_iter().collect();
289
290        // Include all block headers except the reference block itself.
291        let mut block_headers = Vec::new();
292
293        for block_ref_num in &reference_blocks {
294            let block_index = block_ref_num.as_usize();
295            let block = self
296                .blocks
297                .get(block_index)
298                .ok_or_else(|| anyhow::anyhow!("block {} not found in chain", block_ref_num))?;
299            let block_header = block.header().clone();
300            // Exclude the reference block header.
301            if block_header.commitment() != reference_block_header.commitment() {
302                block_headers.push(block_header);
303            }
304        }
305
306        let partial_blockchain =
307            PartialBlockchain::from_blockchain_at(&self.chain, reference_block, block_headers)?;
308
309        Ok((reference_block_header, partial_blockchain))
310    }
311
312    /// Returns a map of [`AccountWitness`]es for the requested account IDs from the current
313    /// [`AccountTree`] in the chain.
314    pub fn account_witnesses(
315        &self,
316        account_ids: impl IntoIterator<Item = AccountId>,
317    ) -> BTreeMap<AccountId, AccountWitness> {
318        let mut account_witnesses = BTreeMap::new();
319
320        for account_id in account_ids {
321            let witness = self.account_tree.open(account_id);
322            account_witnesses.insert(account_id, witness);
323        }
324
325        account_witnesses
326    }
327
328    /// Returns a map of [`NullifierWitness`]es for the requested nullifiers from the current
329    /// [`NullifierTree`] in the chain.
330    pub fn nullifier_witnesses(
331        &self,
332        nullifiers: impl IntoIterator<Item = Nullifier>,
333    ) -> BTreeMap<Nullifier, NullifierWitness> {
334        let mut nullifier_proofs = BTreeMap::new();
335
336        for nullifier in nullifiers {
337            let witness = self.nullifier_tree.open(&nullifier);
338            nullifier_proofs.insert(nullifier, witness);
339        }
340
341        nullifier_proofs
342    }
343
344    /// Returns all note inclusion proofs for the requested note IDs, **if they are available for
345    /// consumption**. Therefore, not all of the requested notes will be guaranteed to have an entry
346    /// in the returned map.
347    pub fn unauthenticated_note_proofs(
348        &self,
349        notes: impl IntoIterator<Item = NoteId>,
350    ) -> BTreeMap<NoteId, NoteInclusionProof> {
351        let mut proofs = BTreeMap::default();
352        for note in notes {
353            if let Some(input_note) = self.committed_notes.get(&note) {
354                proofs.insert(note, input_note.inclusion_proof().clone());
355            }
356        }
357
358        proofs
359    }
360
361    /// Returns the genesis [`BlockHeader`] of the chain.
362    pub fn genesis_block_header(&self) -> BlockHeader {
363        self.block_header(BlockNumber::GENESIS.as_usize())
364    }
365
366    /// Returns the latest [`BlockHeader`] in the chain.
367    pub fn latest_block_header(&self) -> BlockHeader {
368        let chain_tip =
369            self.chain.chain_tip().expect("chain should contain at least the genesis block");
370        self.blocks[chain_tip.as_usize()].header().clone()
371    }
372
373    /// Returns the [`BlockHeader`] with the specified `block_number`.
374    ///
375    /// # Panics
376    ///
377    /// - If the block number does not exist in the chain.
378    pub fn block_header(&self, block_number: usize) -> BlockHeader {
379        self.blocks[block_number].header().clone()
380    }
381
382    /// Returns a reference to slice of all created proven blocks.
383    pub fn proven_blocks(&self) -> &[ProvenBlock] {
384        &self.blocks
385    }
386
387    /// Returns the [`AccountId`] of the faucet whose assets are accepted for fee payments in the
388    /// transaction kernel, or in other words, the native asset of the blockchain.
389    ///
390    /// This value is taken from the genesis block because it is assumed not to change throughout
391    /// the chain's lifecycle.
392    pub fn native_asset_id(&self) -> AccountId {
393        self.genesis_block_header().fee_parameters().native_asset_id()
394    }
395
396    /// Returns a reference to the nullifier tree.
397    pub fn nullifier_tree(&self) -> &NullifierTree {
398        &self.nullifier_tree
399    }
400
401    /// Returns the map of note IDs to committed notes.
402    ///
403    /// These notes are committed for authenticated consumption.
404    pub fn committed_notes(&self) -> &BTreeMap<NoteId, MockChainNote> {
405        &self.committed_notes
406    }
407
408    /// Returns an [`InputNote`] for the given note ID. If the note does not exist or is not
409    /// public, `None` is returned.
410    pub fn get_public_note(&self, note_id: &NoteId) -> Option<InputNote> {
411        let note = self.committed_notes.get(note_id)?;
412        note.clone().try_into().ok()
413    }
414
415    /// Returns a reference to the account identified by the given account ID.
416    ///
417    /// The account is retrieved with the latest state known to the [`MockChain`].
418    pub fn committed_account(&self, account_id: AccountId) -> anyhow::Result<&Account> {
419        self.committed_accounts
420            .get(&account_id)
421            .with_context(|| format!("account {account_id} not found in committed accounts"))
422    }
423
424    /// Returns a reference to the [`AccountTree`] of the chain.
425    pub fn account_tree(&self) -> &AccountTree {
426        &self.account_tree
427    }
428
429    // BATCH APIS
430    // ----------------------------------------------------------------------------------------
431
432    /// Proposes a new transaction batch from the provided transactions and returns it.
433    ///
434    /// This method does not modify the chain state.
435    pub fn propose_transaction_batch<I>(
436        &self,
437        txs: impl IntoIterator<Item = ProvenTransaction, IntoIter = I>,
438    ) -> anyhow::Result<ProposedBatch>
439    where
440        I: Iterator<Item = ProvenTransaction> + Clone,
441    {
442        let transactions: Vec<_> = txs.into_iter().map(alloc::sync::Arc::new).collect();
443
444        let (batch_reference_block, partial_blockchain, unauthenticated_note_proofs) = self
445            .get_batch_inputs(
446                transactions.iter().map(|tx| tx.ref_block_num()),
447                transactions
448                    .iter()
449                    .flat_map(|tx| tx.unauthenticated_notes().map(NoteHeader::id)),
450            )?;
451
452        Ok(ProposedBatch::new(
453            transactions,
454            batch_reference_block,
455            partial_blockchain,
456            unauthenticated_note_proofs,
457        )?)
458    }
459
460    /// Mock-proves a proposed transaction batch from the provided [`ProposedBatch`] and returns it.
461    ///
462    /// This method does not modify the chain state.
463    pub fn prove_transaction_batch(
464        &self,
465        proposed_batch: ProposedBatch,
466    ) -> anyhow::Result<ProvenBatch> {
467        let batch_prover = LocalBatchProver::new(0);
468        Ok(batch_prover.prove_dummy(proposed_batch)?)
469    }
470
471    // BLOCK APIS
472    // ----------------------------------------------------------------------------------------
473
474    /// Proposes a new block from the provided batches with the given timestamp and returns it.
475    ///
476    /// This method does not modify the chain state.
477    pub fn propose_block_at<I>(
478        &self,
479        batches: impl IntoIterator<Item = ProvenBatch, IntoIter = I>,
480        timestamp: u32,
481    ) -> anyhow::Result<ProposedBlock>
482    where
483        I: Iterator<Item = ProvenBatch> + Clone,
484    {
485        let batches: Vec<_> = batches.into_iter().collect();
486
487        let block_inputs = self
488            .get_block_inputs(batches.iter())
489            .context("could not retrieve block inputs")?;
490
491        let proposed_block = ProposedBlock::new_at(block_inputs, batches, timestamp)
492            .context("failed to create proposed block")?;
493
494        Ok(proposed_block)
495    }
496
497    /// Proposes a new block from the provided batches and returns it.
498    ///
499    /// This method does not modify the chain state.
500    pub fn propose_block<I>(
501        &self,
502        batches: impl IntoIterator<Item = ProvenBatch, IntoIter = I>,
503    ) -> anyhow::Result<ProposedBlock>
504    where
505        I: Iterator<Item = ProvenBatch> + Clone,
506    {
507        // We can't access system time because we are in a no-std environment, so we use the
508        // minimally correct next timestamp.
509        let timestamp = self.latest_block_header().timestamp() + 1;
510
511        self.propose_block_at(batches, timestamp)
512    }
513
514    /// Mock-proves a proposed block into a proven block and returns it.
515    ///
516    /// This method does not modify the chain state.
517    pub fn prove_block(
518        &self,
519        proposed_block: ProposedBlock,
520    ) -> Result<ProvenBlock, ProvenBlockError> {
521        LocalBlockProver::new(0).prove_dummy(proposed_block)
522    }
523
524    // TRANSACTION APIS
525    // ----------------------------------------------------------------------------------------
526
527    /// Initializes a [`TransactionContextBuilder`] for executing against a specific block number.
528    ///
529    /// Depending on the provided `input`, the builder is initialized differently:
530    /// - [`TxContextInput::AccountId`]: Initialize the builder with [`TransactionInputs`] fetched
531    ///   from the chain for the public account identified by the ID.
532    /// - [`TxContextInput::Account`]: Initialize the builder with [`TransactionInputs`] where the
533    ///   account is passed as-is to the inputs.
534    ///
535    /// In all cases, if the chain contains authenticator for the account, they are added to the
536    /// builder.
537    ///
538    /// [`TxContextInput::Account`] can be used to build a chain of transactions against the same
539    /// account that build on top of each other. For example, transaction A modifies an account
540    /// from state 0 to 1, and transaction B modifies it from state 1 to 2.
541    pub fn build_tx_context_at(
542        &self,
543        reference_block: impl Into<BlockNumber>,
544        input: impl Into<TxContextInput>,
545        note_ids: &[NoteId],
546        unauthenticated_notes: &[Note],
547    ) -> anyhow::Result<TransactionContextBuilder> {
548        let input = input.into();
549        let reference_block = reference_block.into();
550
551        let authenticator = self.account_authenticators.get(&input.id());
552        let authenticator =
553            authenticator.and_then(|authenticator| authenticator.authenticator().cloned());
554
555        anyhow::ensure!(
556            reference_block.as_usize() < self.blocks.len(),
557            "reference block {reference_block} is out of range (latest {})",
558            self.latest_block_header().block_num()
559        );
560
561        let account = match input {
562            TxContextInput::AccountId(account_id) => {
563                if account_id.is_private() {
564                    return Err(anyhow::anyhow!(
565                        "transaction contexts for private accounts should be created with TxContextInput::Account"
566                    ));
567                }
568
569                self.committed_accounts
570                    .get(&account_id)
571                    .with_context(|| {
572                        format!("account {account_id} not found in committed accounts")
573                    })?
574                    .clone()
575            },
576            TxContextInput::Account(account) => account,
577        };
578
579        let tx_inputs = self
580            .get_transaction_inputs_at(reference_block, &account, note_ids, unauthenticated_notes)
581            .context("failed to gather transaction inputs")?;
582
583        let tx_context_builder = TransactionContextBuilder::new(account)
584            .authenticator(authenticator)
585            .tx_inputs(tx_inputs);
586
587        Ok(tx_context_builder)
588    }
589
590    /// Initializes a [`TransactionContextBuilder`] for executing against the last block header.
591    ///
592    /// This is a wrapper around [`Self::build_tx_context_at`] which uses the latest block as the
593    /// reference block. See that function's docs for details.
594    pub fn build_tx_context(
595        &self,
596        input: impl Into<TxContextInput>,
597        note_ids: &[NoteId],
598        unauthenticated_notes: &[Note],
599    ) -> anyhow::Result<TransactionContextBuilder> {
600        let reference_block = self.latest_block_header().block_num();
601        self.build_tx_context_at(reference_block, input, note_ids, unauthenticated_notes)
602    }
603
604    // INPUTS APIS
605    // ----------------------------------------------------------------------------------------
606
607    /// Returns a valid [`TransactionInputs`] for the specified entities, executing against
608    /// a specific block number.
609    pub fn get_transaction_inputs_at(
610        &self,
611        reference_block: BlockNumber,
612        account: impl Into<PartialAccount>,
613        notes: &[NoteId],
614        unauthenticated_notes: &[Note],
615    ) -> anyhow::Result<TransactionInputs> {
616        let ref_block = self.block_header(reference_block.as_usize());
617
618        let mut input_notes = vec![];
619        let mut block_headers_map: BTreeMap<BlockNumber, BlockHeader> = BTreeMap::new();
620        for note in notes {
621            let input_note: InputNote = self
622                .committed_notes
623                .get(note)
624                .with_context(|| format!("note with id {note} not found"))?
625                .clone()
626                .try_into()
627                .with_context(|| {
628                    format!("failed to convert mock chain note with id {note} into input note")
629                })?;
630
631            let note_block_num = input_note
632                .location()
633                .with_context(|| format!("note location not available: {note}"))?
634                .block_num();
635
636            if note_block_num > ref_block.block_num() {
637                anyhow::bail!(
638                    "note with ID {note} was created in block {note_block_num} which is larger than the reference block number {}",
639                    ref_block.block_num()
640                )
641            }
642
643            if note_block_num != ref_block.block_num() {
644                let block_header = self
645                    .blocks
646                    .get(note_block_num.as_usize())
647                    .with_context(|| format!("block {note_block_num} not found in chain"))?
648                    .header()
649                    .clone();
650                block_headers_map.insert(note_block_num, block_header);
651            }
652
653            input_notes.push(input_note);
654        }
655
656        for note in unauthenticated_notes {
657            input_notes.push(InputNote::Unauthenticated { note: note.clone() })
658        }
659
660        let block_headers = block_headers_map.values();
661        let (_, partial_blockchain) = self.selective_partial_blockchain(
662            reference_block,
663            block_headers.map(BlockHeader::block_num),
664        )?;
665
666        let input_notes = InputNotes::new(input_notes)?;
667
668        Ok(TransactionInputs::new(
669            account.into(),
670            ref_block.clone(),
671            partial_blockchain,
672            input_notes,
673        )?)
674    }
675
676    /// Returns a valid [`TransactionInputs`] for the specified entities.
677    pub fn get_transaction_inputs(
678        &self,
679        account: impl Into<PartialAccount>,
680        notes: &[NoteId],
681        unauthenticated_notes: &[Note],
682    ) -> anyhow::Result<TransactionInputs> {
683        let latest_block_num = self.latest_block_header().block_num();
684        self.get_transaction_inputs_at(latest_block_num, account, notes, unauthenticated_notes)
685    }
686
687    /// Returns inputs for a transaction batch for all the reference blocks of the provided
688    /// transactions.
689    pub fn get_batch_inputs(
690        &self,
691        tx_reference_blocks: impl IntoIterator<Item = BlockNumber>,
692        unauthenticated_notes: impl Iterator<Item = NoteId>,
693    ) -> anyhow::Result<(BlockHeader, PartialBlockchain, BTreeMap<NoteId, NoteInclusionProof>)>
694    {
695        // Fetch note proofs for notes that exist in the chain.
696        let unauthenticated_note_proofs = self.unauthenticated_note_proofs(unauthenticated_notes);
697
698        // We also need to fetch block inclusion proofs for any of the blocks that contain
699        // unauthenticated notes for which we want to prove inclusion.
700        let required_blocks = tx_reference_blocks.into_iter().chain(
701            unauthenticated_note_proofs
702                .values()
703                .map(|note_proof| note_proof.location().block_num()),
704        );
705
706        let (batch_reference_block, partial_block_chain) =
707            self.latest_selective_partial_blockchain(required_blocks)?;
708
709        Ok((batch_reference_block, partial_block_chain, unauthenticated_note_proofs))
710    }
711
712    /// Gets foreign account inputs to execute FPI transactions.
713    ///
714    /// Only used internally and so does not need to be public.
715    #[cfg(test)]
716    pub(crate) fn get_foreign_account_inputs(
717        &self,
718        account_id: AccountId,
719    ) -> anyhow::Result<(Account, AccountWitness)> {
720        let account = self.committed_account(account_id)?.clone();
721
722        let account_witness = self.account_tree().open(account_id);
723        assert_eq!(account_witness.state_commitment(), account.commitment());
724
725        Ok((account, account_witness))
726    }
727
728    /// Gets the inputs for a block for the provided batches.
729    pub fn get_block_inputs<'batch, I>(
730        &self,
731        batch_iter: impl IntoIterator<Item = &'batch ProvenBatch, IntoIter = I>,
732    ) -> anyhow::Result<BlockInputs>
733    where
734        I: Iterator<Item = &'batch ProvenBatch> + Clone,
735    {
736        let batch_iterator = batch_iter.into_iter();
737
738        let unauthenticated_note_proofs =
739            self.unauthenticated_note_proofs(batch_iterator.clone().flat_map(|batch| {
740                batch.input_notes().iter().filter_map(|note| note.header().map(NoteHeader::id))
741            }));
742
743        let (block_reference_block, partial_blockchain) = self
744            .latest_selective_partial_blockchain(
745                batch_iterator.clone().map(ProvenBatch::reference_block_num).chain(
746                    unauthenticated_note_proofs.values().map(|proof| proof.location().block_num()),
747                ),
748            )?;
749
750        let account_witnesses =
751            self.account_witnesses(batch_iterator.clone().flat_map(ProvenBatch::updated_accounts));
752
753        let nullifier_proofs =
754            self.nullifier_witnesses(batch_iterator.flat_map(ProvenBatch::created_nullifiers));
755
756        Ok(BlockInputs::new(
757            block_reference_block,
758            partial_blockchain,
759            account_witnesses,
760            nullifier_proofs,
761            unauthenticated_note_proofs,
762        ))
763    }
764
765    // PUBLIC MUTATORS
766    // ----------------------------------------------------------------------------------------
767
768    /// Proves the next block in the mock chain.
769    ///
770    /// This will commit all the currently pending transactions into the chain state.
771    pub fn prove_next_block(&mut self) -> anyhow::Result<ProvenBlock> {
772        self.prove_and_apply_block(None)
773    }
774
775    /// Proves the next block in the mock chain at the given timestamp.
776    ///
777    /// This will commit all the currently pending transactions into the chain state.
778    pub fn prove_next_block_at(&mut self, timestamp: u32) -> anyhow::Result<ProvenBlock> {
779        self.prove_and_apply_block(Some(timestamp))
780    }
781
782    /// Proves new blocks until the block with the given target block number has been created.
783    ///
784    /// For example, if the latest block is `5` and this function is called with `10`, then blocks
785    /// `6..=10` will be created and block 10 will be returned.
786    ///
787    /// # Panics
788    ///
789    /// Panics if:
790    /// - the given block number is smaller or equal to the number of the latest block in the chain.
791    pub fn prove_until_block(
792        &mut self,
793        target_block_num: impl Into<BlockNumber>,
794    ) -> anyhow::Result<ProvenBlock> {
795        let target_block_num = target_block_num.into();
796        let latest_block_num = self.latest_block_header().block_num();
797        assert!(
798            target_block_num > latest_block_num,
799            "target block number must be greater than the number of the latest block in the chain"
800        );
801
802        let mut last_block = None;
803        for _ in latest_block_num.as_usize()..target_block_num.as_usize() {
804            last_block = Some(self.prove_next_block()?);
805        }
806
807        Ok(last_block.expect("at least one block should have been created"))
808    }
809
810    // PUBLIC MUTATORS (PENDING APIS)
811    // ----------------------------------------------------------------------------------------
812
813    /// Adds the given [`ExecutedTransaction`] to the list of pending transactions.
814    ///
815    /// A block has to be created to apply the transaction effects to the chain state, e.g. using
816    /// [`MockChain::prove_next_block`].
817    pub fn add_pending_executed_transaction(
818        &mut self,
819        transaction: &ExecutedTransaction,
820    ) -> anyhow::Result<()> {
821        // Transform the executed tx into a proven tx with a dummy proof.
822        let proven_tx = LocalTransactionProver::default()
823            .prove_dummy(transaction.clone())
824            .context("failed to dummy-prove executed transaction into proven transaction")?;
825
826        self.pending_transactions.push(proven_tx);
827
828        Ok(())
829    }
830
831    /// Adds the given [`ProvenTransaction`] to the list of pending transactions.
832    ///
833    /// A block has to be created to apply the transaction effects to the chain state, e.g. using
834    /// [`MockChain::prove_next_block`].
835    pub fn add_pending_proven_transaction(&mut self, transaction: ProvenTransaction) {
836        self.pending_transactions.push(transaction);
837    }
838
839    // PRIVATE HELPERS
840    // ----------------------------------------------------------------------------------------
841
842    /// Applies the given block to the chain state, which means:
843    ///
844    /// - Insert account and nullifiers into the respective trees.
845    /// - Updated accounts from the block are updated in the committed accounts.
846    /// - Created notes are inserted into the committed notes.
847    /// - Consumed notes are removed from the committed notes.
848    /// - The block is appended to the [`BlockChain`] and the list of proven blocks.
849    fn apply_block(&mut self, proven_block: ProvenBlock) -> anyhow::Result<()> {
850        for account_update in proven_block.updated_accounts() {
851            self.account_tree
852                .insert(account_update.account_id(), account_update.final_state_commitment())
853                .context("failed to insert account update into account tree")?;
854        }
855
856        for nullifier in proven_block.created_nullifiers() {
857            self.nullifier_tree
858                .mark_spent(*nullifier, proven_block.header().block_num())
859                .context("failed to mark block nullifier as spent")?;
860
861            // TODO: Remove from self.committed_notes. This is not critical to have for now. It is
862            // not straightforward, because committed_notes are indexed by note IDs rather than
863            // nullifiers, so we'll have to create a second index to do this.
864        }
865
866        for account_update in proven_block.updated_accounts() {
867            match account_update.details() {
868                AccountUpdateDetails::Delta(account_delta) => {
869                    if account_delta.is_full_state() {
870                        let account = Account::try_from(account_delta)
871                            .context("failed to convert full state delta into full account")?;
872                        self.committed_accounts.insert(account.id(), account.clone());
873                    } else {
874                        let committed_account = self
875                            .committed_accounts
876                            .get_mut(&account_update.account_id())
877                            .ok_or_else(|| {
878                                anyhow::anyhow!("account delta in block for non-existent account")
879                            })?;
880                        committed_account
881                            .apply_delta(account_delta)
882                            .context("failed to apply account delta")?;
883                    }
884                },
885                // No state to keep for private accounts other than the commitment on the account
886                // tree
887                AccountUpdateDetails::Private => {},
888            }
889        }
890
891        let notes_tree = proven_block.build_output_note_tree();
892        for (block_note_index, created_note) in proven_block.output_notes() {
893            let note_path = notes_tree.open(block_note_index);
894            let note_inclusion_proof = NoteInclusionProof::new(
895                proven_block.header().block_num(),
896                block_note_index.leaf_index_value(),
897                note_path,
898            )
899            .context("failed to create inclusion proof for output note")?;
900
901            if let OutputNote::Full(note) = created_note {
902                self.committed_notes
903                    .insert(note.id(), MockChainNote::Public(note.clone(), note_inclusion_proof));
904            } else {
905                self.committed_notes.insert(
906                    created_note.id(),
907                    MockChainNote::Private(
908                        created_note.id(),
909                        *created_note.metadata(),
910                        note_inclusion_proof,
911                    ),
912                );
913            }
914        }
915
916        debug_assert_eq!(
917            self.chain.commitment(),
918            proven_block.header().chain_commitment(),
919            "current mock chain commitment and new block's chain commitment should match"
920        );
921        debug_assert_eq!(
922            BlockNumber::from(self.chain.as_mmr().forest().num_leaves() as u32),
923            proven_block.header().block_num(),
924            "current mock chain length and new block's number should match"
925        );
926
927        self.chain.push(proven_block.header().commitment());
928        self.blocks.push(proven_block);
929
930        Ok(())
931    }
932
933    fn pending_transactions_to_batches(&mut self) -> anyhow::Result<Vec<ProvenBatch>> {
934        // Batches must contain at least one transaction, so if there are no pending transactions,
935        // return early.
936        if self.pending_transactions.is_empty() {
937            return Ok(vec![]);
938        }
939
940        let pending_transactions = core::mem::take(&mut self.pending_transactions);
941
942        // TODO: Distribute the transactions into multiple batches if the transactions would not fit
943        // into a single batch (according to max input notes, max output notes and max accounts).
944        let proposed_batch = self.propose_transaction_batch(pending_transactions)?;
945        let proven_batch = self.prove_transaction_batch(proposed_batch)?;
946
947        Ok(vec![proven_batch])
948    }
949
950    /// Creates a new block in the mock chain.
951    ///
952    /// Block building is divided into two steps:
953    ///
954    /// 1. Build batches from pending transactions and a block from those batches. This results in a
955    ///    block.
956    /// 2. Insert all the account updates, nullifiers and notes from the block into the chain state.
957    ///
958    /// If a `timestamp` is provided, it will be set on the block.
959    fn prove_and_apply_block(&mut self, timestamp: Option<u32>) -> anyhow::Result<ProvenBlock> {
960        // Create batches from pending transactions.
961        // ----------------------------------------------------------------------------------------
962
963        let batches = self.pending_transactions_to_batches()?;
964
965        // Create block.
966        // ----------------------------------------------------------------------------------------
967
968        let block_timestamp =
969            timestamp.unwrap_or(self.latest_block_header().timestamp() + Self::TIMESTAMP_STEP_SECS);
970
971        let proposed_block = self
972            .propose_block_at(batches, block_timestamp)
973            .context("failed to create proposed block")?;
974        let proven_block = self.prove_block(proposed_block).context("failed to prove block")?;
975
976        // Apply block.
977        // ----------------------------------------------------------------------------------------
978
979        self.apply_block(proven_block.clone()).context("failed to apply block")?;
980
981        Ok(proven_block)
982    }
983}
984
985impl Default for MockChain {
986    fn default() -> Self {
987        MockChain::new()
988    }
989}
990
991// SERIALIZATION
992// ================================================================================================
993
994impl Serializable for MockChain {
995    fn write_into<W: ByteWriter>(&self, target: &mut W) {
996        self.chain.write_into(target);
997        self.blocks.write_into(target);
998        self.nullifier_tree.write_into(target);
999        self.account_tree.write_into(target);
1000        self.pending_transactions.write_into(target);
1001        self.committed_accounts.write_into(target);
1002        self.committed_notes.write_into(target);
1003        self.account_authenticators.write_into(target);
1004    }
1005}
1006
1007impl Deserializable for MockChain {
1008    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
1009        let chain = Blockchain::read_from(source)?;
1010        let blocks = Vec::<ProvenBlock>::read_from(source)?;
1011        let nullifier_tree = NullifierTree::read_from(source)?;
1012        let account_tree = AccountTree::read_from(source)?;
1013        let pending_transactions = Vec::<ProvenTransaction>::read_from(source)?;
1014        let committed_accounts = BTreeMap::<AccountId, Account>::read_from(source)?;
1015        let committed_notes = BTreeMap::<NoteId, MockChainNote>::read_from(source)?;
1016        let account_authenticators =
1017            BTreeMap::<AccountId, AccountAuthenticator>::read_from(source)?;
1018
1019        Ok(Self {
1020            chain,
1021            blocks,
1022            nullifier_tree,
1023            account_tree,
1024            pending_transactions,
1025            committed_notes,
1026            committed_accounts,
1027            account_authenticators,
1028        })
1029    }
1030}
1031
1032// ACCOUNT STATE
1033// ================================================================================================
1034
1035/// Helper type for increased readability at call-sites. Indicates whether to build a new (nonce =
1036/// ZERO) or existing account (nonce = ONE).
1037pub enum AccountState {
1038    New,
1039    Exists,
1040}
1041
1042// ACCOUNT AUTHENTICATOR
1043// ================================================================================================
1044
1045/// A wrapper around the authenticator of an account.
1046#[derive(Debug, Clone)]
1047pub(super) struct AccountAuthenticator {
1048    authenticator: Option<BasicAuthenticator>,
1049}
1050
1051impl AccountAuthenticator {
1052    pub fn new(authenticator: Option<BasicAuthenticator>) -> Self {
1053        Self { authenticator }
1054    }
1055
1056    pub fn authenticator(&self) -> Option<&BasicAuthenticator> {
1057        self.authenticator.as_ref()
1058    }
1059}
1060
1061impl PartialEq for AccountAuthenticator {
1062    fn eq(&self, other: &Self) -> bool {
1063        match (&self.authenticator, &other.authenticator) {
1064            (Some(a), Some(b)) => {
1065                a.keys().keys().zip(b.keys().keys()).all(|(a_key, b_key)| a_key == b_key)
1066            },
1067            (None, None) => true,
1068            _ => false,
1069        }
1070    }
1071}
1072
1073// SERIALIZATION
1074// ================================================================================================
1075
1076impl Serializable for AccountAuthenticator {
1077    fn write_into<W: ByteWriter>(&self, target: &mut W) {
1078        self.authenticator
1079            .as_ref()
1080            .map(|auth| auth.keys().values().collect::<Vec<_>>())
1081            .write_into(target);
1082    }
1083}
1084
1085impl Deserializable for AccountAuthenticator {
1086    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
1087        let authenticator = Option::<Vec<AuthSecretKey>>::read_from(source)?;
1088
1089        let authenticator = authenticator.map(|keys| BasicAuthenticator::new(&keys));
1090
1091        Ok(Self { authenticator })
1092    }
1093}
1094
1095// TX CONTEXT INPUT
1096// ================================================================================================
1097
1098/// Helper type to abstract over the inputs to [`MockChain::build_tx_context`]. See that method's
1099/// docs for details.
1100#[allow(clippy::large_enum_variant)]
1101#[derive(Debug, Clone)]
1102pub enum TxContextInput {
1103    AccountId(AccountId),
1104    Account(Account),
1105}
1106
1107impl TxContextInput {
1108    /// Returns the account ID that this input references.
1109    fn id(&self) -> AccountId {
1110        match self {
1111            TxContextInput::AccountId(account_id) => *account_id,
1112            TxContextInput::Account(account) => account.id(),
1113        }
1114    }
1115}
1116
1117impl From<AccountId> for TxContextInput {
1118    fn from(account: AccountId) -> Self {
1119        Self::AccountId(account)
1120    }
1121}
1122
1123impl From<Account> for TxContextInput {
1124    fn from(account: Account) -> Self {
1125        Self::Account(account)
1126    }
1127}
1128
1129// TESTS
1130// ================================================================================================
1131
1132#[cfg(test)]
1133mod tests {
1134    use miden_lib::account::wallets::BasicWallet;
1135    use miden_objects::account::{AccountBuilder, AccountStorageMode};
1136    use miden_objects::asset::{Asset, FungibleAsset};
1137    use miden_objects::note::NoteType;
1138    use miden_objects::testing::account_id::{
1139        ACCOUNT_ID_PRIVATE_FUNGIBLE_FAUCET,
1140        ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET,
1141        ACCOUNT_ID_SENDER,
1142    };
1143
1144    use super::*;
1145    use crate::Auth;
1146
1147    #[test]
1148    fn prove_until_block() -> anyhow::Result<()> {
1149        let mut chain = MockChain::new();
1150        let block = chain.prove_until_block(5)?;
1151        assert_eq!(block.header().block_num(), 5u32.into());
1152        assert_eq!(chain.proven_blocks().len(), 6);
1153
1154        Ok(())
1155    }
1156
1157    #[tokio::test]
1158    async fn private_account_state_update() -> anyhow::Result<()> {
1159        let faucet_id = ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET.try_into()?;
1160        let account_builder = AccountBuilder::new([4; 32])
1161            .storage_mode(AccountStorageMode::Private)
1162            .with_component(BasicWallet);
1163
1164        let mut builder = MockChain::builder();
1165        let account = builder.add_account_from_builder(
1166            Auth::BasicAuth,
1167            account_builder,
1168            AccountState::New,
1169        )?;
1170
1171        let account_id = account.id();
1172        assert_eq!(account.nonce().as_int(), 0);
1173
1174        let note_1 = builder.add_p2id_note(
1175            ACCOUNT_ID_SENDER.try_into().unwrap(),
1176            account.id(),
1177            &[Asset::Fungible(FungibleAsset::new(faucet_id, 1000u64).unwrap())],
1178            NoteType::Private,
1179        )?;
1180
1181        let mut mock_chain = builder.build()?;
1182        mock_chain.prove_next_block()?;
1183
1184        let tx = mock_chain
1185            .build_tx_context(TxContextInput::Account(account), &[], &[note_1])?
1186            .build()?
1187            .execute()
1188            .await?;
1189
1190        mock_chain.add_pending_executed_transaction(&tx)?;
1191        mock_chain.prove_next_block()?;
1192
1193        assert!(tx.final_account().nonce().as_int() > 0);
1194        assert_eq!(
1195            tx.final_account().commitment(),
1196            mock_chain.account_tree.open(account_id).state_commitment()
1197        );
1198
1199        Ok(())
1200    }
1201
1202    #[tokio::test]
1203    async fn mock_chain_serialization() {
1204        let mut builder = MockChain::builder();
1205
1206        let mut notes = vec![];
1207        for i in 0..10 {
1208            let account = builder
1209                .add_account_from_builder(
1210                    Auth::BasicAuth,
1211                    AccountBuilder::new([i; 32]).with_component(BasicWallet),
1212                    AccountState::New,
1213                )
1214                .unwrap();
1215            let note = builder
1216                .add_p2id_note(
1217                    ACCOUNT_ID_SENDER.try_into().unwrap(),
1218                    account.id(),
1219                    &[Asset::Fungible(
1220                        FungibleAsset::new(
1221                            ACCOUNT_ID_PRIVATE_FUNGIBLE_FAUCET.try_into().unwrap(),
1222                            1000u64,
1223                        )
1224                        .unwrap(),
1225                    )],
1226                    NoteType::Private,
1227                )
1228                .unwrap();
1229            notes.push((account, note));
1230        }
1231
1232        let mut chain = builder.build().unwrap();
1233        for (account, note) in notes {
1234            let tx = chain
1235                .build_tx_context(TxContextInput::Account(account), &[], &[note])
1236                .unwrap()
1237                .build()
1238                .unwrap()
1239                .execute()
1240                .await
1241                .unwrap();
1242            chain.add_pending_executed_transaction(&tx).unwrap();
1243            chain.prove_next_block().unwrap();
1244        }
1245
1246        let bytes = chain.to_bytes();
1247
1248        let deserialized = MockChain::read_from_bytes(&bytes).unwrap();
1249
1250        assert_eq!(chain.chain.as_mmr().peaks(), deserialized.chain.as_mmr().peaks());
1251        assert_eq!(chain.blocks, deserialized.blocks);
1252        assert_eq!(chain.nullifier_tree, deserialized.nullifier_tree);
1253        assert_eq!(chain.account_tree, deserialized.account_tree);
1254        assert_eq!(chain.pending_transactions, deserialized.pending_transactions);
1255        assert_eq!(chain.committed_accounts, deserialized.committed_accounts);
1256        assert_eq!(chain.committed_notes, deserialized.committed_notes);
1257        assert_eq!(chain.account_authenticators, deserialized.account_authenticators);
1258    }
1259}
miden_testing/mock_chain/chain.rs

miden_testing/mock_chain/
chain.rs
