miden_testing/mock_chain/
chain_builder.rs

1use alloc::collections::BTreeMap;
2use alloc::vec::Vec;
3
4use anyhow::Context;
5use itertools::Itertools;
6use miden_lib::account::faucets::BasicFungibleFaucet;
7use miden_lib::account::wallets::BasicWallet;
8use miden_lib::note::{create_p2id_note, create_p2ide_note, create_swap_note};
9use miden_lib::testing::account_component::MockAccountComponent;
10use miden_lib::transaction::{TransactionKernel, memory};
11use miden_objects::account::delta::AccountUpdateDetails;
12use miden_objects::account::{
13    Account,
14    AccountBuilder,
15    AccountId,
16    AccountStorageMode,
17    AccountType,
18    StorageSlot,
19};
20use miden_objects::asset::{Asset, FungibleAsset, TokenSymbol};
21use miden_objects::block::{
22    AccountTree,
23    BlockAccountUpdate,
24    BlockHeader,
25    BlockNoteTree,
26    BlockNumber,
27    Blockchain,
28    FeeParameters,
29    NullifierTree,
30    OutputNoteBatch,
31    ProvenBlock,
32};
33use miden_objects::note::{Note, NoteDetails, NoteType};
34use miden_objects::testing::account_id::ACCOUNT_ID_NATIVE_ASSET_FAUCET;
35use miden_objects::transaction::{OrderedTransactionHeaders, OutputNote};
36use miden_objects::{Felt, FieldElement, MAX_OUTPUT_NOTES_PER_BATCH, NoteError, Word, ZERO};
37use miden_processor::crypto::RpoRandomCoin;
38use rand::Rng;
39
40use crate::mock_chain::chain::AccountCredentials;
41use crate::utils::{create_p2any_note, create_spawn_note};
42use crate::{AccountState, Auth, MockChain};
43
44/// A builder for a [`MockChain`].
45#[derive(Debug, Clone)]
46pub struct MockChainBuilder {
47    accounts: BTreeMap<AccountId, Account>,
48    account_credentials: BTreeMap<AccountId, AccountCredentials>,
49    notes: Vec<OutputNote>,
50    rng: RpoRandomCoin,
51    // Fee parameters.
52    native_asset_id: AccountId,
53    verification_base_fee: u32,
54}
55
56impl MockChainBuilder {
57    // CONSTRUCTORS
58    // ----------------------------------------------------------------------------------------
59
60    /// Initializes a new mock chain builder with an empty state.
61    ///
62    /// By default, the `native_asset_id` is set to [`ACCOUNT_ID_NATIVE_ASSET_FAUCET`] and can be
63    /// overwritten using [`Self::native_asset_id`].
64    ///
65    /// The `verification_base_fee` is initialized to 0 which means no fees are required by default.
66    pub fn new() -> Self {
67        let native_asset_id =
68            ACCOUNT_ID_NATIVE_ASSET_FAUCET.try_into().expect("account ID should be valid");
69
70        Self {
71            accounts: BTreeMap::new(),
72            account_credentials: BTreeMap::new(),
73            notes: Vec::new(),
74            rng: RpoRandomCoin::new(Default::default()),
75            native_asset_id,
76            verification_base_fee: 0,
77        }
78    }
79
80    /// Initializes a new mock chain builder with the provided accounts.
81    ///
82    /// This method only adds the accounts and cannot not register any seed or authenticator for it.
83    /// Calling [`MockChain::build_tx_context`] on accounts added in this way will not work if the
84    /// account is new or if they need an authenticator.
85    ///
86    /// Due to these limitations, prefer using other methods to add accounts to the chain, e.g.
87    /// [`MockChainBuilder::add_account_from_builder`].
88    pub fn with_accounts(accounts: impl IntoIterator<Item = Account>) -> anyhow::Result<Self> {
89        let mut builder = Self::new();
90
91        for account in accounts {
92            builder.add_account(account)?;
93        }
94
95        Ok(builder)
96    }
97
98    // BUILDER METHODS
99    // ----------------------------------------------------------------------------------------
100
101    /// Sets the native asset ID of the chain.
102    ///
103    /// This must be a fungible faucet [`AccountId`] and is the asset in which fees will be accepted
104    /// by the transaction kernel.
105    pub fn native_asset_id(mut self, native_asset_id: AccountId) -> Self {
106        self.native_asset_id = native_asset_id;
107        self
108    }
109
110    /// Sets the `verification_base_fee` of the chain.
111    ///
112    /// See [`FeeParameters`] for more details.
113    pub fn verification_base_fee(mut self, verification_base_fee: u32) -> Self {
114        self.verification_base_fee = verification_base_fee;
115        self
116    }
117
118    /// Consumes the builder, creates the genesis block of the chain and returns the [`MockChain`].
119    pub fn build(self) -> anyhow::Result<MockChain> {
120        // Create the genesis block, consisting of the provided accounts and notes.
121        let block_account_updates: Vec<BlockAccountUpdate> = self
122            .accounts
123            .into_values()
124            .map(|account| {
125                BlockAccountUpdate::new(
126                    account.id(),
127                    account.commitment(),
128                    AccountUpdateDetails::New(account),
129                )
130            })
131            .collect();
132
133        let account_tree = AccountTree::with_entries(
134            block_account_updates
135                .iter()
136                .map(|account| (account.account_id(), account.final_state_commitment())),
137        )
138        .context("failed to create genesis account tree")?;
139
140        let note_chunks = self.notes.into_iter().chunks(MAX_OUTPUT_NOTES_PER_BATCH);
141        let output_note_batches: Vec<OutputNoteBatch> = note_chunks
142            .into_iter()
143            .map(|batch_notes| batch_notes.into_iter().enumerate().collect::<Vec<_>>())
144            .collect();
145
146        let created_nullifiers = Vec::new();
147        let transactions = OrderedTransactionHeaders::new_unchecked(Vec::new());
148
149        let note_tree = BlockNoteTree::from_note_batches(&output_note_batches)
150            .context("failed to create block note tree")?;
151
152        let version = 0;
153        let prev_block_commitment = Word::empty();
154        let block_num = BlockNumber::from(0u32);
155        let chain_commitment = Blockchain::new().commitment();
156        let account_root = account_tree.root();
157        let nullifier_root = NullifierTree::new().root();
158        let note_root = note_tree.root();
159        let tx_commitment = transactions.commitment();
160        let tx_kernel_commitment = TransactionKernel::kernel_commitment();
161        let proof_commitment = Word::empty();
162        let timestamp = MockChain::TIMESTAMP_START_SECS;
163        let fee_parameters = FeeParameters::new(self.native_asset_id, self.verification_base_fee)
164            .context("failed to construct fee parameters")?;
165
166        let header = BlockHeader::new(
167            version,
168            prev_block_commitment,
169            block_num,
170            chain_commitment,
171            account_root,
172            nullifier_root,
173            note_root,
174            tx_commitment,
175            tx_kernel_commitment,
176            proof_commitment,
177            fee_parameters,
178            timestamp,
179        );
180
181        let genesis_block = ProvenBlock::new_unchecked(
182            header,
183            block_account_updates,
184            output_note_batches,
185            created_nullifiers,
186            transactions,
187        );
188
189        MockChain::from_genesis_block(genesis_block, account_tree, self.account_credentials)
190    }
191
192    // ACCOUNT METHODS
193    // ----------------------------------------------------------------------------------------
194
195    /// Creates a new public [`BasicWallet`] account and registers the authenticator (if any) and
196    /// seed.
197    ///
198    /// This does not add the account to the chain state, but it can still be used to call
199    /// [`MockChain::build_tx_context`] to automatically handle the authenticator and seed.
200    pub fn create_new_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
201        let account_builder = AccountBuilder::new(self.rng.random())
202            .storage_mode(AccountStorageMode::Public)
203            .with_component(BasicWallet);
204
205        self.add_account_from_builder(auth_method, account_builder, AccountState::New)
206    }
207
208    /// Adds an existing public [`BasicWallet`] account to the initial chain state and registers the
209    /// authenticator (if any).
210    pub fn add_existing_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
211        self.add_existing_wallet_with_assets(auth_method, [])
212    }
213
214    /// Adds an existing public [`BasicWallet`] account to the initial chain state and registers the
215    /// authenticator (if any).
216    pub fn add_existing_wallet_with_assets(
217        &mut self,
218        auth_method: Auth,
219        assets: impl IntoIterator<Item = Asset>,
220    ) -> anyhow::Result<Account> {
221        let account_builder = Account::builder(self.rng.random())
222            .storage_mode(AccountStorageMode::Public)
223            .with_component(BasicWallet)
224            .with_assets(assets);
225
226        self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
227    }
228
229    /// Creates a new public [`BasicFungibleFaucet`] account and registers the authenticator (if
230    /// any) and seed.
231    ///
232    /// This does not add the account to the chain state, but it can still be used to call
233    /// [`MockChain::build_tx_context`] to automatically handle the authenticator and seed.
234    pub fn create_new_faucet(
235        &mut self,
236        auth_method: Auth,
237        token_symbol: &str,
238        max_supply: u64,
239    ) -> anyhow::Result<Account> {
240        let token_symbol = TokenSymbol::new(token_symbol)
241            .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
242        let max_supply_felt = max_supply.try_into().map_err(|_| {
243            anyhow::anyhow!("max supply value cannot be converted to Felt: {max_supply}")
244        })?;
245        let basic_faucet = BasicFungibleFaucet::new(token_symbol, 10, max_supply_felt)
246            .context("failed to create BasicFungibleFaucet")?;
247
248        let account_builder = AccountBuilder::new(self.rng.random())
249            .storage_mode(AccountStorageMode::Public)
250            .account_type(AccountType::FungibleFaucet)
251            .with_component(basic_faucet);
252
253        self.add_account_from_builder(auth_method, account_builder, AccountState::New)
254    }
255
256    /// Adds an existing public [`BasicFungibleFaucet`] account to the initial chain state and
257    /// registers the authenticator (if the given [`Auth`] results in the creation of one).
258    pub fn add_existing_faucet(
259        &mut self,
260        auth_method: Auth,
261        token_symbol: &str,
262        max_supply: u64,
263        total_issuance: Option<u64>,
264    ) -> anyhow::Result<Account> {
265        let token_symbol = TokenSymbol::new(token_symbol).context("invalid argument")?;
266        let basic_faucet = BasicFungibleFaucet::new(token_symbol, 10u8, Felt::new(max_supply))
267            .context("invalid argument")?;
268
269        let account_builder = AccountBuilder::new(self.rng.random())
270            .storage_mode(AccountStorageMode::Public)
271            .with_component(basic_faucet)
272            .account_type(AccountType::FungibleFaucet);
273
274        let mut account =
275            self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)?;
276
277        // The faucet's reserved slot is initialized to an empty word by default.
278        // If total_issuance is set, overwrite it and reinsert the account.
279        if let Some(issuance) = total_issuance {
280            account
281                .storage_mut()
282                .set_item(
283                    memory::FAUCET_STORAGE_DATA_SLOT,
284                    Word::from([ZERO, ZERO, ZERO, Felt::new(issuance)]),
285                )
286                .context("failed to set faucet storage")?;
287            self.accounts.insert(account.id(), account.clone());
288        }
289
290        Ok(account)
291    }
292
293    /// Creates a new public account with an [`MockAccountComponent`] and registers the
294    /// authenticator (if any).
295    pub fn create_new_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
296        let account_builder = Account::builder(self.rng.random())
297            .storage_mode(AccountStorageMode::Public)
298            .with_component(MockAccountComponent::with_empty_slots());
299
300        self.add_account_from_builder(auth_method, account_builder, AccountState::New)
301    }
302
303    /// Adds an existing public account with an [`MockAccountComponent`] to the initial chain state
304    /// and registers the authenticator (if any).
305    pub fn add_existing_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
306        self.add_existing_mock_account_with_storage_and_assets(auth_method, [], [])
307    }
308
309    /// Adds an existing public account with an [`MockAccountComponent`] to the initial chain state
310    /// and registers the authenticator (if any).
311    pub fn add_existing_mock_account_with_storage(
312        &mut self,
313        auth_method: Auth,
314        slots: impl IntoIterator<Item = StorageSlot>,
315    ) -> anyhow::Result<Account> {
316        self.add_existing_mock_account_with_storage_and_assets(auth_method, slots, [])
317    }
318
319    /// Adds an existing public account with an [`MockAccountComponent`] to the initial chain state
320    /// and registers the authenticator (if any).
321    pub fn add_existing_mock_account_with_assets(
322        &mut self,
323        auth_method: Auth,
324        assets: impl IntoIterator<Item = Asset>,
325    ) -> anyhow::Result<Account> {
326        self.add_existing_mock_account_with_storage_and_assets(auth_method, [], assets)
327    }
328
329    /// Adds an existing public account with an [`MockAccountComponent`] to the initial chain state
330    /// and registers the authenticator (if any).
331    pub fn add_existing_mock_account_with_storage_and_assets(
332        &mut self,
333        auth_method: Auth,
334        slots: impl IntoIterator<Item = StorageSlot>,
335        assets: impl IntoIterator<Item = Asset>,
336    ) -> anyhow::Result<Account> {
337        let account_builder = Account::builder(self.rng.random())
338            .storage_mode(AccountStorageMode::Public)
339            .with_component(MockAccountComponent::with_slots(slots.into_iter().collect()))
340            .with_assets(assets);
341
342        self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
343    }
344
345    /// Builds the provided [`AccountBuilder`] with the provided auth method and registers the
346    /// authenticator (if any).
347    ///
348    /// - If [`AccountState::Exists`] is given the account is built as an existing account and added
349    ///   to the initial chain state. It can then be used in a transaction without having to
350    ///   validate its seed.
351    /// - If [`AccountState::New`] is given the account is built as a new account and is **not**
352    ///   added to the chain. Its seed and authenticator are registered (if any). Its first
353    ///   transaction will be its creation transaction. [`MockChain::build_tx_context`] can be
354    ///   called with the account to automatically handle the authenticator and seed.
355    pub fn add_account_from_builder(
356        &mut self,
357        auth_method: Auth,
358        mut account_builder: AccountBuilder,
359        account_state: AccountState,
360    ) -> anyhow::Result<Account> {
361        let (auth_component, authenticator) = auth_method.build_component();
362        account_builder = account_builder.with_auth_component(auth_component);
363
364        let (account, seed) = if let AccountState::New = account_state {
365            let (account, seed) =
366                account_builder.build().context("failed to build account from builder")?;
367            (account, Some(seed))
368        } else {
369            let account = account_builder
370                .build_existing()
371                .context("failed to build account from builder")?;
372            (account, None)
373        };
374
375        self.account_credentials
376            .insert(account.id(), AccountCredentials::new(seed, authenticator));
377
378        if let AccountState::Exists = account_state {
379            self.accounts.insert(account.id(), account.clone());
380        }
381
382        Ok(account)
383    }
384
385    /// Adds the provided account to the list of genesis accounts.
386    ///
387    /// This method only adds the account and does not store its account credentials (seed and
388    /// authenticator) for it. Calling [`MockChain::build_tx_context`] on accounts added in this
389    /// way will not work if the account is new or if they need an authenticator.
390    ///
391    /// Due to these limitations, prefer using other methods to add accounts to the chain, e.g.
392    /// [`MockChainBuilder::add_account_from_builder`].
393    pub fn add_account(&mut self, account: Account) -> anyhow::Result<()> {
394        self.accounts.insert(account.id(), account);
395
396        // This returns a Result to be conservative in case we need to return an error in the future
397        // and do not want to break this API.
398        Ok(())
399    }
400
401    // NOTE METHODS
402    // ----------------------------------------------------------------------------------------
403
404    /// Adds the provided note to the initial chain state.
405    pub fn add_note(&mut self, note: impl Into<OutputNote>) {
406        self.notes.push(note.into());
407    }
408
409    /// Creates a new P2ANY note from the provided parameters and adds it to the list of genesis
410    /// notes. This note is similar to a P2ID note but can be consumed by any account.
411    ///
412    /// In the created [`MockChain`], the note will be immediately spendable by `target_account_id`
413    /// and carries no additional reclaim or timelock conditions.
414    pub fn add_p2any_note(
415        &mut self,
416        sender_account_id: AccountId,
417        asset: &[Asset],
418    ) -> anyhow::Result<Note> {
419        let note = create_p2any_note(sender_account_id, asset);
420
421        self.add_note(OutputNote::Full(note.clone()));
422
423        Ok(note)
424    }
425
426    /// Creates a new P2ID note from the provided parameters and adds it to the list of genesis
427    /// notes.
428    ///
429    /// In the created [`MockChain`], the note will be immediately spendable by `target_account_id`
430    /// and carries no additional reclaim or timelock conditions.
431    pub fn add_p2id_note(
432        &mut self,
433        sender_account_id: AccountId,
434        target_account_id: AccountId,
435        asset: &[Asset],
436        note_type: NoteType,
437    ) -> Result<Note, NoteError> {
438        let note = create_p2id_note(
439            sender_account_id,
440            target_account_id,
441            asset.to_vec(),
442            note_type,
443            Default::default(),
444            &mut self.rng,
445        )?;
446
447        self.add_note(OutputNote::Full(note.clone()));
448
449        Ok(note)
450    }
451
452    /// Adds a P2IDE [`OutputNote`] (pay‑to‑ID‑extended) to the list of genesis notes.
453    ///
454    /// A P2IDE note can include an optional `timelock_height` and/or an optional
455    /// `reclaim_height` after which the `sender_account_id` may reclaim the
456    /// funds.
457    pub fn add_p2ide_note(
458        &mut self,
459        sender_account_id: AccountId,
460        target_account_id: AccountId,
461        asset: &[Asset],
462        note_type: NoteType,
463        reclaim_height: Option<BlockNumber>,
464        timelock_height: Option<BlockNumber>,
465    ) -> Result<Note, NoteError> {
466        let note = create_p2ide_note(
467            sender_account_id,
468            target_account_id,
469            asset.to_vec(),
470            reclaim_height,
471            timelock_height,
472            note_type,
473            Default::default(),
474            &mut self.rng,
475        )?;
476
477        self.add_note(OutputNote::Full(note.clone()));
478
479        Ok(note)
480    }
481
482    /// Adds a public SWAP [`OutputNote`] to the list of genesis notes.
483    pub fn add_swap_note(
484        &mut self,
485        sender: AccountId,
486        offered_asset: Asset,
487        requested_asset: Asset,
488        payback_note_type: NoteType,
489    ) -> anyhow::Result<(Note, NoteDetails)> {
490        let (swap_note, payback_note) = create_swap_note(
491            sender,
492            offered_asset,
493            requested_asset,
494            NoteType::Public,
495            Felt::ZERO,
496            payback_note_type,
497            Felt::ZERO,
498            &mut self.rng,
499        )?;
500
501        self.add_note(OutputNote::Full(swap_note.clone()));
502
503        Ok((swap_note, payback_note))
504    }
505
506    /// Adds a public `SPAWN` note to the list of genesis notes.
507    ///
508    /// A `SPAWN` note contains a note script that creates all `output_notes` that get passed as a
509    /// parameter.
510    pub fn add_spawn_note<'note>(
511        &mut self,
512        sender_id: AccountId,
513        output_notes: impl IntoIterator<Item = &'note Note>,
514    ) -> anyhow::Result<Note> {
515        let output_notes = output_notes.into_iter().collect();
516        let note = create_spawn_note(sender_id, output_notes)?;
517
518        self.add_note(OutputNote::Full(note.clone()));
519
520        Ok(note)
521    }
522
523    /// Creates a new P2ID note with the provided amount of the native fee asset of the chain.
524    ///
525    /// The native asset ID of the asset can be set using [`Self::native_asset_id`]. By default it
526    /// is [`ACCOUNT_ID_NATIVE_ASSET_FAUCET`].
527    ///
528    /// In the created [`MockChain`], the note will be immediately spendable by `target_account_id`.
529    pub fn add_p2id_note_with_fee(
530        &mut self,
531        target_account_id: AccountId,
532        amount: u64,
533    ) -> anyhow::Result<Note> {
534        let fee_asset = self.native_fee_asset(amount)?;
535        let note = self.add_p2id_note(
536            self.native_asset_id,
537            target_account_id,
538            &[Asset::from(fee_asset)],
539            NoteType::Public,
540        )?;
541
542        Ok(note)
543    }
544
545    // HELPER FUNCTIONS
546    // ----------------------------------------------------------------------------------------
547
548    /// Constructs a fungible asset based on the native asset ID and the provided amount.
549    fn native_fee_asset(&self, amount: u64) -> anyhow::Result<FungibleAsset> {
550        FungibleAsset::new(self.native_asset_id, amount).context("failed to create fee asset")
551    }
552}
553
554impl Default for MockChainBuilder {
555    fn default() -> Self {
556        Self::new()
557    }
558}
miden_testing/mock_chain/chain_builder.rs

miden_testing/mock_chain/
chain_builder.rs
