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_lib::account::faucets::{BasicFungibleFaucet, NetworkFungibleFaucet};
17use miden_lib::account::wallets::BasicWallet;
18use miden_lib::note::{create_p2id_note, create_p2ide_note, create_swap_note};
19use miden_lib::testing::account_component::MockAccountComponent;
20use miden_lib::transaction::{TransactionKernel, memory};
21use miden_objects::account::delta::AccountUpdateDetails;
22use miden_objects::account::{
23 Account,
24 AccountBuilder,
25 AccountDelta,
26 AccountId,
27 AccountStorageMode,
28 AccountType,
29 StorageSlot,
30};
31use miden_objects::asset::{Asset, FungibleAsset, TokenSymbol};
32use miden_objects::block::account_tree::AccountTree;
33use miden_objects::block::{
34 BlockAccountUpdate,
35 BlockHeader,
36 BlockNoteTree,
37 BlockNumber,
38 Blockchain,
39 FeeParameters,
40 NullifierTree,
41 OutputNoteBatch,
42 ProvenBlock,
43};
44use miden_objects::note::{Note, NoteDetails, NoteType};
45use miden_objects::testing::account_id::ACCOUNT_ID_NATIVE_ASSET_FAUCET;
46use miden_objects::transaction::{OrderedTransactionHeaders, OutputNote};
47use miden_objects::{Felt, FieldElement, MAX_OUTPUT_NOTES_PER_BATCH, NoteError, Word, ZERO};
48use miden_processor::crypto::RpoRandomCoin;
49use rand::Rng;
50
51use crate::mock_chain::chain::AccountAuthenticator;
52use crate::utils::{create_p2any_note, create_spawn_note};
53use crate::{AccountState, Auth, MockChain};
54
55#[derive(Debug, Clone)]
99pub struct MockChainBuilder {
100 accounts: BTreeMap<AccountId, Account>,
101 account_authenticators: BTreeMap<AccountId, AccountAuthenticator>,
102 notes: Vec<OutputNote>,
103 rng: RpoRandomCoin,
104 native_asset_id: AccountId,
106 verification_base_fee: u32,
107}
108
109impl MockChainBuilder {
110 pub fn new() -> Self {
120 let native_asset_id =
121 ACCOUNT_ID_NATIVE_ASSET_FAUCET.try_into().expect("account ID should be valid");
122
123 Self {
124 accounts: BTreeMap::new(),
125 account_authenticators: BTreeMap::new(),
126 notes: Vec::new(),
127 rng: RpoRandomCoin::new(Default::default()),
128 native_asset_id,
129 verification_base_fee: 0,
130 }
131 }
132
133 pub fn with_accounts(accounts: impl IntoIterator<Item = Account>) -> anyhow::Result<Self> {
142 let mut builder = Self::new();
143
144 for account in accounts {
145 builder.add_account(account)?;
146 }
147
148 Ok(builder)
149 }
150
151 pub fn native_asset_id(mut self, native_asset_id: AccountId) -> Self {
159 self.native_asset_id = native_asset_id;
160 self
161 }
162
163 pub fn verification_base_fee(mut self, verification_base_fee: u32) -> Self {
167 self.verification_base_fee = verification_base_fee;
168 self
169 }
170
171 pub fn build(self) -> anyhow::Result<MockChain> {
173 let block_account_updates: Vec<BlockAccountUpdate> = self
175 .accounts
176 .into_values()
177 .map(|account| {
178 let account_id = account.id();
179 let account_commitment = account.commitment();
180 let account_delta = AccountDelta::try_from(account)
181 .expect("chain builder should only store existing accounts without seeds");
182 let update_details = AccountUpdateDetails::Delta(account_delta);
183
184 BlockAccountUpdate::new(account_id, account_commitment, update_details)
185 })
186 .collect();
187
188 let account_tree = AccountTree::with_entries(
189 block_account_updates
190 .iter()
191 .map(|account| (account.account_id(), account.final_state_commitment())),
192 )
193 .context("failed to create genesis account tree")?;
194
195 let note_chunks = self.notes.into_iter().chunks(MAX_OUTPUT_NOTES_PER_BATCH);
196 let output_note_batches: Vec<OutputNoteBatch> = note_chunks
197 .into_iter()
198 .map(|batch_notes| batch_notes.into_iter().enumerate().collect::<Vec<_>>())
199 .collect();
200
201 let created_nullifiers = Vec::new();
202 let transactions = OrderedTransactionHeaders::new_unchecked(Vec::new());
203
204 let note_tree = BlockNoteTree::from_note_batches(&output_note_batches)
205 .context("failed to create block note tree")?;
206
207 let version = 0;
208 let prev_block_commitment = Word::empty();
209 let block_num = BlockNumber::from(0u32);
210 let chain_commitment = Blockchain::new().commitment();
211 let account_root = account_tree.root();
212 let nullifier_root = NullifierTree::new().root();
213 let note_root = note_tree.root();
214 let tx_commitment = transactions.commitment();
215 let tx_kernel_commitment = TransactionKernel.to_commitment();
216 let proof_commitment = Word::empty();
217 let timestamp = MockChain::TIMESTAMP_START_SECS;
218 let fee_parameters = FeeParameters::new(self.native_asset_id, self.verification_base_fee)
219 .context("failed to construct fee parameters")?;
220
221 let header = BlockHeader::new(
222 version,
223 prev_block_commitment,
224 block_num,
225 chain_commitment,
226 account_root,
227 nullifier_root,
228 note_root,
229 tx_commitment,
230 tx_kernel_commitment,
231 proof_commitment,
232 fee_parameters,
233 timestamp,
234 );
235
236 let genesis_block = ProvenBlock::new_unchecked(
237 header,
238 block_account_updates,
239 output_note_batches,
240 created_nullifiers,
241 transactions,
242 );
243
244 MockChain::from_genesis_block(genesis_block, account_tree, self.account_authenticators)
245 }
246
247 pub fn create_new_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
256 let account_builder = AccountBuilder::new(self.rng.random())
257 .storage_mode(AccountStorageMode::Public)
258 .with_component(BasicWallet);
259
260 self.add_account_from_builder(auth_method, account_builder, AccountState::New)
261 }
262
263 pub fn add_existing_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
266 self.add_existing_wallet_with_assets(auth_method, [])
267 }
268
269 pub fn add_existing_wallet_with_assets(
272 &mut self,
273 auth_method: Auth,
274 assets: impl IntoIterator<Item = Asset>,
275 ) -> anyhow::Result<Account> {
276 let account_builder = Account::builder(self.rng.random())
277 .storage_mode(AccountStorageMode::Public)
278 .with_component(BasicWallet)
279 .with_assets(assets);
280
281 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
282 }
283
284 pub fn create_new_faucet(
290 &mut self,
291 auth_method: Auth,
292 token_symbol: &str,
293 max_supply: u64,
294 ) -> anyhow::Result<Account> {
295 let token_symbol = TokenSymbol::new(token_symbol)
296 .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
297 let max_supply_felt = max_supply.try_into().map_err(|_| {
298 anyhow::anyhow!("max supply value cannot be converted to Felt: {max_supply}")
299 })?;
300 let basic_faucet =
301 BasicFungibleFaucet::new(token_symbol, DEFAULT_FAUCET_DECIMALS, max_supply_felt)
302 .context("failed to create BasicFungibleFaucet")?;
303
304 let account_builder = AccountBuilder::new(self.rng.random())
305 .storage_mode(AccountStorageMode::Public)
306 .account_type(AccountType::FungibleFaucet)
307 .with_component(basic_faucet);
308
309 self.add_account_from_builder(auth_method, account_builder, AccountState::New)
310 }
311
312 pub fn add_existing_basic_faucet(
317 &mut self,
318 auth_method: Auth,
319 token_symbol: &str,
320 max_supply: u64,
321 total_issuance: Option<u64>,
322 ) -> anyhow::Result<Account> {
323 let token_symbol = TokenSymbol::new(token_symbol).context("invalid argument")?;
324 let basic_faucet =
325 BasicFungibleFaucet::new(token_symbol, DEFAULT_FAUCET_DECIMALS, Felt::new(max_supply))
326 .context("invalid argument")?;
327
328 let account_builder = AccountBuilder::new(self.rng.random())
329 .storage_mode(AccountStorageMode::Public)
330 .with_component(basic_faucet)
331 .account_type(AccountType::FungibleFaucet);
332
333 let mut account =
334 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)?;
335
336 if let Some(issuance) = total_issuance {
339 account
340 .storage_mut()
341 .set_item(
342 memory::FAUCET_STORAGE_DATA_SLOT,
343 Word::from([ZERO, ZERO, ZERO, Felt::new(issuance)]),
344 )
345 .context("failed to set faucet storage")?;
346 self.accounts.insert(account.id(), account.clone());
347 }
348
349 Ok(account)
350 }
351
352 pub fn add_existing_network_faucet(
356 &mut self,
357 token_symbol: &str,
358 max_supply: u64,
359 owner_account_id: AccountId,
360 total_issuance: Option<u64>,
361 ) -> anyhow::Result<Account> {
362 let token_symbol = TokenSymbol::new(token_symbol).context("invalid argument")?;
363 let network_faucet = NetworkFungibleFaucet::new(
364 token_symbol,
365 DEFAULT_FAUCET_DECIMALS,
366 Felt::new(max_supply),
367 owner_account_id,
368 )
369 .context("invalid argument")?;
370
371 let account_builder = AccountBuilder::new(self.rng.random())
372 .storage_mode(AccountStorageMode::Network)
373 .with_component(network_faucet)
374 .account_type(AccountType::FungibleFaucet);
375
376 let mut account =
378 self.add_account_from_builder(Auth::IncrNonce, account_builder, AccountState::Exists)?;
379
380 if let Some(issuance) = total_issuance {
383 account
384 .storage_mut()
385 .set_item(
386 memory::FAUCET_STORAGE_DATA_SLOT,
387 Word::from([ZERO, ZERO, ZERO, Felt::new(issuance)]),
388 )
389 .context("failed to set faucet storage")?;
390 self.accounts.insert(account.id(), account.clone());
391 }
392
393 Ok(account)
394 }
395
396 pub fn create_new_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
399 let account_builder = Account::builder(self.rng.random())
400 .storage_mode(AccountStorageMode::Public)
401 .with_component(MockAccountComponent::with_empty_slots());
402
403 self.add_account_from_builder(auth_method, account_builder, AccountState::New)
404 }
405
406 pub fn add_existing_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
409 self.add_existing_mock_account_with_storage_and_assets(auth_method, [], [])
410 }
411
412 pub fn add_existing_mock_account_with_storage(
415 &mut self,
416 auth_method: Auth,
417 slots: impl IntoIterator<Item = StorageSlot>,
418 ) -> anyhow::Result<Account> {
419 self.add_existing_mock_account_with_storage_and_assets(auth_method, slots, [])
420 }
421
422 pub fn add_existing_mock_account_with_assets(
425 &mut self,
426 auth_method: Auth,
427 assets: impl IntoIterator<Item = Asset>,
428 ) -> anyhow::Result<Account> {
429 self.add_existing_mock_account_with_storage_and_assets(auth_method, [], assets)
430 }
431
432 pub fn add_existing_mock_account_with_storage_and_assets(
435 &mut self,
436 auth_method: Auth,
437 slots: impl IntoIterator<Item = StorageSlot>,
438 assets: impl IntoIterator<Item = Asset>,
439 ) -> anyhow::Result<Account> {
440 let account_builder = Account::builder(self.rng.random())
441 .storage_mode(AccountStorageMode::Public)
442 .with_component(MockAccountComponent::with_slots(slots.into_iter().collect()))
443 .with_assets(assets);
444
445 self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
446 }
447
448 pub fn add_account_from_builder(
459 &mut self,
460 auth_method: Auth,
461 mut account_builder: AccountBuilder,
462 account_state: AccountState,
463 ) -> anyhow::Result<Account> {
464 let (auth_component, authenticator) = auth_method.build_component();
465 account_builder = account_builder.with_auth_component(auth_component);
466
467 let account = if let AccountState::New = account_state {
468 account_builder.build().context("failed to build account from builder")?
469 } else {
470 account_builder
471 .build_existing()
472 .context("failed to build account from builder")?
473 };
474
475 self.account_authenticators
476 .insert(account.id(), AccountAuthenticator::new(authenticator));
477
478 if let AccountState::Exists = account_state {
479 self.accounts.insert(account.id(), account.clone());
480 }
481
482 Ok(account)
483 }
484
485 pub fn add_account(&mut self, account: Account) -> anyhow::Result<()> {
494 self.accounts.insert(account.id(), account);
495
496 Ok(())
499 }
500
501 pub fn add_output_note(&mut self, note: impl Into<OutputNote>) {
506 self.notes.push(note.into());
507 }
508
509 pub fn add_p2any_note(
514 &mut self,
515 sender_account_id: AccountId,
516 note_type: NoteType,
517 assets: impl IntoIterator<Item = Asset>,
518 ) -> anyhow::Result<Note> {
519 let note = create_p2any_note(sender_account_id, note_type, assets, &mut self.rng);
520 self.add_output_note(OutputNote::Full(note.clone()));
521
522 Ok(note)
523 }
524
525 pub fn add_p2id_note(
531 &mut self,
532 sender_account_id: AccountId,
533 target_account_id: AccountId,
534 asset: &[Asset],
535 note_type: NoteType,
536 ) -> Result<Note, NoteError> {
537 let note = create_p2id_note(
538 sender_account_id,
539 target_account_id,
540 asset.to_vec(),
541 note_type,
542 Felt::ZERO,
543 &mut self.rng,
544 )?;
545 self.add_output_note(OutputNote::Full(note.clone()));
546
547 Ok(note)
548 }
549
550 pub fn add_p2ide_note(
556 &mut self,
557 sender_account_id: AccountId,
558 target_account_id: AccountId,
559 asset: &[Asset],
560 note_type: NoteType,
561 reclaim_height: Option<BlockNumber>,
562 timelock_height: Option<BlockNumber>,
563 ) -> Result<Note, NoteError> {
564 let note = create_p2ide_note(
565 sender_account_id,
566 target_account_id,
567 asset.to_vec(),
568 reclaim_height,
569 timelock_height,
570 note_type,
571 Default::default(),
572 &mut self.rng,
573 )?;
574
575 self.add_output_note(OutputNote::Full(note.clone()));
576
577 Ok(note)
578 }
579
580 pub fn add_swap_note(
582 &mut self,
583 sender: AccountId,
584 offered_asset: Asset,
585 requested_asset: Asset,
586 payback_note_type: NoteType,
587 ) -> anyhow::Result<(Note, NoteDetails)> {
588 let (swap_note, payback_note) = create_swap_note(
589 sender,
590 offered_asset,
591 requested_asset,
592 NoteType::Public,
593 Felt::ZERO,
594 payback_note_type,
595 Felt::ZERO,
596 &mut self.rng,
597 )?;
598
599 self.add_output_note(OutputNote::Full(swap_note.clone()));
600
601 Ok((swap_note, payback_note))
602 }
603
604 pub fn add_spawn_note<'note, I>(
615 &mut self,
616 output_notes: impl IntoIterator<Item = &'note Note, IntoIter = I>,
617 ) -> anyhow::Result<Note>
618 where
619 I: ExactSizeIterator<Item = &'note Note>,
620 {
621 let note = create_spawn_note(output_notes)?;
622 self.add_output_note(OutputNote::Full(note.clone()));
623
624 Ok(note)
625 }
626
627 pub fn add_p2id_note_with_fee(
634 &mut self,
635 target_account_id: AccountId,
636 amount: u64,
637 ) -> anyhow::Result<Note> {
638 let fee_asset = self.native_fee_asset(amount)?;
639 let note = self.add_p2id_note(
640 self.native_asset_id,
641 target_account_id,
642 &[Asset::from(fee_asset)],
643 NoteType::Public,
644 )?;
645
646 Ok(note)
647 }
648
649 pub fn rng_mut(&mut self) -> &mut RpoRandomCoin {
656 &mut self.rng
657 }
658
659 fn native_fee_asset(&self, amount: u64) -> anyhow::Result<FungibleAsset> {
661 FungibleAsset::new(self.native_asset_id, amount).context("failed to create fee asset")
662 }
663}
664
665impl Default for MockChainBuilder {
666 fn default() -> Self {
667 Self::new()
668 }
669}