miden_testing/mock_chain/
chain.rs

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

miden_testing/mock_chain/
chain.rs
