1use alloc::collections::{BTreeMap, BTreeSet};
2use alloc::vec::Vec;
3
4use anyhow::Context;
5
6const DEFAULT_FAUCET_DECIMALS: u8 = 10;
11
12use 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#[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_faucet_id: AccountId,
118 verification_base_fee: u32,
119}
120
121impl MockChainBuilder {
122 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 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 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 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 pub fn build(self) -> anyhow::Result<MockChain> {
184 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 pub fn add_account(&mut self, account: Account) -> anyhow::Result<()> {
654 self.accounts.insert(account.id(), account);
655
656 Ok(())
659 }
660
661 pub fn add_output_note(&mut self, note: impl Into<RawOutputNote>) {
666 self.notes.push(note.into());
667 }
668
669 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 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 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 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 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 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 pub fn rng_mut(&mut self) -> &mut RandomCoin {
815 &mut self.rng
816 }
817
818 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}