1use alloc::collections::{BTreeMap, BTreeSet};
2use alloc::vec::Vec;
3
4use anyhow::Context;
5
6const DEFAULT_FAUCET_DECIMALS: u8 = 10;
11
12const DEFAULT_VALIDATOR_COUNT: usize = 3;
16
17use itertools::Itertools;
21use miden_processor::crypto::random::RandomCoin;
22use miden_protocol::account::{
23 Account,
24 AccountBuilder,
25 AccountComponent,
26 AccountId,
27 AccountPatch,
28 AccountType,
29 AccountUpdateDetails,
30 AssetCallbackFlag,
31 StorageSlot,
32};
33use miden_protocol::asset::{Asset, AssetAmount, FungibleAsset, TokenSymbol};
34use miden_protocol::block::account_tree::AccountTree;
35use miden_protocol::block::nullifier_tree::NullifierTree;
36use miden_protocol::block::{
37 BlockAccountUpdate,
38 BlockBody,
39 BlockHeader,
40 BlockNoteTree,
41 BlockNumber,
42 BlockProof,
43 BlockSignatures,
44 Blockchain,
45 FeeParameters,
46 OutputNoteBatch,
47 ProvenBlock,
48 ValidatorKeys,
49};
50use miden_protocol::crypto::dsa::ecdsa_k256_keccak::SigningKey;
51use miden_protocol::crypto::merkle::smt::Smt;
52use miden_protocol::errors::NoteError;
53use miden_protocol::note::{Note, NoteAttachments, NoteDetails, NoteScriptRoot, NoteType};
54use miden_protocol::testing::account_id::ACCOUNT_ID_FEE_FAUCET;
55use miden_protocol::testing::random_secret_key::random_secret_key;
56use miden_protocol::transaction::{OrderedTransactionHeaders, RawOutputNote, TransactionKernel};
57use miden_protocol::{MAX_OUTPUT_NOTES_PER_BATCH, Word};
58use miden_standards::account::access::{AccessControl, Authority, Pausable, PausableManager};
59use miden_standards::account::faucets::{FungibleFaucet, TokenName};
60use miden_standards::account::policies::{
61 BurnPolicy,
62 MintPolicy,
63 TokenPolicyManager,
64 TransferPolicy,
65};
66use miden_standards::account::wallets::{BasicWallet, NoteCreator};
67use miden_standards::note::{BurnNote, MintNote, P2idNote, P2ideNote, SwapNote};
68use miden_standards::testing::account_component::MockAccountComponent;
69use rand::RngExt;
70
71use crate::mock_chain::chain::AccountAuthenticator;
72use crate::utils::{create_p2any_note, create_spawn_note};
73use crate::{AccountState, Auth, MockChain};
74
75#[derive(Debug, Clone)]
119pub struct MockChainBuilder {
120 accounts: BTreeMap<AccountId, Account>,
121 account_authenticators: BTreeMap<AccountId, AccountAuthenticator>,
122 notes: Vec<RawOutputNote>,
123 rng: RandomCoin,
124 fee_faucet_id: AccountId,
126 verification_base_fee: u32,
127}
128
129impl MockChainBuilder {
130 pub fn new() -> Self {
140 let fee_faucet_id = ACCOUNT_ID_FEE_FAUCET.try_into().expect("account ID should be valid");
141
142 Self {
143 accounts: BTreeMap::new(),
144 account_authenticators: BTreeMap::new(),
145 notes: Vec::new(),
146 rng: RandomCoin::new(Default::default()),
147 fee_faucet_id,
148 verification_base_fee: 0,
149 }
150 }
151
152 pub fn with_accounts(accounts: impl IntoIterator<Item = Account>) -> anyhow::Result<Self> {
161 let mut builder = Self::new();
162
163 for account in accounts {
164 builder.add_account(account)?;
165 }
166
167 Ok(builder)
168 }
169
170 pub fn fee_faucet_id(mut self, fee_faucet_id: AccountId) -> Self {
178 self.fee_faucet_id = fee_faucet_id;
179 self
180 }
181
182 pub fn verification_base_fee(mut self, verification_base_fee: u32) -> Self {
186 self.verification_base_fee = verification_base_fee;
187 self
188 }
189
190 pub fn build(self) -> anyhow::Result<MockChain> {
192 let block_account_updates: Vec<BlockAccountUpdate> = self
194 .accounts
195 .into_values()
196 .map(|account| {
197 let account_id = account.id();
198 let account_commitment = account.to_commitment();
199 let account_patch = AccountPatch::try_from(account)
200 .expect("chain builder should only store existing accounts without seeds");
201 let update_details = AccountUpdateDetails::Public(account_patch);
202
203 BlockAccountUpdate::new(account_id, account_commitment, update_details)
204 })
205 .collect();
206
207 let account_tree = AccountTree::with_entries(
208 block_account_updates
209 .iter()
210 .map(|account| (account.account_id(), account.final_state_commitment())),
211 )
212 .context("failed to create genesis account tree")?;
213
214 let full_notes: Vec<Note> = self
216 .notes
217 .iter()
218 .filter_map(|note| match note {
219 RawOutputNote::Full(n) => Some(n.clone()),
220 _ => None,
221 })
222 .collect();
223
224 let proven_notes: Vec<_> = self
225 .notes
226 .into_iter()
227 .map(|note| note.into_output_note().expect("genesis note should be valid"))
228 .collect();
229 let note_chunks = proven_notes.into_iter().chunks(MAX_OUTPUT_NOTES_PER_BATCH);
230 let output_note_batches: Vec<OutputNoteBatch> = note_chunks
231 .into_iter()
232 .map(|batch_notes| batch_notes.into_iter().enumerate().collect::<Vec<_>>())
233 .collect();
234
235 let created_nullifiers = Vec::new();
236 let transactions = OrderedTransactionHeaders::new_unchecked(Vec::new());
237
238 let note_tree = BlockNoteTree::from_note_batches(&output_note_batches)
239 .context("failed to create block note tree")?;
240
241 let version = 0;
242 let prev_block_commitment = Word::empty();
243 let block_num = BlockNumber::from(0u32);
244 let chain_commitment = Blockchain::new().commitment();
245 let account_root = account_tree.root();
246 let nullifier_root = NullifierTree::<Smt>::default().root();
247 let note_root = note_tree.root();
248 let tx_commitment = transactions.commitment();
249 let tx_kernel_commitment = TransactionKernel.to_commitment();
250 let timestamp = MockChain::TIMESTAMP_START_SECS;
251 let fee_parameters = FeeParameters::new(self.fee_faucet_id, self.verification_base_fee);
252 let validator_secret_keys: Vec<SigningKey> =
253 (0..DEFAULT_VALIDATOR_COUNT).map(|_| random_secret_key()).collect();
254 let validator_keys =
255 ValidatorKeys::new(validator_secret_keys.iter().map(|sk| sk.public_key()).collect())
256 .expect("randomly generated genesis validator keys should be distinct");
257
258 let header = BlockHeader::new(
259 version,
260 prev_block_commitment,
261 block_num,
262 chain_commitment,
263 account_root,
264 nullifier_root,
265 note_root,
266 tx_commitment,
267 tx_kernel_commitment,
268 validator_keys.clone(),
269 fee_parameters,
270 timestamp,
271 );
272
273 let body = BlockBody::new_unchecked(
274 block_account_updates,
275 output_note_batches,
276 created_nullifiers,
277 transactions,
278 );
279
280 let signatures = BlockSignatures::new(
283 validator_keys
284 .as_keys()
285 .iter()
286 .map(|key| {
287 let signer = validator_secret_keys
288 .iter()
289 .find(|sk| &sk.public_key() == key)
290 .expect("a signer should exist for every validator key");
291 signer.sign(header.commitment())
292 })
293 .collect(),
294 )
295 .expect("signature count same as validator key count");
296 let block_proof = BlockProof::new_dummy();
297 let genesis_block = ProvenBlock::new_unchecked(header, body, signatures, block_proof);
298
299 MockChain::from_genesis_block(
300 genesis_block,
301 account_tree,
302 self.account_authenticators,
303 validator_secret_keys,
304 full_notes,
305 )
306 }
307
308 pub fn create_new_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
317 let account_builder = AccountBuilder::new(self.rng.random())
318 .account_type(AccountType::Public)
319 .with_component(BasicWallet);
320
321 self.add_account_from_builder(auth_method, account_builder, AccountState::New)
322 }
323
324 pub fn add_existing_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
327 self.add_existing_wallet_with_assets(auth_method, [])
328 }
329
330 pub fn add_existing_wallet_with_assets(
333 &mut self,
334 auth_method: Auth,
335 assets: impl IntoIterator<Item = Asset>,
336 ) -> anyhow::Result<Account> {
337 let account_builder = Account::builder(self.rng.random())
338 .account_type(AccountType::Public)
339 .with_component(BasicWallet)
340 .with_assets(assets);
341
342 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
343 }
344
345 pub fn add_existing_note_creator(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
351 let account_builder = Account::builder(self.rng.random())
352 .account_type(AccountType::Public)
353 .with_component(NoteCreator);
354
355 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
356 }
357
358 fn add_existing_network_fungible_faucet(
361 &mut self,
362 auth_method: Auth,
363 faucet: FungibleFaucet,
364 account_type: AccountType,
365 access_control: AccessControl,
366 token_policy_manager: TokenPolicyManager,
367 ) -> anyhow::Result<Account> {
368 let account_builder = AccountBuilder::new(self.rng.random())
369 .account_type(account_type)
370 .with_component(faucet)
371 .with_components(access_control)
372 .with_asset_callbacks(AssetCallbackFlag::from(
373 token_policy_manager.has_transfer_policy(),
374 ))
375 .with_components(token_policy_manager)
376 .with_component(Pausable::unpaused())
377 .with_component(PausableManager);
378
379 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
380 }
381
382 pub fn add_existing_basic_faucet(
392 &mut self,
393 auth_method: Auth,
394 token_symbol: &str,
395 max_supply: u64,
396 token_supply: Option<u64>,
397 ) -> anyhow::Result<Account> {
398 let token_supply = token_supply.unwrap_or(0);
399 let name = TokenName::new(token_symbol)?;
400 let symbol = TokenSymbol::new(token_symbol)
401 .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
402 let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
403 let token_supply = AssetAmount::new(token_supply).context("invalid token_supply")?;
404 let faucet = FungibleFaucet::builder()
405 .name(name)
406 .symbol(symbol)
407 .decimals(DEFAULT_FAUCET_DECIMALS)
408 .max_supply(max_supply)
409 .token_supply(token_supply)
410 .build()
411 .context("failed to build FungibleFaucet")?;
412
413 let token_policy_manager = TokenPolicyManager::builder()
414 .active_mint_policy(MintPolicy::allow_all())
415 .active_burn_policy(BurnPolicy::allow_all())
416 .build();
417
418 let account_builder = AccountBuilder::new(self.rng.random())
419 .account_type(AccountType::Public)
420 .with_component(faucet)
421 .with_component(Authority::AuthControlled)
422 .with_asset_callbacks(AssetCallbackFlag::Disabled)
423 .with_components(token_policy_manager)
424 .with_component(Pausable::unpaused())
425 .with_component(PausableManager);
426
427 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
428 }
429
430 pub fn add_existing_network_faucet(
441 &mut self,
442 token_symbol: &str,
443 max_supply: u64,
444 owner_account_id: AccountId,
445 token_supply: Option<u64>,
446 mint_policy: MintPolicy,
447 allowed_script_roots: impl IntoIterator<Item = NoteScriptRoot>,
448 ) -> anyhow::Result<Account> {
449 let token_supply = token_supply.unwrap_or(0);
450 let name = TokenName::new(token_symbol)?;
451 let symbol = TokenSymbol::new(token_symbol)
452 .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
453 let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
454 let token_supply = AssetAmount::new(token_supply).context("invalid token_supply")?;
455 let faucet = FungibleFaucet::builder()
456 .name(name)
457 .symbol(symbol)
458 .decimals(DEFAULT_FAUCET_DECIMALS)
459 .max_supply(max_supply)
460 .token_supply(token_supply)
461 .build()
462 .context("failed to build FungibleFaucet")?;
463
464 let token_policy_manager = TokenPolicyManager::builder()
465 .active_mint_policy(mint_policy)
466 .active_burn_policy(BurnPolicy::allow_all())
467 .active_send_policy(TransferPolicy::allow_all())
468 .active_receive_policy(TransferPolicy::allow_all())
469 .build();
470
471 let allowed_script_roots: BTreeSet<NoteScriptRoot> = allowed_script_roots
472 .into_iter()
473 .chain([MintNote::script_root(), BurnNote::script_root()])
474 .collect();
475
476 self.add_existing_network_fungible_faucet(
477 Auth::NetworkAccount {
478 allowed_script_roots,
479 allowed_tx_script_roots: BTreeSet::new(),
480 },
481 faucet,
482 AccountType::Public,
483 AccessControl::Ownable2Step { owner: owner_account_id },
484 token_policy_manager,
485 )
486 }
487
488 pub fn add_existing_network_faucet_with_metadata(
495 &mut self,
496 owner_account_id: AccountId,
497 faucet: FungibleFaucet,
498 allowed_script_roots: impl IntoIterator<Item = NoteScriptRoot>,
499 ) -> anyhow::Result<Account> {
500 let token_policy_manager = TokenPolicyManager::builder()
501 .active_mint_policy(MintPolicy::owner_only())
502 .active_burn_policy(BurnPolicy::allow_all())
503 .active_send_policy(TransferPolicy::allow_all())
504 .active_receive_policy(TransferPolicy::allow_all())
505 .build();
506
507 let allowed_script_roots: BTreeSet<NoteScriptRoot> = allowed_script_roots
508 .into_iter()
509 .chain([MintNote::script_root(), BurnNote::script_root()])
510 .collect();
511
512 self.add_existing_network_fungible_faucet(
513 Auth::NetworkAccount {
514 allowed_script_roots,
515 allowed_tx_script_roots: BTreeSet::new(),
516 },
517 faucet,
518 AccountType::Public,
519 AccessControl::Ownable2Step { owner: owner_account_id },
520 token_policy_manager,
521 )
522 }
523
524 pub fn create_new_faucet(
528 &mut self,
529 auth_method: Auth,
530 token_symbol: &str,
531 max_supply: u64,
532 ) -> anyhow::Result<Account> {
533 let name = TokenName::new(token_symbol)?;
534 let symbol = TokenSymbol::new(token_symbol)
535 .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
536 let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
537 let faucet = FungibleFaucet::builder()
538 .name(name)
539 .symbol(symbol)
540 .decimals(DEFAULT_FAUCET_DECIMALS)
541 .max_supply(max_supply)
542 .build()
543 .context("failed to build FungibleFaucet")?;
544
545 let token_policy_manager = TokenPolicyManager::builder()
546 .active_mint_policy(MintPolicy::allow_all())
547 .active_burn_policy(BurnPolicy::allow_all())
548 .build();
549
550 let account_builder = AccountBuilder::new(self.rng.random())
551 .account_type(AccountType::Public)
552 .with_component(faucet)
553 .with_component(Authority::AuthControlled)
554 .with_asset_callbacks(AssetCallbackFlag::Disabled)
555 .with_components(token_policy_manager)
556 .with_component(Pausable::unpaused())
557 .with_component(PausableManager);
558
559 self.add_account_from_builder(auth_method, account_builder, AccountState::New)
560 }
561
562 pub fn create_new_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
565 let account_builder = Account::builder(self.rng.random())
566 .account_type(AccountType::Public)
567 .with_component(MockAccountComponent::with_empty_slots());
568
569 self.add_account_from_builder(auth_method, account_builder, AccountState::New)
570 }
571
572 pub fn add_existing_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
575 self.add_existing_mock_account_with_storage_and_assets(auth_method, [], [])
576 }
577
578 pub fn add_existing_mock_account_with_storage(
581 &mut self,
582 auth_method: Auth,
583 slots: impl IntoIterator<Item = StorageSlot>,
584 ) -> anyhow::Result<Account> {
585 self.add_existing_mock_account_with_storage_and_assets(auth_method, slots, [])
586 }
587
588 pub fn add_existing_mock_account_with_assets(
591 &mut self,
592 auth_method: Auth,
593 assets: impl IntoIterator<Item = Asset>,
594 ) -> anyhow::Result<Account> {
595 self.add_existing_mock_account_with_storage_and_assets(auth_method, [], assets)
596 }
597
598 pub fn add_existing_mock_account_with_storage_and_assets(
601 &mut self,
602 auth_method: Auth,
603 slots: impl IntoIterator<Item = StorageSlot>,
604 assets: impl IntoIterator<Item = Asset>,
605 ) -> anyhow::Result<Account> {
606 let account_builder = Account::builder(self.rng.random())
607 .account_type(AccountType::Public)
608 .with_component(MockAccountComponent::with_slots(slots.into_iter().collect()))
609 .with_assets(assets);
610
611 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
612 }
613
614 pub fn add_account_from_builder(
625 &mut self,
626 auth_method: Auth,
627 mut account_builder: AccountBuilder,
628 account_state: AccountState,
629 ) -> anyhow::Result<Account> {
630 let (auth_component, authenticator) = auth_method.build_component();
631 account_builder = account_builder.with_auth_component(auth_component);
632
633 let account = if let AccountState::New = account_state {
634 account_builder.build().context("failed to build account from builder")?
635 } else {
636 account_builder
637 .build_existing()
638 .context("failed to build account from builder")?
639 };
640
641 self.account_authenticators
642 .insert(account.id(), AccountAuthenticator::new(authenticator));
643
644 if let AccountState::Exists = account_state {
645 self.accounts.insert(account.id(), account.clone());
646 }
647
648 Ok(account)
649 }
650 pub fn add_existing_account_from_components(
651 &mut self,
652 auth: Auth,
653 components: impl IntoIterator<Item = AccountComponent>,
654 ) -> anyhow::Result<Account> {
655 let mut account_builder =
656 Account::builder(rand::rng().random()).account_type(AccountType::Public);
657
658 for component in components {
659 account_builder = account_builder.with_component(component);
660 }
661
662 self.add_account_from_builder(auth, account_builder, AccountState::Exists)
663 }
664
665 pub fn add_account(&mut self, account: Account) -> anyhow::Result<()> {
674 self.accounts.insert(account.id(), account);
675
676 Ok(())
679 }
680
681 pub fn add_output_note(&mut self, note: impl Into<RawOutputNote>) {
686 self.notes.push(note.into());
687 }
688
689 pub fn add_p2any_note(
694 &mut self,
695 sender_account_id: AccountId,
696 note_type: NoteType,
697 assets: impl IntoIterator<Item = Asset>,
698 ) -> anyhow::Result<Note> {
699 let note = create_p2any_note(sender_account_id, note_type, assets, &mut self.rng);
700 self.add_output_note(RawOutputNote::Full(note.clone()));
701
702 Ok(note)
703 }
704
705 pub fn add_p2id_note(
711 &mut self,
712 sender_account_id: AccountId,
713 target_account_id: AccountId,
714 asset: &[Asset],
715 note_type: NoteType,
716 ) -> Result<Note, NoteError> {
717 let note: Note = P2idNote::builder()
718 .sender(sender_account_id)
719 .target(target_account_id)
720 .assets(asset.iter().copied())
721 .note_type(note_type)
722 .generate_serial_number(&mut self.rng)
723 .build()?
724 .into();
725 self.add_output_note(RawOutputNote::Full(note.clone()));
726
727 Ok(note)
728 }
729
730 pub fn add_p2ide_note(
738 &mut self,
739 sender_account_id: AccountId,
740 target_account_id: AccountId,
741 reclaimer: Option<AccountId>,
742 asset: &[Asset],
743 note_type: NoteType,
744 reclaim_height: Option<BlockNumber>,
745 timelock_height: Option<BlockNumber>,
746 ) -> Result<Note, NoteError> {
747 let note: Note = P2ideNote::builder()
748 .sender(sender_account_id)
749 .target(target_account_id)
750 .maybe_reclaimer(reclaimer)
751 .assets(asset.iter().copied())
752 .note_type(note_type)
753 .maybe_reclaim_height(reclaim_height)
754 .maybe_timelock_height(timelock_height)
755 .generate_serial_number(&mut self.rng)
756 .build()?
757 .into();
758
759 self.add_output_note(RawOutputNote::Full(note.clone()));
760
761 Ok(note)
762 }
763
764 pub fn add_swap_note(
766 &mut self,
767 sender: AccountId,
768 offered_asset: Asset,
769 requested_asset: Asset,
770 payback_note_type: NoteType,
771 ) -> anyhow::Result<(Note, NoteDetails)> {
772 let (swap_note, payback_note) = SwapNote::create(
773 sender,
774 offered_asset,
775 requested_asset,
776 NoteType::Public,
777 NoteAttachments::default(),
778 payback_note_type,
779 &mut self.rng,
780 )?;
781
782 self.add_output_note(RawOutputNote::Full(swap_note.clone()));
783
784 Ok((swap_note, payback_note))
785 }
786
787 pub fn add_spawn_note<'note, I>(
798 &mut self,
799 output_notes: impl IntoIterator<Item = &'note Note, IntoIter = I>,
800 ) -> anyhow::Result<Note>
801 where
802 I: ExactSizeIterator<Item = &'note Note>,
803 {
804 let note = create_spawn_note(output_notes)?;
805 self.add_output_note(RawOutputNote::Full(note.clone()));
806
807 Ok(note)
808 }
809
810 pub fn add_p2id_note_with_fee(
817 &mut self,
818 target_account_id: AccountId,
819 amount: u64,
820 ) -> anyhow::Result<Note> {
821 let fee_asset = self.fee_asset(amount)?;
822 let note = self.add_p2id_note(
823 self.fee_faucet_id,
824 target_account_id,
825 &[Asset::from(fee_asset)],
826 NoteType::Public,
827 )?;
828
829 Ok(note)
830 }
831
832 pub fn rng_mut(&mut self) -> &mut RandomCoin {
839 &mut self.rng
840 }
841
842 fn fee_asset(&self, amount: u64) -> anyhow::Result<FungibleAsset> {
844 FungibleAsset::new(self.fee_faucet_id, amount).context("failed to create fee asset")
845 }
846}
847
848impl Default for MockChainBuilder {
849 fn default() -> Self {
850 Self::new()
851 }
852}