Skip to main content

miden_testing/mock_chain/
chain_builder.rs

1use alloc::collections::{BTreeMap, BTreeSet};
2use alloc::vec::Vec;
3
4use anyhow::Context;
5
6// CONSTANTS
7// ================================================================================================
8
9/// Default number of decimals for faucets created in tests.
10const DEFAULT_FAUCET_DECIMALS: u8 = 10;
11
12// IMPORTS
13// ================================================================================================
14
15use itertools::Itertools;
16use miden_processor::crypto::random::RandomCoin;
17use miden_protocol::account::delta::AccountUpdateDetails;
18use miden_protocol::account::{
19    Account,
20    AccountBuilder,
21    AccountComponent,
22    AccountDelta,
23    AccountId,
24    AccountType,
25    StorageSlot,
26};
27use miden_protocol::asset::{Asset, AssetAmount, FungibleAsset, TokenSymbol};
28use miden_protocol::block::account_tree::AccountTree;
29use miden_protocol::block::nullifier_tree::NullifierTree;
30use miden_protocol::block::{
31    BlockAccountUpdate,
32    BlockBody,
33    BlockHeader,
34    BlockNoteTree,
35    BlockNumber,
36    BlockProof,
37    Blockchain,
38    FeeParameters,
39    OutputNoteBatch,
40    ProvenBlock,
41};
42use miden_protocol::crypto::merkle::smt::Smt;
43use miden_protocol::errors::NoteError;
44use miden_protocol::note::{Note, NoteAttachments, NoteDetails, NoteScriptRoot, NoteType};
45use miden_protocol::testing::account_id::ACCOUNT_ID_FEE_FAUCET;
46use miden_protocol::testing::random_secret_key::random_secret_key;
47use miden_protocol::transaction::{OrderedTransactionHeaders, RawOutputNote, TransactionKernel};
48use miden_protocol::{MAX_OUTPUT_NOTES_PER_BATCH, Word};
49use miden_standards::account::access::AccessControl;
50use miden_standards::account::faucets::{FungibleFaucet, TokenName};
51use miden_standards::account::policies::{
52    BurnPolicyConfig,
53    MintPolicyConfig,
54    PolicyRegistration,
55    TokenPolicyManager,
56    TransferPolicy,
57};
58use miden_standards::account::wallets::BasicWallet;
59use miden_standards::note::{BurnNote, MintNote, P2idNote, P2ideNote, P2ideNoteStorage, SwapNote};
60use miden_standards::testing::account_component::MockAccountComponent;
61use rand::Rng;
62
63use crate::mock_chain::chain::AccountAuthenticator;
64use crate::utils::{create_p2any_note, create_spawn_note};
65use crate::{AccountState, Auth, MockChain};
66
67/// A builder for a [`MockChain`]'s genesis block.
68///
69/// ## Example
70///
71/// ```
72/// # use anyhow::Result;
73/// # use miden_protocol::{
74/// #    asset::{Asset, FungibleAsset},
75/// #    note::NoteType,
76/// # };
77/// # use miden_testing::{Auth, MockChain};
78/// #
79/// # fn main() -> Result<()> {
80/// let mut builder = MockChain::builder();
81/// let existing_wallet =
82///     builder.add_existing_wallet_with_assets(Auth::IncrNonce, [FungibleAsset::mock(500)])?;
83/// let new_wallet = builder.create_new_wallet(Auth::IncrNonce)?;
84///
85/// let existing_note = builder.add_p2id_note(
86///     existing_wallet.id(),
87///     new_wallet.id(),
88///     &[FungibleAsset::mock(100)],
89///     NoteType::Private,
90/// )?;
91/// let chain = builder.build()?;
92///
93/// // The existing wallet and note should be part of the chain state.
94/// assert!(chain.committed_account(existing_wallet.id()).is_ok());
95/// assert!(chain.committed_notes().get(&existing_note.id()).is_some());
96///
97/// // The new wallet should *not* be part of the chain state - it must be created in
98/// // a transaction first.
99/// assert!(chain.committed_account(new_wallet.id()).is_err());
100///
101/// # Ok(())
102/// # }
103/// ```
104///
105/// Note the distinction between `add_` and `create_` APIs. Any `add_` APIs will add something to
106/// the genesis chain state while `create_` APIs do not mutate the genesis state. The latter are
107/// simply convenient for creating accounts or notes that will be created by transactions.
108///
109/// See also the [`MockChain`] docs for examples on using the mock chain.
110#[derive(Debug, Clone)]
111pub struct MockChainBuilder {
112    accounts: BTreeMap<AccountId, Account>,
113    account_authenticators: BTreeMap<AccountId, AccountAuthenticator>,
114    notes: Vec<RawOutputNote>,
115    rng: RandomCoin,
116    // Fee parameters.
117    fee_faucet_id: AccountId,
118    verification_base_fee: u32,
119}
120
121impl MockChainBuilder {
122    // CONSTRUCTORS
123    // ----------------------------------------------------------------------------------------
124
125    /// Initializes a new mock chain builder with an empty state.
126    ///
127    /// By default, the `fee_faucet_id` is set to [`ACCOUNT_ID_FEE_FAUCET`] and can be
128    /// overwritten using [`Self::fee_faucet_id`].
129    ///
130    /// The `verification_base_fee` is initialized to 0 which means no fees are required by default.
131    pub fn new() -> Self {
132        let fee_faucet_id = ACCOUNT_ID_FEE_FAUCET.try_into().expect("account ID should be valid");
133
134        Self {
135            accounts: BTreeMap::new(),
136            account_authenticators: BTreeMap::new(),
137            notes: Vec::new(),
138            rng: RandomCoin::new(Default::default()),
139            fee_faucet_id,
140            verification_base_fee: 0,
141        }
142    }
143
144    /// Initializes a new mock chain builder with the provided accounts.
145    ///
146    /// This method only adds the accounts and cannot not register any authenticators for them.
147    /// Calling [`MockChain::build_tx_context`] on accounts added in this way will not work if the
148    /// account needs an authenticator.
149    ///
150    /// Due to these limitations, prefer using other methods to add accounts to the chain, e.g.
151    /// [`MockChainBuilder::add_account_from_builder`].
152    pub fn with_accounts(accounts: impl IntoIterator<Item = Account>) -> anyhow::Result<Self> {
153        let mut builder = Self::new();
154
155        for account in accounts {
156            builder.add_account(account)?;
157        }
158
159        Ok(builder)
160    }
161
162    // BUILDER METHODS
163    // ----------------------------------------------------------------------------------------
164
165    /// Sets the fee faucet ID of the chain.
166    ///
167    /// This must be a fungible faucet [`AccountId`] and is the asset in which fees will be accepted
168    /// by the transaction kernel.
169    pub fn fee_faucet_id(mut self, fee_faucet_id: AccountId) -> Self {
170        self.fee_faucet_id = fee_faucet_id;
171        self
172    }
173
174    /// Sets the `verification_base_fee` of the chain.
175    ///
176    /// See [`FeeParameters`] for more details.
177    pub fn verification_base_fee(mut self, verification_base_fee: u32) -> Self {
178        self.verification_base_fee = verification_base_fee;
179        self
180    }
181
182    /// Consumes the builder, creates the genesis block of the chain and returns the [`MockChain`].
183    pub fn build(self) -> anyhow::Result<MockChain> {
184        // Create the genesis block, consisting of the provided accounts and notes.
185        let block_account_updates: Vec<BlockAccountUpdate> = self
186            .accounts
187            .into_values()
188            .map(|account| {
189                let account_id = account.id();
190                let account_commitment = account.to_commitment();
191                let account_delta = AccountDelta::try_from(account)
192                    .expect("chain builder should only store existing accounts without seeds");
193                let update_details = AccountUpdateDetails::Delta(account_delta);
194
195                BlockAccountUpdate::new(account_id, account_commitment, update_details)
196            })
197            .collect();
198
199        let account_tree = AccountTree::with_entries(
200            block_account_updates
201                .iter()
202                .map(|account| (account.account_id(), account.final_state_commitment())),
203        )
204        .context("failed to create genesis account tree")?;
205
206        // Extract full notes before shrinking for later use in MockChain
207        let full_notes: Vec<Note> = self
208            .notes
209            .iter()
210            .filter_map(|note| match note {
211                RawOutputNote::Full(n) => Some(n.clone()),
212                _ => None,
213            })
214            .collect();
215
216        let proven_notes: Vec<_> = self
217            .notes
218            .into_iter()
219            .map(|note| note.into_output_note().expect("genesis note should be valid"))
220            .collect();
221        let note_chunks = proven_notes.into_iter().chunks(MAX_OUTPUT_NOTES_PER_BATCH);
222        let output_note_batches: Vec<OutputNoteBatch> = note_chunks
223            .into_iter()
224            .map(|batch_notes| batch_notes.into_iter().enumerate().collect::<Vec<_>>())
225            .collect();
226
227        let created_nullifiers = Vec::new();
228        let transactions = OrderedTransactionHeaders::new_unchecked(Vec::new());
229
230        let note_tree = BlockNoteTree::from_note_batches(&output_note_batches)
231            .context("failed to create block note tree")?;
232
233        let version = 0;
234        let prev_block_commitment = Word::empty();
235        let block_num = BlockNumber::from(0u32);
236        let chain_commitment = Blockchain::new().commitment();
237        let account_root = account_tree.root();
238        let nullifier_root = NullifierTree::<Smt>::default().root();
239        let note_root = note_tree.root();
240        let tx_commitment = transactions.commitment();
241        let tx_kernel_commitment = TransactionKernel.to_commitment();
242        let timestamp = MockChain::TIMESTAMP_START_SECS;
243        let fee_parameters = FeeParameters::new(self.fee_faucet_id, self.verification_base_fee);
244        let validator_secret_key = random_secret_key();
245        let validator_public_key = validator_secret_key.public_key();
246
247        let header = BlockHeader::new(
248            version,
249            prev_block_commitment,
250            block_num,
251            chain_commitment,
252            account_root,
253            nullifier_root,
254            note_root,
255            tx_commitment,
256            tx_kernel_commitment,
257            validator_public_key,
258            fee_parameters,
259            timestamp,
260        );
261
262        let body = BlockBody::new_unchecked(
263            block_account_updates,
264            output_note_batches,
265            created_nullifiers,
266            transactions,
267        );
268
269        let signature = validator_secret_key.sign(header.commitment());
270        let block_proof = BlockProof::new_dummy();
271        let genesis_block = ProvenBlock::new_unchecked(header, body, signature, block_proof);
272
273        MockChain::from_genesis_block(
274            genesis_block,
275            account_tree,
276            self.account_authenticators,
277            validator_secret_key,
278            full_notes,
279        )
280    }
281
282    // ACCOUNT METHODS
283    // ----------------------------------------------------------------------------------------
284
285    /// Creates a new public [`BasicWallet`] account and registers the authenticator (if any) for
286    /// it.
287    ///
288    /// This does not add the account to the chain state, but it can still be used to call
289    /// [`MockChain::build_tx_context`] to automatically add the authenticator.
290    pub fn create_new_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
291        let account_builder = AccountBuilder::new(self.rng.random())
292            .account_type(AccountType::Public)
293            .with_component(BasicWallet);
294
295        self.add_account_from_builder(auth_method, account_builder, AccountState::New)
296    }
297
298    /// Adds an existing public [`BasicWallet`] account to the initial chain state and registers the
299    /// authenticator (if any).
300    pub fn add_existing_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
301        self.add_existing_wallet_with_assets(auth_method, [])
302    }
303
304    /// Adds an existing public [`BasicWallet`] account to the initial chain state and registers the
305    /// authenticator (if any).
306    pub fn add_existing_wallet_with_assets(
307        &mut self,
308        auth_method: Auth,
309        assets: impl IntoIterator<Item = Asset>,
310    ) -> anyhow::Result<Account> {
311        let account_builder = Account::builder(self.rng.random())
312            .account_type(AccountType::Public)
313            .with_component(BasicWallet)
314            .with_assets(assets);
315
316        self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
317    }
318
319    /// Creates a new public [`FungibleFaucet`] account and registers the authenticator (if
320    /// any) for it.
321    ///
322    /// This does not add the account to the chain state, but it can still be used to call
323    /// [`MockChain::build_tx_context`] to automatically add the authenticator.
324    fn create_new_fungible_faucet(
325        &mut self,
326        auth_method: Auth,
327        faucet: FungibleFaucet,
328        account_type: AccountType,
329        access_control: AccessControl,
330        token_policy_manager: TokenPolicyManager,
331    ) -> anyhow::Result<Account> {
332        let account_builder = AccountBuilder::new(self.rng.random())
333            .account_type(account_type)
334            .with_component(faucet)
335            .with_components(access_control)
336            .with_components(token_policy_manager);
337
338        self.add_account_from_builder(auth_method, account_builder, AccountState::New)
339    }
340
341    /// Adds an existing fungible faucet account to the initial chain state and registers the
342    /// authenticator (if any).
343    ///
344    /// The behaviour of the faucet (basic vs network-style) is determined entirely by the
345    /// combination of arguments:
346    /// - `account_type`: [`AccountType::Public`] for basic faucets, or [`AccountType::Private`] for
347    ///   off-chain accounts.
348    /// - `auth_method`: typically a [`Auth::BasicAuth`] for basic faucets, or [`Auth::IncrNonce`]
349    ///   for network-style faucets.
350    /// - `access_control`: [`AccessControl::AuthControlled`] for basic faucets;
351    ///   [`AccessControl::Ownable2Step`] / [`AccessControl::Rbac`] for owner-controlled faucets.
352    ///   The matching `Authority` component is auto-installed by `AccessControl`.
353    /// - `token_policy_manager`: the unified [`TokenPolicyManager`] holding both mint and burn
354    ///   policy.
355    fn add_existing_fungible_faucet(
356        &mut self,
357        auth_method: Auth,
358        faucet: FungibleFaucet,
359        account_type: AccountType,
360        access_control: AccessControl,
361        token_policy_manager: TokenPolicyManager,
362    ) -> anyhow::Result<Account> {
363        let account_builder = AccountBuilder::new(self.rng.random())
364            .account_type(account_type)
365            .with_component(faucet)
366            .with_components(access_control)
367            .with_components(token_policy_manager);
368
369        self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
370    }
371
372    /// Convenience: builds a basic auth-controlled fungible faucet from a token-symbol shorthand
373    /// using default decimals and `AllowAll` policies, then adds it via
374    /// `Self::add_existing_fungible_faucet`.
375    ///
376    /// For full control over the faucet's metadata, decimals, and policies, construct a
377    /// [`FungibleFaucet`] manually and call `Self::add_existing_fungible_faucet`.
378    pub fn add_existing_basic_faucet(
379        &mut self,
380        auth_method: Auth,
381        token_symbol: &str,
382        max_supply: u64,
383        token_supply: Option<u64>,
384    ) -> anyhow::Result<Account> {
385        let token_supply = token_supply.unwrap_or(0);
386        let name = TokenName::new(token_symbol)?;
387        let symbol = TokenSymbol::new(token_symbol)
388            .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
389        let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
390        let token_supply = AssetAmount::new(token_supply).context("invalid token_supply")?;
391        let faucet = FungibleFaucet::builder()
392            .name(name)
393            .symbol(symbol)
394            .decimals(DEFAULT_FAUCET_DECIMALS)
395            .max_supply(max_supply)
396            .token_supply(token_supply)
397            .build()
398            .context("failed to build FungibleFaucet")?;
399
400        let token_policy_manager = TokenPolicyManager::new()
401            .with_mint_policy(MintPolicyConfig::AllowAll, PolicyRegistration::Active)?
402            .with_burn_policy(BurnPolicyConfig::AllowAll, PolicyRegistration::Active)?
403            .with_send_policy(TransferPolicy::AllowAll, PolicyRegistration::Active)?
404            .with_receive_policy(TransferPolicy::AllowAll, PolicyRegistration::Active)?;
405
406        self.add_existing_fungible_faucet(
407            auth_method,
408            faucet,
409            AccountType::Public,
410            AccessControl::AuthControlled,
411            token_policy_manager,
412        )
413    }
414
415    /// Convenience: builds an owner-controlled (network-style) fungible faucet from a
416    /// token-symbol shorthand using default decimals, the given `mint_policy`, and `BurnAllowAll`.
417    ///
418    /// The faucet is added with [`AccountType::Public`] and [`Auth::IncrNonce`].
419    ///
420    /// `mint_policy` selects the initial active mint policy on the faucet. The installed
421    /// [`TokenPolicyManager`] is always owner-controlled.
422    ///
423    /// The [`MintNote`] and [`BurnNote`] script roots are always added to `allowed_script_roots`,
424    /// so callers only need to provide any additional roots their test scripts require.
425    pub fn add_existing_network_faucet(
426        &mut self,
427        token_symbol: &str,
428        max_supply: u64,
429        owner_account_id: AccountId,
430        token_supply: Option<u64>,
431        mint_policy: MintPolicyConfig,
432        allowed_script_roots: impl IntoIterator<Item = NoteScriptRoot>,
433    ) -> anyhow::Result<Account> {
434        let token_supply = token_supply.unwrap_or(0);
435        let name = TokenName::new(token_symbol)?;
436        let symbol = TokenSymbol::new(token_symbol)
437            .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
438        let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
439        let token_supply = AssetAmount::new(token_supply).context("invalid token_supply")?;
440        let faucet = FungibleFaucet::builder()
441            .name(name)
442            .symbol(symbol)
443            .decimals(DEFAULT_FAUCET_DECIMALS)
444            .max_supply(max_supply)
445            .token_supply(token_supply)
446            .build()
447            .context("failed to build FungibleFaucet")?;
448
449        let token_policy_manager = TokenPolicyManager::new()
450            .with_mint_policy(mint_policy, PolicyRegistration::Active)?
451            .with_burn_policy(BurnPolicyConfig::AllowAll, PolicyRegistration::Active)?
452            .with_send_policy(TransferPolicy::AllowAll, PolicyRegistration::Active)?
453            .with_receive_policy(TransferPolicy::AllowAll, PolicyRegistration::Active)?;
454
455        let allowed_script_roots = allowed_script_roots
456            .into_iter()
457            .chain([MintNote::script_root(), BurnNote::script_root()])
458            .collect();
459
460        self.add_existing_fungible_faucet(
461            Auth::NetworkAccount {
462                allowed_script_roots,
463                allowed_tx_script_roots: BTreeSet::new(),
464            },
465            faucet,
466            AccountType::Public,
467            AccessControl::Ownable2Step { owner: owner_account_id },
468            token_policy_manager,
469        )
470    }
471
472    /// Convenience: adds an existing owner-controlled (network-style) fungible faucet whose token
473    /// metadata is fully provided by the caller. Uses `OwnerOnly` mint policy and `AllowAll`
474    /// burn policy by default.
475    ///
476    /// The [`MintNote`] and [`BurnNote`] script roots are always added to `allowed_script_roots`,
477    /// so callers only need to provide any additional roots their test scripts require.
478    pub fn add_existing_network_faucet_with_metadata(
479        &mut self,
480        owner_account_id: AccountId,
481        faucet: FungibleFaucet,
482        allowed_script_roots: impl IntoIterator<Item = NoteScriptRoot>,
483    ) -> anyhow::Result<Account> {
484        let token_policy_manager = TokenPolicyManager::new()
485            .with_mint_policy(MintPolicyConfig::OwnerOnly, PolicyRegistration::Active)?
486            .with_burn_policy(BurnPolicyConfig::AllowAll, PolicyRegistration::Active)?
487            .with_send_policy(TransferPolicy::AllowAll, PolicyRegistration::Active)?
488            .with_receive_policy(TransferPolicy::AllowAll, PolicyRegistration::Active)?;
489
490        let allowed_script_roots = allowed_script_roots
491            .into_iter()
492            .chain([MintNote::script_root(), BurnNote::script_root()])
493            .collect();
494
495        self.add_existing_fungible_faucet(
496            Auth::NetworkAccount {
497                allowed_script_roots,
498                allowed_tx_script_roots: BTreeSet::new(),
499            },
500            faucet,
501            AccountType::Public,
502            AccessControl::Ownable2Step { owner: owner_account_id },
503            token_policy_manager,
504        )
505    }
506
507    /// Convenience: builds a new (uncreated) basic auth-controlled fungible faucet from a
508    /// token-symbol shorthand using default decimals and `AllowAll` policies.
509    pub fn create_new_faucet(
510        &mut self,
511        auth_method: Auth,
512        token_symbol: &str,
513        max_supply: u64,
514    ) -> anyhow::Result<Account> {
515        let name = TokenName::new(token_symbol)?;
516        let symbol = TokenSymbol::new(token_symbol)
517            .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
518        let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
519        let faucet = FungibleFaucet::builder()
520            .name(name)
521            .symbol(symbol)
522            .decimals(DEFAULT_FAUCET_DECIMALS)
523            .max_supply(max_supply)
524            .build()
525            .context("failed to build FungibleFaucet")?;
526
527        let token_policy_manager = TokenPolicyManager::new()
528            .with_mint_policy(MintPolicyConfig::AllowAll, PolicyRegistration::Active)?
529            .with_burn_policy(BurnPolicyConfig::AllowAll, PolicyRegistration::Active)?
530            .with_send_policy(TransferPolicy::AllowAll, PolicyRegistration::Active)?
531            .with_receive_policy(TransferPolicy::AllowAll, PolicyRegistration::Active)?;
532
533        self.create_new_fungible_faucet(
534            auth_method,
535            faucet,
536            AccountType::Public,
537            AccessControl::AuthControlled,
538            token_policy_manager,
539        )
540    }
541
542    /// Creates a new public account with an [`MockAccountComponent`] and registers the
543    /// authenticator (if any).
544    pub fn create_new_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
545        let account_builder = Account::builder(self.rng.random())
546            .account_type(AccountType::Public)
547            .with_component(MockAccountComponent::with_empty_slots());
548
549        self.add_account_from_builder(auth_method, account_builder, AccountState::New)
550    }
551
552    /// Adds an existing public account with an [`MockAccountComponent`] to the initial chain state
553    /// and registers the authenticator (if any).
554    pub fn add_existing_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
555        self.add_existing_mock_account_with_storage_and_assets(auth_method, [], [])
556    }
557
558    /// Adds an existing public account with an [`MockAccountComponent`] to the initial chain state
559    /// and registers the authenticator (if any).
560    pub fn add_existing_mock_account_with_storage(
561        &mut self,
562        auth_method: Auth,
563        slots: impl IntoIterator<Item = StorageSlot>,
564    ) -> anyhow::Result<Account> {
565        self.add_existing_mock_account_with_storage_and_assets(auth_method, slots, [])
566    }
567
568    /// Adds an existing public account with an [`MockAccountComponent`] to the initial chain state
569    /// and registers the authenticator (if any).
570    pub fn add_existing_mock_account_with_assets(
571        &mut self,
572        auth_method: Auth,
573        assets: impl IntoIterator<Item = Asset>,
574    ) -> anyhow::Result<Account> {
575        self.add_existing_mock_account_with_storage_and_assets(auth_method, [], assets)
576    }
577
578    /// Adds an existing public account with an [`MockAccountComponent`] to the initial chain state
579    /// and registers the authenticator (if any).
580    pub fn add_existing_mock_account_with_storage_and_assets(
581        &mut self,
582        auth_method: Auth,
583        slots: impl IntoIterator<Item = StorageSlot>,
584        assets: impl IntoIterator<Item = Asset>,
585    ) -> anyhow::Result<Account> {
586        let account_builder = Account::builder(self.rng.random())
587            .account_type(AccountType::Public)
588            .with_component(MockAccountComponent::with_slots(slots.into_iter().collect()))
589            .with_assets(assets);
590
591        self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
592    }
593
594    /// Builds the provided [`AccountBuilder`] with the provided auth method and registers the
595    /// authenticator (if any).
596    ///
597    /// - If [`AccountState::Exists`] is given the account is built as an existing account and added
598    ///   to the initial chain state. It can then be used in a transaction without having to
599    ///   validate its seed.
600    /// - If [`AccountState::New`] is given the account is built as a new account and is **not**
601    ///   added to the chain. Its authenticator is registered (if present). Its first transaction
602    ///   will be its creation transaction. [`MockChain::build_tx_context`] can be called with the
603    ///   account to automatically add the authenticator.
604    pub fn add_account_from_builder(
605        &mut self,
606        auth_method: Auth,
607        mut account_builder: AccountBuilder,
608        account_state: AccountState,
609    ) -> anyhow::Result<Account> {
610        let (auth_component, authenticator) = auth_method.build_component();
611        account_builder = account_builder.with_auth_component(auth_component);
612
613        let account = if let AccountState::New = account_state {
614            account_builder.build().context("failed to build account from builder")?
615        } else {
616            account_builder
617                .build_existing()
618                .context("failed to build account from builder")?
619        };
620
621        self.account_authenticators
622            .insert(account.id(), AccountAuthenticator::new(authenticator));
623
624        if let AccountState::Exists = account_state {
625            self.accounts.insert(account.id(), account.clone());
626        }
627
628        Ok(account)
629    }
630    pub fn add_existing_account_from_components(
631        &mut self,
632        auth: Auth,
633        components: impl IntoIterator<Item = AccountComponent>,
634    ) -> anyhow::Result<Account> {
635        let mut account_builder =
636            Account::builder(rand::rng().random()).account_type(AccountType::Public);
637
638        for component in components {
639            account_builder = account_builder.with_component(component);
640        }
641
642        self.add_account_from_builder(auth, account_builder, AccountState::Exists)
643    }
644
645    /// Adds the provided account to the list of genesis accounts.
646    ///
647    /// This method only adds the account and does not store its account authenticator for it.
648    /// Calling [`MockChain::build_tx_context`] on accounts added in this way will not work if
649    /// the account needs an authenticator.
650    ///
651    /// Due to these limitations, prefer using other methods to add accounts to the chain, e.g.
652    /// [`MockChainBuilder::add_account_from_builder`].
653    pub fn add_account(&mut self, account: Account) -> anyhow::Result<()> {
654        self.accounts.insert(account.id(), account);
655
656        // This returns a Result to be conservative in case we need to return an error in the future
657        // and do not want to break this API.
658        Ok(())
659    }
660
661    // NOTE ADD METHODS
662    // ----------------------------------------------------------------------------------------
663
664    /// Adds the provided note to the initial chain state.
665    pub fn add_output_note(&mut self, note: impl Into<RawOutputNote>) {
666        self.notes.push(note.into());
667    }
668
669    /// Creates a new P2ANY note from the provided parameters and adds it to the list of
670    /// genesis notes.
671    ///
672    /// This note is similar to a P2ID note but can be consumed by any account.
673    pub fn add_p2any_note(
674        &mut self,
675        sender_account_id: AccountId,
676        note_type: NoteType,
677        assets: impl IntoIterator<Item = Asset>,
678    ) -> anyhow::Result<Note> {
679        let note = create_p2any_note(sender_account_id, note_type, assets, &mut self.rng);
680        self.add_output_note(RawOutputNote::Full(note.clone()));
681
682        Ok(note)
683    }
684
685    /// Creates a new P2ID note from the provided parameters and adds it to the list of genesis
686    /// notes.
687    ///
688    /// In the created [`MockChain`], the note will be immediately spendable by `target_account_id`
689    /// and carries no additional reclaim or timelock conditions.
690    pub fn add_p2id_note(
691        &mut self,
692        sender_account_id: AccountId,
693        target_account_id: AccountId,
694        asset: &[Asset],
695        note_type: NoteType,
696    ) -> Result<Note, NoteError> {
697        let note = P2idNote::create(
698            sender_account_id,
699            target_account_id,
700            asset.to_vec(),
701            note_type,
702            NoteAttachments::default(),
703            &mut self.rng,
704        )?;
705        self.add_output_note(RawOutputNote::Full(note.clone()));
706
707        Ok(note)
708    }
709
710    /// Adds a P2IDE note (pay‑to‑ID‑extended) to the list of genesis notes.
711    ///
712    /// A P2IDE note can include an optional `timelock_height` and/or an optional
713    /// `reclaim_height` after which the `sender_account_id` may reclaim the
714    /// funds.
715    pub fn add_p2ide_note(
716        &mut self,
717        sender_account_id: AccountId,
718        target_account_id: AccountId,
719        asset: &[Asset],
720        note_type: NoteType,
721        reclaim_height: Option<BlockNumber>,
722        timelock_height: Option<BlockNumber>,
723    ) -> Result<Note, NoteError> {
724        let storage = P2ideNoteStorage::new(target_account_id, reclaim_height, timelock_height);
725
726        let note = P2ideNote::create(
727            sender_account_id,
728            storage,
729            asset.to_vec(),
730            note_type,
731            NoteAttachments::default(),
732            &mut self.rng,
733        )?;
734
735        self.add_output_note(RawOutputNote::Full(note.clone()));
736
737        Ok(note)
738    }
739
740    /// Adds a public SWAP note to the list of genesis notes.
741    pub fn add_swap_note(
742        &mut self,
743        sender: AccountId,
744        offered_asset: Asset,
745        requested_asset: Asset,
746        payback_note_type: NoteType,
747    ) -> anyhow::Result<(Note, NoteDetails)> {
748        let (swap_note, payback_note) = SwapNote::create(
749            sender,
750            offered_asset,
751            requested_asset,
752            NoteType::Public,
753            NoteAttachments::default(),
754            payback_note_type,
755            &mut self.rng,
756        )?;
757
758        self.add_output_note(RawOutputNote::Full(swap_note.clone()));
759
760        Ok((swap_note, payback_note))
761    }
762
763    /// Adds a public `SPAWN` note to the list of genesis notes.
764    ///
765    /// A `SPAWN` note contains a note script that creates all `output_notes` that get passed as a
766    /// parameter.
767    ///
768    /// # Errors
769    ///
770    /// Returns an error if:
771    /// - the sender account ID of the provided output notes is not consistent or does not match the
772    ///   transaction's sender.
773    pub fn add_spawn_note<'note, I>(
774        &mut self,
775        output_notes: impl IntoIterator<Item = &'note Note, IntoIter = I>,
776    ) -> anyhow::Result<Note>
777    where
778        I: ExactSizeIterator<Item = &'note Note>,
779    {
780        let note = create_spawn_note(output_notes)?;
781        self.add_output_note(RawOutputNote::Full(note.clone()));
782
783        Ok(note)
784    }
785
786    /// Creates a new P2ID note with the provided amount of the fee asset of the chain.
787    ///
788    /// The fee faucet ID of the asset can be set using [`Self::fee_faucet_id`]. By default it
789    /// is [`ACCOUNT_ID_FEE_FAUCET`].
790    ///
791    /// In the created [`MockChain`], the note will be immediately spendable by `target_account_id`.
792    pub fn add_p2id_note_with_fee(
793        &mut self,
794        target_account_id: AccountId,
795        amount: u64,
796    ) -> anyhow::Result<Note> {
797        let fee_asset = self.fee_asset(amount)?;
798        let note = self.add_p2id_note(
799            self.fee_faucet_id,
800            target_account_id,
801            &[Asset::from(fee_asset)],
802            NoteType::Public,
803        )?;
804
805        Ok(note)
806    }
807
808    // HELPER FUNCTIONS
809    // ----------------------------------------------------------------------------------------
810
811    /// Returns a mutable reference to the builder's RNG.
812    ///
813    /// This can be used when creating accounts or notes and randomness is required.
814    pub fn rng_mut(&mut self) -> &mut RandomCoin {
815        &mut self.rng
816    }
817
818    /// Constructs a fungible asset based on the fee faucet ID and the provided amount.
819    fn fee_asset(&self, amount: u64) -> anyhow::Result<FungibleAsset> {
820        FungibleAsset::new(self.fee_faucet_id, amount).context("failed to create fee asset")
821    }
822}
823
824impl Default for MockChainBuilder {
825    fn default() -> Self {
826        Self::new()
827    }
828}