use alloc::collections::BTreeMap;
use alloc::vec::Vec;
use anyhow::Context;
const DEFAULT_FAUCET_DECIMALS: u8 = 10;
use itertools::Itertools;
use miden_processor::crypto::random::RandomCoin;
use miden_protocol::account::delta::AccountUpdateDetails;
use miden_protocol::account::{
Account,
AccountBuilder,
AccountComponent,
AccountDelta,
AccountId,
AccountStorageMode,
AccountType,
StorageSlot,
};
use miden_protocol::asset::{Asset, FungibleAsset, TokenSymbol};
use miden_protocol::block::account_tree::AccountTree;
use miden_protocol::block::nullifier_tree::NullifierTree;
use miden_protocol::block::{
BlockAccountUpdate,
BlockBody,
BlockHeader,
BlockNoteTree,
BlockNumber,
BlockProof,
Blockchain,
FeeParameters,
OutputNoteBatch,
ProvenBlock,
};
use miden_protocol::crypto::merkle::smt::Smt;
use miden_protocol::errors::NoteError;
use miden_protocol::note::{Note, NoteAttachment, NoteDetails, NoteType};
use miden_protocol::testing::account_id::ACCOUNT_ID_NATIVE_ASSET_FAUCET;
use miden_protocol::testing::random_secret_key::random_secret_key;
use miden_protocol::transaction::{OrderedTransactionHeaders, RawOutputNote, TransactionKernel};
use miden_protocol::{Felt, MAX_OUTPUT_NOTES_PER_BATCH, Word};
use miden_standards::account::access::Ownable2Step;
use miden_standards::account::faucets::{BasicFungibleFaucet, NetworkFungibleFaucet};
use miden_standards::account::mint_policies::{
AuthControlled,
OwnerControlled,
OwnerControlledInitConfig,
};
use miden_standards::account::wallets::BasicWallet;
use miden_standards::note::{P2idNote, P2ideNote, P2ideNoteStorage, SwapNote};
use miden_standards::testing::account_component::MockAccountComponent;
use rand::Rng;
use crate::mock_chain::chain::AccountAuthenticator;
use crate::utils::{create_p2any_note, create_spawn_note};
use crate::{AccountState, Auth, MockChain};
#[derive(Debug, Clone)]
pub struct MockChainBuilder {
accounts: BTreeMap<AccountId, Account>,
account_authenticators: BTreeMap<AccountId, AccountAuthenticator>,
notes: Vec<RawOutputNote>,
rng: RandomCoin,
native_asset_id: AccountId,
verification_base_fee: u32,
}
impl MockChainBuilder {
pub fn new() -> Self {
let native_asset_id =
ACCOUNT_ID_NATIVE_ASSET_FAUCET.try_into().expect("account ID should be valid");
Self {
accounts: BTreeMap::new(),
account_authenticators: BTreeMap::new(),
notes: Vec::new(),
rng: RandomCoin::new(Default::default()),
native_asset_id,
verification_base_fee: 0,
}
}
pub fn with_accounts(accounts: impl IntoIterator<Item = Account>) -> anyhow::Result<Self> {
let mut builder = Self::new();
for account in accounts {
builder.add_account(account)?;
}
Ok(builder)
}
pub fn native_asset_id(mut self, native_asset_id: AccountId) -> Self {
self.native_asset_id = native_asset_id;
self
}
pub fn verification_base_fee(mut self, verification_base_fee: u32) -> Self {
self.verification_base_fee = verification_base_fee;
self
}
pub fn build(self) -> anyhow::Result<MockChain> {
let block_account_updates: Vec<BlockAccountUpdate> = self
.accounts
.into_values()
.map(|account| {
let account_id = account.id();
let account_commitment = account.to_commitment();
let account_delta = AccountDelta::try_from(account)
.expect("chain builder should only store existing accounts without seeds");
let update_details = AccountUpdateDetails::Delta(account_delta);
BlockAccountUpdate::new(account_id, account_commitment, update_details)
})
.collect();
let account_tree = AccountTree::with_entries(
block_account_updates
.iter()
.map(|account| (account.account_id(), account.final_state_commitment())),
)
.context("failed to create genesis account tree")?;
let full_notes: Vec<Note> = self
.notes
.iter()
.filter_map(|note| match note {
RawOutputNote::Full(n) => Some(n.clone()),
_ => None,
})
.collect();
let proven_notes: Vec<_> = self
.notes
.into_iter()
.map(|note| note.into_output_note().expect("genesis note should be valid"))
.collect();
let note_chunks = proven_notes.into_iter().chunks(MAX_OUTPUT_NOTES_PER_BATCH);
let output_note_batches: Vec<OutputNoteBatch> = note_chunks
.into_iter()
.map(|batch_notes| batch_notes.into_iter().enumerate().collect::<Vec<_>>())
.collect();
let created_nullifiers = Vec::new();
let transactions = OrderedTransactionHeaders::new_unchecked(Vec::new());
let note_tree = BlockNoteTree::from_note_batches(&output_note_batches)
.context("failed to create block note tree")?;
let version = 0;
let prev_block_commitment = Word::empty();
let block_num = BlockNumber::from(0u32);
let chain_commitment = Blockchain::new().commitment();
let account_root = account_tree.root();
let nullifier_root = NullifierTree::<Smt>::default().root();
let note_root = note_tree.root();
let tx_commitment = transactions.commitment();
let tx_kernel_commitment = TransactionKernel.to_commitment();
let timestamp = MockChain::TIMESTAMP_START_SECS;
let fee_parameters = FeeParameters::new(self.native_asset_id, self.verification_base_fee)
.context("failed to construct fee parameters")?;
let validator_secret_key = random_secret_key();
let validator_public_key = validator_secret_key.public_key();
let header = BlockHeader::new(
version,
prev_block_commitment,
block_num,
chain_commitment,
account_root,
nullifier_root,
note_root,
tx_commitment,
tx_kernel_commitment,
validator_public_key,
fee_parameters,
timestamp,
);
let body = BlockBody::new_unchecked(
block_account_updates,
output_note_batches,
created_nullifiers,
transactions,
);
let signature = validator_secret_key.sign(header.commitment());
let block_proof = BlockProof::new_dummy();
let genesis_block = ProvenBlock::new_unchecked(header, body, signature, block_proof);
MockChain::from_genesis_block(
genesis_block,
account_tree,
self.account_authenticators,
validator_secret_key,
full_notes,
)
}
pub fn create_new_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
let account_builder = AccountBuilder::new(self.rng.random())
.storage_mode(AccountStorageMode::Public)
.with_component(BasicWallet);
self.add_account_from_builder(auth_method, account_builder, AccountState::New)
}
pub fn add_existing_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
self.add_existing_wallet_with_assets(auth_method, [])
}
pub fn add_existing_wallet_with_assets(
&mut self,
auth_method: Auth,
assets: impl IntoIterator<Item = Asset>,
) -> anyhow::Result<Account> {
let account_builder = Account::builder(self.rng.random())
.storage_mode(AccountStorageMode::Public)
.with_component(BasicWallet)
.with_assets(assets);
self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
}
pub fn create_new_faucet(
&mut self,
auth_method: Auth,
token_symbol: &str,
max_supply: u64,
) -> anyhow::Result<Account> {
let token_symbol = TokenSymbol::new(token_symbol)
.with_context(|| format!("invalid token symbol: {token_symbol}"))?;
let max_supply_felt = Felt::try_from(max_supply)?;
let basic_faucet =
BasicFungibleFaucet::new(token_symbol, DEFAULT_FAUCET_DECIMALS, max_supply_felt)
.context("failed to create BasicFungibleFaucet")?;
let account_builder = AccountBuilder::new(self.rng.random())
.storage_mode(AccountStorageMode::Public)
.account_type(AccountType::FungibleFaucet)
.with_component(basic_faucet)
.with_component(AuthControlled::allow_all());
self.add_account_from_builder(auth_method, account_builder, AccountState::New)
}
pub fn add_existing_basic_faucet(
&mut self,
auth_method: Auth,
token_symbol: &str,
max_supply: u64,
token_supply: Option<u64>,
) -> anyhow::Result<Account> {
let max_supply = Felt::try_from(max_supply)?;
let token_supply = Felt::try_from(token_supply.unwrap_or(0))?;
let token_symbol =
TokenSymbol::new(token_symbol).context("failed to create token symbol")?;
let basic_faucet =
BasicFungibleFaucet::new(token_symbol, DEFAULT_FAUCET_DECIMALS, max_supply)
.and_then(|fungible_faucet| fungible_faucet.with_token_supply(token_supply))
.context("failed to create basic fungible faucet")?;
let account_builder = AccountBuilder::new(self.rng.random())
.storage_mode(AccountStorageMode::Public)
.with_component(basic_faucet)
.with_component(AuthControlled::allow_all())
.account_type(AccountType::FungibleFaucet);
self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
}
pub fn add_existing_network_faucet(
&mut self,
token_symbol: &str,
max_supply: u64,
owner_account_id: AccountId,
token_supply: Option<u64>,
mint_policy: OwnerControlledInitConfig,
) -> anyhow::Result<Account> {
let max_supply = Felt::try_from(max_supply)?;
let token_supply = Felt::try_from(token_supply.unwrap_or(0))?;
let token_symbol =
TokenSymbol::new(token_symbol).context("failed to create token symbol")?;
let network_faucet =
NetworkFungibleFaucet::new(token_symbol, DEFAULT_FAUCET_DECIMALS, max_supply)
.and_then(|fungible_faucet| fungible_faucet.with_token_supply(token_supply))
.context("failed to create network fungible faucet")?;
let account_builder = AccountBuilder::new(self.rng.random())
.storage_mode(AccountStorageMode::Network)
.with_component(network_faucet)
.with_component(Ownable2Step::new(owner_account_id))
.with_component(OwnerControlled::new(mint_policy))
.account_type(AccountType::FungibleFaucet);
self.add_account_from_builder(Auth::IncrNonce, account_builder, AccountState::Exists)
}
pub fn create_new_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
let account_builder = Account::builder(self.rng.random())
.storage_mode(AccountStorageMode::Public)
.with_component(MockAccountComponent::with_empty_slots());
self.add_account_from_builder(auth_method, account_builder, AccountState::New)
}
pub fn add_existing_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
self.add_existing_mock_account_with_storage_and_assets(auth_method, [], [])
}
pub fn add_existing_mock_account_with_storage(
&mut self,
auth_method: Auth,
slots: impl IntoIterator<Item = StorageSlot>,
) -> anyhow::Result<Account> {
self.add_existing_mock_account_with_storage_and_assets(auth_method, slots, [])
}
pub fn add_existing_mock_account_with_assets(
&mut self,
auth_method: Auth,
assets: impl IntoIterator<Item = Asset>,
) -> anyhow::Result<Account> {
self.add_existing_mock_account_with_storage_and_assets(auth_method, [], assets)
}
pub fn add_existing_mock_account_with_storage_and_assets(
&mut self,
auth_method: Auth,
slots: impl IntoIterator<Item = StorageSlot>,
assets: impl IntoIterator<Item = Asset>,
) -> anyhow::Result<Account> {
let account_builder = Account::builder(self.rng.random())
.storage_mode(AccountStorageMode::Public)
.with_component(MockAccountComponent::with_slots(slots.into_iter().collect()))
.with_assets(assets);
self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
}
pub fn add_account_from_builder(
&mut self,
auth_method: Auth,
mut account_builder: AccountBuilder,
account_state: AccountState,
) -> anyhow::Result<Account> {
let (auth_component, authenticator) = auth_method.build_component();
account_builder = account_builder.with_auth_component(auth_component);
let account = if let AccountState::New = account_state {
account_builder.build().context("failed to build account from builder")?
} else {
account_builder
.build_existing()
.context("failed to build account from builder")?
};
self.account_authenticators
.insert(account.id(), AccountAuthenticator::new(authenticator));
if let AccountState::Exists = account_state {
self.accounts.insert(account.id(), account.clone());
}
Ok(account)
}
pub fn add_existing_account_from_components(
&mut self,
auth: Auth,
components: impl IntoIterator<Item = AccountComponent>,
) -> anyhow::Result<Account> {
let mut account_builder =
Account::builder(rand::rng().random()).storage_mode(AccountStorageMode::Public);
for component in components {
account_builder = account_builder.with_component(component);
}
self.add_account_from_builder(auth, account_builder, AccountState::Exists)
}
pub fn add_account(&mut self, account: Account) -> anyhow::Result<()> {
self.accounts.insert(account.id(), account);
Ok(())
}
pub fn add_output_note(&mut self, note: impl Into<RawOutputNote>) {
self.notes.push(note.into());
}
pub fn add_p2any_note(
&mut self,
sender_account_id: AccountId,
note_type: NoteType,
assets: impl IntoIterator<Item = Asset>,
) -> anyhow::Result<Note> {
let note = create_p2any_note(sender_account_id, note_type, assets, &mut self.rng);
self.add_output_note(RawOutputNote::Full(note.clone()));
Ok(note)
}
pub fn add_p2id_note(
&mut self,
sender_account_id: AccountId,
target_account_id: AccountId,
asset: &[Asset],
note_type: NoteType,
) -> Result<Note, NoteError> {
let note = P2idNote::create(
sender_account_id,
target_account_id,
asset.to_vec(),
note_type,
NoteAttachment::default(),
&mut self.rng,
)?;
self.add_output_note(RawOutputNote::Full(note.clone()));
Ok(note)
}
pub fn add_p2ide_note(
&mut self,
sender_account_id: AccountId,
target_account_id: AccountId,
asset: &[Asset],
note_type: NoteType,
reclaim_height: Option<BlockNumber>,
timelock_height: Option<BlockNumber>,
) -> Result<Note, NoteError> {
let storage = P2ideNoteStorage::new(target_account_id, reclaim_height, timelock_height);
let note = P2ideNote::create(
sender_account_id,
storage,
asset.to_vec(),
note_type,
Default::default(),
&mut self.rng,
)?;
self.add_output_note(RawOutputNote::Full(note.clone()));
Ok(note)
}
pub fn add_swap_note(
&mut self,
sender: AccountId,
offered_asset: Asset,
requested_asset: Asset,
payback_note_type: NoteType,
) -> anyhow::Result<(Note, NoteDetails)> {
let (swap_note, payback_note) = SwapNote::create(
sender,
offered_asset,
requested_asset,
NoteType::Public,
NoteAttachment::default(),
payback_note_type,
NoteAttachment::default(),
&mut self.rng,
)?;
self.add_output_note(RawOutputNote::Full(swap_note.clone()));
Ok((swap_note, payback_note))
}
pub fn add_spawn_note<'note, I>(
&mut self,
output_notes: impl IntoIterator<Item = &'note Note, IntoIter = I>,
) -> anyhow::Result<Note>
where
I: ExactSizeIterator<Item = &'note Note>,
{
let note = create_spawn_note(output_notes)?;
self.add_output_note(RawOutputNote::Full(note.clone()));
Ok(note)
}
pub fn add_p2id_note_with_fee(
&mut self,
target_account_id: AccountId,
amount: u64,
) -> anyhow::Result<Note> {
let fee_asset = self.native_fee_asset(amount)?;
let note = self.add_p2id_note(
self.native_asset_id,
target_account_id,
&[Asset::from(fee_asset)],
NoteType::Public,
)?;
Ok(note)
}
pub fn rng_mut(&mut self) -> &mut RandomCoin {
&mut self.rng
}
fn native_fee_asset(&self, amount: u64) -> anyhow::Result<FungibleAsset> {
FungibleAsset::new(self.native_asset_id, amount).context("failed to create fee asset")
}
}
impl Default for MockChainBuilder {
fn default() -> Self {
Self::new()
}
}