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::auth::AuthNetworkAccount;
60use miden_standards::account::faucets::{FungibleFaucet, TokenName};
61use miden_standards::account::fees::{ConstantFeePolicy, FeeManager};
62use miden_standards::account::policies::{
63 BurnPolicy,
64 MintPolicy,
65 TokenPolicyManager,
66 TransferPolicy,
67};
68use miden_standards::account::wallets::{BasicWallet, NoteCreator};
69use miden_standards::note::{BurnNote, MintNote, P2idNote, P2ideNote, SwapNote, TxFeeNote};
70use miden_standards::testing::account_component::MockAccountComponent;
71use rand::RngExt;
72
73use crate::mock_chain::chain::AccountAuthenticator;
74use crate::utils::{create_p2any_note, create_spawn_note};
75use crate::{AccountState, Auth, MockChain};
76
77#[derive(Debug, Clone)]
121pub struct MockChainBuilder {
122 accounts: BTreeMap<AccountId, Account>,
123 account_authenticators: BTreeMap<AccountId, AccountAuthenticator>,
124 notes: Vec<RawOutputNote>,
125 rng: RandomCoin,
126 fee_faucet_id: AccountId,
128 verification_base_fee: u32,
129}
130
131impl MockChainBuilder {
132 pub fn new() -> Self {
142 let fee_faucet_id = ACCOUNT_ID_FEE_FAUCET.try_into().expect("account ID should be valid");
143
144 Self {
145 accounts: BTreeMap::new(),
146 account_authenticators: BTreeMap::new(),
147 notes: Vec::new(),
148 rng: RandomCoin::new(Default::default()),
149 fee_faucet_id,
150 verification_base_fee: 0,
151 }
152 }
153
154 pub fn with_accounts(accounts: impl IntoIterator<Item = Account>) -> anyhow::Result<Self> {
163 let mut builder = Self::new();
164
165 for account in accounts {
166 builder.add_account(account)?;
167 }
168
169 Ok(builder)
170 }
171
172 pub fn fee_faucet_id(mut self, fee_faucet_id: AccountId) -> Self {
180 self.fee_faucet_id = fee_faucet_id;
181 self
182 }
183
184 pub fn verification_base_fee(mut self, verification_base_fee: u32) -> Self {
188 self.verification_base_fee = verification_base_fee;
189 self
190 }
191
192 pub fn build(self) -> anyhow::Result<MockChain> {
194 let block_account_updates: Vec<BlockAccountUpdate> = self
196 .accounts
197 .into_values()
198 .map(|account| {
199 let account_id = account.id();
200 let account_commitment = account.to_commitment();
201 let account_patch = AccountPatch::try_from(account)
202 .expect("chain builder should only store existing accounts without seeds");
203 let update_details = AccountUpdateDetails::Public(account_patch);
204
205 BlockAccountUpdate::new(account_id, account_commitment, update_details)
206 })
207 .collect();
208
209 let account_tree = AccountTree::with_entries(
210 block_account_updates
211 .iter()
212 .map(|account| (account.account_id(), account.final_state_commitment())),
213 )
214 .context("failed to create genesis account tree")?;
215
216 let full_notes: Vec<Note> = self
218 .notes
219 .iter()
220 .filter_map(|note| match note {
221 RawOutputNote::Full(n) => Some(n.clone()),
222 _ => None,
223 })
224 .collect();
225
226 let proven_notes: Vec<_> = self
227 .notes
228 .into_iter()
229 .map(|note| note.into_output_note().expect("genesis note should be valid"))
230 .collect();
231 let note_chunks = proven_notes.into_iter().chunks(MAX_OUTPUT_NOTES_PER_BATCH);
232 let output_note_batches: Vec<OutputNoteBatch> = note_chunks
233 .into_iter()
234 .map(|batch_notes| batch_notes.into_iter().enumerate().collect::<Vec<_>>())
235 .collect();
236
237 let created_nullifiers = Vec::new();
238 let transactions = OrderedTransactionHeaders::new_unchecked(Vec::new());
239
240 let note_tree = BlockNoteTree::from_note_batches(&output_note_batches)
241 .context("failed to create block note tree")?;
242
243 let version = 0;
244 let prev_block_commitment = Word::empty();
245 let block_num = BlockNumber::from(0u32);
246 let chain_commitment = Blockchain::new().commitment();
247 let account_root = account_tree.root();
248 let nullifier_root = NullifierTree::<Smt>::default().root();
249 let note_root = note_tree.root();
250 let tx_commitment = transactions.commitment();
251 let tx_kernel_commitment = TransactionKernel.to_commitment();
252 let timestamp = MockChain::TIMESTAMP_START_SECS;
253 let fee_parameters = FeeParameters::new(self.fee_faucet_id, self.verification_base_fee);
254 let validator_secret_keys: Vec<SigningKey> =
255 (0..DEFAULT_VALIDATOR_COUNT).map(|_| random_secret_key()).collect();
256 let validator_keys =
257 ValidatorKeys::new(validator_secret_keys.iter().map(|sk| sk.public_key()).collect())
258 .expect("randomly generated genesis validator keys should be distinct");
259
260 let header = BlockHeader::new(
261 version,
262 prev_block_commitment,
263 block_num,
264 chain_commitment,
265 account_root,
266 nullifier_root,
267 note_root,
268 tx_commitment,
269 tx_kernel_commitment,
270 validator_keys.clone(),
271 fee_parameters,
272 timestamp,
273 );
274
275 let body = BlockBody::new_unchecked(
276 block_account_updates,
277 output_note_batches,
278 created_nullifiers,
279 transactions,
280 );
281
282 let signatures = BlockSignatures::new(
285 validator_keys
286 .as_keys()
287 .iter()
288 .map(|key| {
289 let signer = validator_secret_keys
290 .iter()
291 .find(|sk| &sk.public_key() == key)
292 .expect("a signer should exist for every validator key");
293 signer.sign(header.commitment())
294 })
295 .collect(),
296 )
297 .expect("signature count same as validator key count");
298 let block_proof = BlockProof::new_dummy();
299 let genesis_block = ProvenBlock::new_unchecked(header, body, signatures, block_proof);
300
301 MockChain::from_genesis_block(
302 genesis_block,
303 account_tree,
304 self.account_authenticators,
305 validator_secret_keys,
306 full_notes,
307 )
308 }
309
310 pub fn create_new_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
319 let account_builder = AccountBuilder::new(self.rng.random())
320 .account_type(AccountType::Public)
321 .with_component(BasicWallet);
322
323 self.add_account_from_builder(auth_method, account_builder, AccountState::New)
324 }
325
326 pub fn add_existing_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
329 self.add_existing_wallet_with_assets(auth_method, [])
330 }
331
332 pub fn add_existing_wallet_with_assets(
335 &mut self,
336 auth_method: Auth,
337 assets: impl IntoIterator<Item = Asset>,
338 ) -> anyhow::Result<Account> {
339 let account_builder = Account::builder(self.rng.random())
340 .account_type(AccountType::Public)
341 .with_component(BasicWallet)
342 .with_assets(assets);
343
344 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
345 }
346
347 pub fn add_existing_note_creator(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
353 let account_builder = Account::builder(self.rng.random())
354 .account_type(AccountType::Public)
355 .with_component(NoteCreator);
356
357 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
358 }
359
360 fn add_existing_network_fungible_faucet(
363 &mut self,
364 auth: AuthNetworkAccount,
365 faucet: FungibleFaucet,
366 account_type: AccountType,
367 access_control: AccessControl,
368 token_policy_manager: TokenPolicyManager,
369 ) -> anyhow::Result<Account> {
370 let mut constant_fee_policy = ConstantFeePolicy::new();
374 for note_script in auth.allowed_notes().allowed_script_roots() {
375 constant_fee_policy = constant_fee_policy.with_fee(*note_script, AssetAmount::ZERO)
376 }
377
378 let fee_manager = FeeManager::builder()
379 .active_fee_policy(constant_fee_policy.into())
380 .fee_faucet_id(self.fee_faucet_id)
381 .build();
382
383 let account_builder = AccountBuilder::new(self.rng.random())
384 .account_type(account_type)
385 .with_component(faucet)
386 .with_components(access_control)
387 .with_asset_callbacks(AssetCallbackFlag::from(
388 token_policy_manager.has_transfer_policy(),
389 ))
390 .with_components(token_policy_manager)
391 .with_component(Pausable::unpaused())
392 .with_component(PausableManager)
393 .with_components(fee_manager);
394
395 let auth = Auth::NetworkAccount {
396 allowed_script_roots: auth.allowed_notes().allowed_script_roots().clone(),
397 allowed_tx_script_roots: auth.allowed_tx_scripts().allowed_script_roots().clone(),
398 };
399
400 self.add_account_from_builder(auth, account_builder, AccountState::Exists)
401 }
402
403 pub fn add_existing_basic_faucet(
413 &mut self,
414 auth_method: Auth,
415 token_symbol: &str,
416 max_supply: u64,
417 token_supply: Option<u64>,
418 ) -> anyhow::Result<Account> {
419 let token_supply = token_supply.unwrap_or(0);
420 let name = TokenName::new(token_symbol)?;
421 let symbol = TokenSymbol::new(token_symbol)
422 .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
423 let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
424 let token_supply = AssetAmount::new(token_supply).context("invalid token_supply")?;
425 let faucet = FungibleFaucet::builder()
426 .name(name)
427 .symbol(symbol)
428 .decimals(DEFAULT_FAUCET_DECIMALS)
429 .max_supply(max_supply)
430 .token_supply(token_supply)
431 .build()
432 .context("failed to build FungibleFaucet")?;
433
434 let token_policy_manager = TokenPolicyManager::builder()
435 .active_mint_policy(MintPolicy::allow_all())
436 .active_burn_policy(BurnPolicy::allow_all())
437 .build();
438
439 let account_builder = AccountBuilder::new(self.rng.random())
440 .account_type(AccountType::Public)
441 .with_component(faucet)
442 .with_component(Authority::AuthControlled)
443 .with_asset_callbacks(AssetCallbackFlag::Disabled)
444 .with_components(token_policy_manager)
445 .with_component(Pausable::unpaused())
446 .with_component(PausableManager);
447
448 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
449 }
450
451 pub fn add_existing_network_faucet(
462 &mut self,
463 token_symbol: &str,
464 max_supply: u64,
465 owner_account_id: AccountId,
466 token_supply: Option<u64>,
467 mint_policy: MintPolicy,
468 allowed_script_roots: impl IntoIterator<Item = NoteScriptRoot>,
469 ) -> anyhow::Result<Account> {
470 let token_supply = token_supply.unwrap_or(0);
471 let name = TokenName::new(token_symbol)?;
472 let symbol = TokenSymbol::new(token_symbol)
473 .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
474 let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
475 let token_supply = AssetAmount::new(token_supply).context("invalid token_supply")?;
476 let faucet = FungibleFaucet::builder()
477 .name(name)
478 .symbol(symbol)
479 .decimals(DEFAULT_FAUCET_DECIMALS)
480 .max_supply(max_supply)
481 .token_supply(token_supply)
482 .build()
483 .context("failed to build FungibleFaucet")?;
484
485 let token_policy_manager = TokenPolicyManager::builder()
486 .active_mint_policy(mint_policy)
487 .active_burn_policy(BurnPolicy::allow_all())
488 .active_send_policy(TransferPolicy::allow_all())
489 .active_receive_policy(TransferPolicy::allow_all())
490 .build();
491
492 let allowed_script_roots: BTreeSet<NoteScriptRoot> = allowed_script_roots
493 .into_iter()
494 .chain([MintNote::script_root(), BurnNote::script_root()])
495 .collect();
496
497 self.add_existing_network_fungible_faucet(
498 AuthNetworkAccount::with_allowed_notes(allowed_script_roots)?,
499 faucet,
500 AccountType::Public,
501 AccessControl::Ownable2Step { owner: owner_account_id },
502 token_policy_manager,
503 )
504 }
505
506 pub fn add_existing_network_faucet_with_metadata(
513 &mut self,
514 owner_account_id: AccountId,
515 faucet: FungibleFaucet,
516 allowed_script_roots: impl IntoIterator<Item = NoteScriptRoot>,
517 ) -> anyhow::Result<Account> {
518 let token_policy_manager = TokenPolicyManager::builder()
519 .active_mint_policy(MintPolicy::owner_only())
520 .active_burn_policy(BurnPolicy::allow_all())
521 .active_send_policy(TransferPolicy::allow_all())
522 .active_receive_policy(TransferPolicy::allow_all())
523 .build();
524
525 let allowed_script_roots: BTreeSet<NoteScriptRoot> = allowed_script_roots
526 .into_iter()
527 .chain([MintNote::script_root(), BurnNote::script_root()])
528 .collect();
529
530 self.add_existing_network_fungible_faucet(
531 AuthNetworkAccount::with_allowed_notes(allowed_script_roots)?,
532 faucet,
533 AccountType::Public,
534 AccessControl::Ownable2Step { owner: owner_account_id },
535 token_policy_manager,
536 )
537 }
538
539 pub fn create_new_faucet(
543 &mut self,
544 auth_method: Auth,
545 token_symbol: &str,
546 max_supply: u64,
547 ) -> anyhow::Result<Account> {
548 let name = TokenName::new(token_symbol)?;
549 let symbol = TokenSymbol::new(token_symbol)
550 .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
551 let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
552 let faucet = FungibleFaucet::builder()
553 .name(name)
554 .symbol(symbol)
555 .decimals(DEFAULT_FAUCET_DECIMALS)
556 .max_supply(max_supply)
557 .build()
558 .context("failed to build FungibleFaucet")?;
559
560 let token_policy_manager = TokenPolicyManager::builder()
561 .active_mint_policy(MintPolicy::allow_all())
562 .active_burn_policy(BurnPolicy::allow_all())
563 .build();
564
565 let account_builder = AccountBuilder::new(self.rng.random())
566 .account_type(AccountType::Public)
567 .with_component(faucet)
568 .with_component(Authority::AuthControlled)
569 .with_asset_callbacks(AssetCallbackFlag::Disabled)
570 .with_components(token_policy_manager)
571 .with_component(Pausable::unpaused())
572 .with_component(PausableManager);
573
574 self.add_account_from_builder(auth_method, account_builder, AccountState::New)
575 }
576
577 pub fn create_new_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
580 let account_builder = Account::builder(self.rng.random())
581 .account_type(AccountType::Public)
582 .with_component(MockAccountComponent::with_empty_slots());
583
584 self.add_account_from_builder(auth_method, account_builder, AccountState::New)
585 }
586
587 pub fn add_existing_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
590 self.add_existing_mock_account_with_storage_and_assets(auth_method, [], [])
591 }
592
593 pub fn add_existing_mock_account_with_storage(
596 &mut self,
597 auth_method: Auth,
598 slots: impl IntoIterator<Item = StorageSlot>,
599 ) -> anyhow::Result<Account> {
600 self.add_existing_mock_account_with_storage_and_assets(auth_method, slots, [])
601 }
602
603 pub fn add_existing_mock_account_with_assets(
606 &mut self,
607 auth_method: Auth,
608 assets: impl IntoIterator<Item = Asset>,
609 ) -> anyhow::Result<Account> {
610 self.add_existing_mock_account_with_storage_and_assets(auth_method, [], assets)
611 }
612
613 pub fn add_existing_mock_account_with_storage_and_assets(
616 &mut self,
617 auth_method: Auth,
618 slots: impl IntoIterator<Item = StorageSlot>,
619 assets: impl IntoIterator<Item = Asset>,
620 ) -> anyhow::Result<Account> {
621 let account_builder = Account::builder(self.rng.random())
622 .account_type(AccountType::Public)
623 .with_component(MockAccountComponent::with_slots(slots.into_iter().collect()))
624 .with_assets(assets);
625
626 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
627 }
628
629 pub fn add_account_from_builder(
640 &mut self,
641 auth_method: Auth,
642 mut account_builder: AccountBuilder,
643 account_state: AccountState,
644 ) -> anyhow::Result<Account> {
645 let (auth_component, authenticator) = auth_method.build_component();
646 account_builder = account_builder.with_auth_component(auth_component);
647
648 let account = if let AccountState::New = account_state {
649 account_builder.build().context("failed to build account from builder")?
650 } else {
651 account_builder
652 .build_existing()
653 .context("failed to build account from builder")?
654 };
655
656 self.account_authenticators
657 .insert(account.id(), AccountAuthenticator::new(authenticator));
658
659 if let AccountState::Exists = account_state {
660 self.accounts.insert(account.id(), account.clone());
661 }
662
663 Ok(account)
664 }
665 pub fn add_existing_account_from_components(
666 &mut self,
667 auth: Auth,
668 components: impl IntoIterator<Item = AccountComponent>,
669 ) -> anyhow::Result<Account> {
670 let mut account_builder =
671 Account::builder(rand::rng().random()).account_type(AccountType::Public);
672
673 for component in components {
674 account_builder = account_builder.with_component(component);
675 }
676
677 self.add_account_from_builder(auth, account_builder, AccountState::Exists)
678 }
679
680 pub fn add_account(&mut self, account: Account) -> anyhow::Result<()> {
689 self.accounts.insert(account.id(), account);
690
691 Ok(())
694 }
695
696 pub fn add_output_note(&mut self, note: impl Into<RawOutputNote>) {
701 self.notes.push(note.into());
702 }
703
704 pub fn add_p2any_note(
709 &mut self,
710 sender_account_id: AccountId,
711 note_type: NoteType,
712 assets: impl IntoIterator<Item = Asset>,
713 ) -> anyhow::Result<Note> {
714 let note = create_p2any_note(sender_account_id, note_type, assets, &mut self.rng);
715 self.add_output_note(RawOutputNote::Full(note.clone()));
716
717 Ok(note)
718 }
719
720 pub fn add_p2id_note(
726 &mut self,
727 sender_account_id: AccountId,
728 target_account_id: AccountId,
729 asset: &[Asset],
730 note_type: NoteType,
731 ) -> Result<Note, NoteError> {
732 let note: Note = P2idNote::builder()
733 .sender(sender_account_id)
734 .target(target_account_id)
735 .assets(asset.iter().copied())
736 .note_type(note_type)
737 .generate_serial_number(&mut self.rng)
738 .build()?
739 .into();
740 self.add_output_note(RawOutputNote::Full(note.clone()));
741
742 Ok(note)
743 }
744
745 pub fn add_tx_fee_note(
750 &mut self,
751 sender_account_id: AccountId,
752 assets: &[Asset],
753 ) -> Result<Note, NoteError> {
754 let note: Note = TxFeeNote::builder()
755 .sender(sender_account_id)
756 .assets(assets.iter().copied())
757 .generate_serial_number(&mut self.rng)
758 .build()?
759 .into();
760 self.add_output_note(RawOutputNote::Full(note.clone()));
761
762 Ok(note)
763 }
764
765 pub fn add_p2ide_note(
773 &mut self,
774 sender_account_id: AccountId,
775 target_account_id: AccountId,
776 reclaimer: Option<AccountId>,
777 asset: &[Asset],
778 note_type: NoteType,
779 reclaim_height: Option<BlockNumber>,
780 timelock_height: Option<BlockNumber>,
781 ) -> Result<Note, NoteError> {
782 let note: Note = P2ideNote::builder()
783 .sender(sender_account_id)
784 .target(target_account_id)
785 .maybe_reclaimer(reclaimer)
786 .assets(asset.iter().copied())
787 .note_type(note_type)
788 .maybe_reclaim_height(reclaim_height)
789 .maybe_timelock_height(timelock_height)
790 .generate_serial_number(&mut self.rng)
791 .build()?
792 .into();
793
794 self.add_output_note(RawOutputNote::Full(note.clone()));
795
796 Ok(note)
797 }
798
799 pub fn add_swap_note(
801 &mut self,
802 sender: AccountId,
803 offered_asset: Asset,
804 requested_asset: Asset,
805 payback_note_type: NoteType,
806 ) -> anyhow::Result<(Note, NoteDetails)> {
807 let (swap_note, payback_note) = SwapNote::create(
808 sender,
809 offered_asset,
810 requested_asset,
811 NoteType::Public,
812 NoteAttachments::default(),
813 payback_note_type,
814 &mut self.rng,
815 )?;
816
817 self.add_output_note(RawOutputNote::Full(swap_note.clone()));
818
819 Ok((swap_note, payback_note))
820 }
821
822 pub fn add_spawn_note<'note, I>(
833 &mut self,
834 output_notes: impl IntoIterator<Item = &'note Note, IntoIter = I>,
835 ) -> anyhow::Result<Note>
836 where
837 I: ExactSizeIterator<Item = &'note Note>,
838 {
839 let note = create_spawn_note(output_notes)?;
840 self.add_output_note(RawOutputNote::Full(note.clone()));
841
842 Ok(note)
843 }
844
845 pub fn add_p2id_note_with_fee(
852 &mut self,
853 target_account_id: AccountId,
854 amount: u64,
855 ) -> anyhow::Result<Note> {
856 let fee_asset = self.fee_asset(amount)?;
857 let note = self.add_p2id_note(
858 self.fee_faucet_id,
859 target_account_id,
860 &[Asset::from(fee_asset)],
861 NoteType::Public,
862 )?;
863
864 Ok(note)
865 }
866
867 pub fn rng_mut(&mut self) -> &mut RandomCoin {
874 &mut self.rng
875 }
876
877 fn fee_asset(&self, amount: u64) -> anyhow::Result<FungibleAsset> {
879 FungibleAsset::new(self.fee_faucet_id, amount).context("failed to create fee asset")
880 }
881}
882
883impl Default for MockChainBuilder {
884 fn default() -> Self {
885 Self::new()
886 }
887}