1use alloc::collections::BTreeMap;
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 { allowed_script_roots },
462 faucet,
463 AccountType::Public,
464 AccessControl::Ownable2Step { owner: owner_account_id },
465 token_policy_manager,
466 )
467 }
468
469 pub fn add_existing_network_faucet_with_metadata(
476 &mut self,
477 owner_account_id: AccountId,
478 faucet: FungibleFaucet,
479 allowed_script_roots: impl IntoIterator<Item = NoteScriptRoot>,
480 ) -> anyhow::Result<Account> {
481 let token_policy_manager = TokenPolicyManager::new()
482 .with_mint_policy(MintPolicyConfig::OwnerOnly, PolicyRegistration::Active)?
483 .with_burn_policy(BurnPolicyConfig::AllowAll, PolicyRegistration::Active)?
484 .with_send_policy(TransferPolicy::AllowAll, PolicyRegistration::Active)?
485 .with_receive_policy(TransferPolicy::AllowAll, PolicyRegistration::Active)?;
486
487 let allowed_script_roots = allowed_script_roots
488 .into_iter()
489 .chain([MintNote::script_root(), BurnNote::script_root()])
490 .collect();
491
492 self.add_existing_fungible_faucet(
493 Auth::NetworkAccount { allowed_script_roots },
494 faucet,
495 AccountType::Public,
496 AccessControl::Ownable2Step { owner: owner_account_id },
497 token_policy_manager,
498 )
499 }
500
501 pub fn create_new_faucet(
504 &mut self,
505 auth_method: Auth,
506 token_symbol: &str,
507 max_supply: u64,
508 ) -> anyhow::Result<Account> {
509 let name = TokenName::new(token_symbol)?;
510 let symbol = TokenSymbol::new(token_symbol)
511 .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
512 let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
513 let faucet = FungibleFaucet::builder()
514 .name(name)
515 .symbol(symbol)
516 .decimals(DEFAULT_FAUCET_DECIMALS)
517 .max_supply(max_supply)
518 .build()
519 .context("failed to build FungibleFaucet")?;
520
521 let token_policy_manager = TokenPolicyManager::new()
522 .with_mint_policy(MintPolicyConfig::AllowAll, PolicyRegistration::Active)?
523 .with_burn_policy(BurnPolicyConfig::AllowAll, PolicyRegistration::Active)?
524 .with_send_policy(TransferPolicy::AllowAll, PolicyRegistration::Active)?
525 .with_receive_policy(TransferPolicy::AllowAll, PolicyRegistration::Active)?;
526
527 self.create_new_fungible_faucet(
528 auth_method,
529 faucet,
530 AccountType::Public,
531 AccessControl::AuthControlled,
532 token_policy_manager,
533 )
534 }
535
536 pub fn create_new_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
539 let account_builder = Account::builder(self.rng.random())
540 .account_type(AccountType::Public)
541 .with_component(MockAccountComponent::with_empty_slots());
542
543 self.add_account_from_builder(auth_method, account_builder, AccountState::New)
544 }
545
546 pub fn add_existing_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
549 self.add_existing_mock_account_with_storage_and_assets(auth_method, [], [])
550 }
551
552 pub fn add_existing_mock_account_with_storage(
555 &mut self,
556 auth_method: Auth,
557 slots: impl IntoIterator<Item = StorageSlot>,
558 ) -> anyhow::Result<Account> {
559 self.add_existing_mock_account_with_storage_and_assets(auth_method, slots, [])
560 }
561
562 pub fn add_existing_mock_account_with_assets(
565 &mut self,
566 auth_method: Auth,
567 assets: impl IntoIterator<Item = Asset>,
568 ) -> anyhow::Result<Account> {
569 self.add_existing_mock_account_with_storage_and_assets(auth_method, [], assets)
570 }
571
572 pub fn add_existing_mock_account_with_storage_and_assets(
575 &mut self,
576 auth_method: Auth,
577 slots: impl IntoIterator<Item = StorageSlot>,
578 assets: impl IntoIterator<Item = Asset>,
579 ) -> anyhow::Result<Account> {
580 let account_builder = Account::builder(self.rng.random())
581 .account_type(AccountType::Public)
582 .with_component(MockAccountComponent::with_slots(slots.into_iter().collect()))
583 .with_assets(assets);
584
585 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
586 }
587
588 pub fn add_account_from_builder(
599 &mut self,
600 auth_method: Auth,
601 mut account_builder: AccountBuilder,
602 account_state: AccountState,
603 ) -> anyhow::Result<Account> {
604 let (auth_component, authenticator) = auth_method.build_component();
605 account_builder = account_builder.with_auth_component(auth_component);
606
607 let account = if let AccountState::New = account_state {
608 account_builder.build().context("failed to build account from builder")?
609 } else {
610 account_builder
611 .build_existing()
612 .context("failed to build account from builder")?
613 };
614
615 self.account_authenticators
616 .insert(account.id(), AccountAuthenticator::new(authenticator));
617
618 if let AccountState::Exists = account_state {
619 self.accounts.insert(account.id(), account.clone());
620 }
621
622 Ok(account)
623 }
624 pub fn add_existing_account_from_components(
625 &mut self,
626 auth: Auth,
627 components: impl IntoIterator<Item = AccountComponent>,
628 ) -> anyhow::Result<Account> {
629 let mut account_builder =
630 Account::builder(rand::rng().random()).account_type(AccountType::Public);
631
632 for component in components {
633 account_builder = account_builder.with_component(component);
634 }
635
636 self.add_account_from_builder(auth, account_builder, AccountState::Exists)
637 }
638
639 pub fn add_account(&mut self, account: Account) -> anyhow::Result<()> {
648 self.accounts.insert(account.id(), account);
649
650 Ok(())
653 }
654
655 pub fn add_output_note(&mut self, note: impl Into<RawOutputNote>) {
660 self.notes.push(note.into());
661 }
662
663 pub fn add_p2any_note(
668 &mut self,
669 sender_account_id: AccountId,
670 note_type: NoteType,
671 assets: impl IntoIterator<Item = Asset>,
672 ) -> anyhow::Result<Note> {
673 let note = create_p2any_note(sender_account_id, note_type, assets, &mut self.rng);
674 self.add_output_note(RawOutputNote::Full(note.clone()));
675
676 Ok(note)
677 }
678
679 pub fn add_p2id_note(
685 &mut self,
686 sender_account_id: AccountId,
687 target_account_id: AccountId,
688 asset: &[Asset],
689 note_type: NoteType,
690 ) -> Result<Note, NoteError> {
691 let note = P2idNote::create(
692 sender_account_id,
693 target_account_id,
694 asset.to_vec(),
695 note_type,
696 NoteAttachments::default(),
697 &mut self.rng,
698 )?;
699 self.add_output_note(RawOutputNote::Full(note.clone()));
700
701 Ok(note)
702 }
703
704 pub fn add_p2ide_note(
710 &mut self,
711 sender_account_id: AccountId,
712 target_account_id: AccountId,
713 asset: &[Asset],
714 note_type: NoteType,
715 reclaim_height: Option<BlockNumber>,
716 timelock_height: Option<BlockNumber>,
717 ) -> Result<Note, NoteError> {
718 let storage = P2ideNoteStorage::new(target_account_id, reclaim_height, timelock_height);
719
720 let note = P2ideNote::create(
721 sender_account_id,
722 storage,
723 asset.to_vec(),
724 note_type,
725 NoteAttachments::default(),
726 &mut self.rng,
727 )?;
728
729 self.add_output_note(RawOutputNote::Full(note.clone()));
730
731 Ok(note)
732 }
733
734 pub fn add_swap_note(
736 &mut self,
737 sender: AccountId,
738 offered_asset: Asset,
739 requested_asset: Asset,
740 payback_note_type: NoteType,
741 ) -> anyhow::Result<(Note, NoteDetails)> {
742 let (swap_note, payback_note) = SwapNote::create(
743 sender,
744 offered_asset,
745 requested_asset,
746 NoteType::Public,
747 NoteAttachments::default(),
748 payback_note_type,
749 &mut self.rng,
750 )?;
751
752 self.add_output_note(RawOutputNote::Full(swap_note.clone()));
753
754 Ok((swap_note, payback_note))
755 }
756
757 pub fn add_spawn_note<'note, I>(
768 &mut self,
769 output_notes: impl IntoIterator<Item = &'note Note, IntoIter = I>,
770 ) -> anyhow::Result<Note>
771 where
772 I: ExactSizeIterator<Item = &'note Note>,
773 {
774 let note = create_spawn_note(output_notes)?;
775 self.add_output_note(RawOutputNote::Full(note.clone()));
776
777 Ok(note)
778 }
779
780 pub fn add_p2id_note_with_fee(
787 &mut self,
788 target_account_id: AccountId,
789 amount: u64,
790 ) -> anyhow::Result<Note> {
791 let fee_asset = self.fee_asset(amount)?;
792 let note = self.add_p2id_note(
793 self.fee_faucet_id,
794 target_account_id,
795 &[Asset::from(fee_asset)],
796 NoteType::Public,
797 )?;
798
799 Ok(note)
800 }
801
802 pub fn rng_mut(&mut self) -> &mut RandomCoin {
809 &mut self.rng
810 }
811
812 fn fee_asset(&self, amount: u64) -> anyhow::Result<FungibleAsset> {
814 FungibleAsset::new(self.fee_faucet_id, amount).context("failed to create fee asset")
815 }
816}
817
818impl Default for MockChainBuilder {
819 fn default() -> Self {
820 Self::new()
821 }
822}