miden_testing/tx_context/
builder.rs1use alloc::collections::BTreeMap;
5use alloc::sync::Arc;
6use alloc::vec::Vec;
7
8use anyhow::Context;
9use miden_lib::testing::account_component::IncrNonceAuthComponent;
10use miden_lib::testing::mock_account::MockAccountExt;
11use miden_objects::EMPTY_WORD;
12use miden_objects::account::auth::{PublicKeyCommitment, Signature};
13use miden_objects::account::{Account, AccountHeader, AccountId};
14use miden_objects::assembly::DefaultSourceManager;
15use miden_objects::assembly::debuginfo::SourceManagerSync;
16use miden_objects::block::AccountWitness;
17use miden_objects::note::{Note, NoteId, NoteScript};
18use miden_objects::testing::account_id::ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE;
19use miden_objects::testing::noop_auth_component::NoopAuthComponent;
20use miden_objects::transaction::{
21 OutputNote,
22 TransactionArgs,
23 TransactionInputs,
24 TransactionScript,
25};
26use miden_processor::{AdviceInputs, Felt, Word};
27use miden_tx::TransactionMastStore;
28use miden_tx::auth::BasicAuthenticator;
29
30use super::TransactionContext;
31use crate::{MockChain, MockChainNote};
32
33pub struct TransactionContextBuilder {
70 source_manager: Arc<dyn SourceManagerSync>,
71 account: Account,
72 advice_inputs: AdviceInputs,
73 authenticator: Option<BasicAuthenticator>,
74 expected_output_notes: Vec<Note>,
75 foreign_account_inputs: BTreeMap<AccountId, (Account, AccountWitness)>,
76 input_notes: Vec<Note>,
77 tx_script: Option<TransactionScript>,
78 tx_script_args: Word,
79 note_args: BTreeMap<NoteId, Word>,
80 tx_inputs: Option<TransactionInputs>,
81 auth_args: Word,
82 signatures: Vec<(PublicKeyCommitment, Word, Signature)>,
83 is_lazy_loading_enabled: bool,
84 note_scripts: BTreeMap<Word, NoteScript>,
85}
86
87impl TransactionContextBuilder {
88 pub fn new(account: Account) -> Self {
89 Self {
90 source_manager: Arc::new(DefaultSourceManager::default()),
91 account,
92 input_notes: Vec::new(),
93 expected_output_notes: Vec::new(),
94 tx_script: None,
95 tx_script_args: EMPTY_WORD,
96 authenticator: None,
97 advice_inputs: Default::default(),
98 tx_inputs: None,
99 note_args: BTreeMap::new(),
100 foreign_account_inputs: BTreeMap::new(),
101 auth_args: EMPTY_WORD,
102 signatures: Vec::new(),
103 is_lazy_loading_enabled: true,
104 note_scripts: BTreeMap::new(),
105 }
106 }
107
108 pub fn with_existing_mock_account() -> Self {
118 Self::new(Account::mock(
119 ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE,
120 IncrNonceAuthComponent,
121 ))
122 }
123
124 pub fn with_noop_auth_account() -> Self {
126 let account =
127 Account::mock(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE, NoopAuthComponent);
128
129 Self::new(account)
130 }
131
132 pub fn with_fungible_faucet(acct_id: u128, initial_balance: Felt) -> Self {
134 let account = Account::mock_fungible_faucet(acct_id, initial_balance);
135 Self::new(account)
136 }
137
138 pub fn with_non_fungible_faucet(acct_id: u128) -> Self {
140 let account = Account::mock_non_fungible_faucet(acct_id);
141 Self::new(account)
142 }
143
144 pub fn extend_advice_inputs(mut self, advice_inputs: AdviceInputs) -> Self {
146 self.advice_inputs.extend(advice_inputs);
147 self
148 }
149
150 pub fn extend_advice_map(
152 mut self,
153 map_entries: impl IntoIterator<Item = (Word, Vec<Felt>)>,
154 ) -> Self {
155 self.advice_inputs.map.extend(map_entries);
156 self
157 }
158
159 pub fn authenticator(mut self, authenticator: Option<BasicAuthenticator>) -> Self {
161 self.authenticator = authenticator;
162 self
163 }
164
165 pub fn foreign_accounts(
167 mut self,
168 inputs: impl IntoIterator<Item = (Account, AccountWitness)>,
169 ) -> Self {
170 self.foreign_account_inputs.extend(
171 inputs.into_iter().map(|(account, witness)| (account.id(), (account, witness))),
172 );
173 self
174 }
175
176 pub fn extend_input_notes(mut self, input_notes: Vec<Note>) -> Self {
178 self.input_notes.extend(input_notes);
179 self
180 }
181
182 pub fn tx_script(mut self, tx_script: TransactionScript) -> Self {
184 self.tx_script = Some(tx_script);
185 self
186 }
187
188 pub fn tx_script_args(mut self, tx_script_args: Word) -> Self {
190 self.tx_script_args = tx_script_args;
191 self
192 }
193
194 pub fn auth_args(mut self, auth_args: Word) -> Self {
196 self.auth_args = auth_args;
197 self
198 }
199
200 pub fn tx_inputs(mut self, tx_inputs: TransactionInputs) -> Self {
202 assert_eq!(
203 AccountHeader::from(&self.account),
204 tx_inputs.account().into(),
205 "account in context and account provided via tx inputs are not the same account"
206 );
207 self.tx_inputs = Some(tx_inputs);
208 self
209 }
210
211 pub fn disable_lazy_loading(mut self) -> Self {
216 self.is_lazy_loading_enabled = false;
217 self
218 }
219
220 pub fn extend_note_args(mut self, note_args: BTreeMap<NoteId, Word>) -> Self {
222 self.note_args.extend(note_args);
223 self
224 }
225
226 pub fn extend_expected_output_notes(mut self, output_notes: Vec<OutputNote>) -> Self {
228 let output_notes = output_notes.into_iter().filter_map(|n| match n {
229 OutputNote::Full(note) => Some(note),
230 OutputNote::Partial(_) => None,
231 OutputNote::Header(_) => None,
232 });
233
234 self.expected_output_notes.extend(output_notes);
235 self
236 }
237
238 pub fn with_source_manager(mut self, source_manager: Arc<dyn SourceManagerSync>) -> Self {
243 self.source_manager = source_manager.clone();
244 self
245 }
246
247 pub fn add_signature(
249 mut self,
250 pub_key: PublicKeyCommitment,
251 message: Word,
252 signature: Signature,
253 ) -> Self {
254 self.signatures.push((pub_key, message, signature));
255 self
256 }
257
258 pub fn add_note_script(mut self, script: NoteScript) -> Self {
260 self.note_scripts.insert(script.root(), script);
261 self
262 }
263
264 pub fn build(self) -> anyhow::Result<TransactionContext> {
269 let mut tx_inputs = match self.tx_inputs {
270 Some(tx_inputs) => tx_inputs,
271 None => {
272 let mut builder = MockChain::builder();
276 for i in self.input_notes {
277 builder.add_output_note(OutputNote::Full(i));
278 }
279 let mut mock_chain = builder.build()?;
280
281 mock_chain.prove_next_block().context("failed to prove first block")?;
282 mock_chain.prove_next_block().context("failed to prove second block")?;
283
284 let input_note_ids: Vec<NoteId> =
285 mock_chain.committed_notes().values().map(MockChainNote::id).collect();
286
287 mock_chain
288 .get_transaction_inputs(&self.account, &input_note_ids, &[])
289 .context("failed to get transaction inputs from mock chain")?
290 },
291 };
292
293 let mut tx_args = TransactionArgs::default().with_note_args(self.note_args);
294
295 tx_args = if let Some(tx_script) = self.tx_script {
296 tx_args.with_tx_script_and_args(tx_script, self.tx_script_args)
297 } else {
298 tx_args
299 };
300 tx_args = tx_args.with_auth_args(self.auth_args);
301 tx_args.extend_advice_inputs(self.advice_inputs.clone());
302 tx_args.extend_output_note_recipients(self.expected_output_notes.clone());
303
304 for (public_key_commitment, message, signature) in self.signatures {
305 tx_args.add_signature(public_key_commitment, message, signature);
306 }
307
308 tx_inputs.set_tx_args(tx_args);
309
310 let mast_store = {
311 let mast_forest_store = TransactionMastStore::new();
312 mast_forest_store.load_account_code(tx_inputs.account().code());
313
314 for (account, _) in self.foreign_account_inputs.values() {
315 mast_forest_store.insert(account.code().mast());
316 }
317
318 mast_forest_store
319 };
320
321 Ok(TransactionContext {
322 account: self.account,
323 expected_output_notes: self.expected_output_notes,
324 foreign_account_inputs: self.foreign_account_inputs,
325 tx_inputs,
326 mast_store,
327 authenticator: self.authenticator,
328 source_manager: self.source_manager,
329 is_lazy_loading_enabled: self.is_lazy_loading_enabled,
330 note_scripts: self.note_scripts,
331 })
332 }
333}
334
335impl Default for TransactionContextBuilder {
336 fn default() -> Self {
337 Self::with_existing_mock_account()
338 }
339}