miden_testing/tx_context/
mock_transaction_builder.rs1use alloc::collections::BTreeMap;
5use alloc::sync::Arc;
6use alloc::vec::Vec;
7
8use anyhow::Context;
9use miden_processor::advice::AdviceInputs;
10use miden_processor::{Felt, Word};
11use miden_protocol::EMPTY_WORD;
12use miden_protocol::account::auth::{PublicKeyCommitment, Signature};
13use miden_protocol::account::{Account, AccountId};
14use miden_protocol::assembly::DefaultSourceManager;
15use miden_protocol::assembly::debuginfo::SourceManagerSync;
16use miden_protocol::block::BlockNumber;
17use miden_protocol::block::account_tree::AccountWitness;
18use miden_protocol::note::{Note, NoteId, NoteScript, NoteScriptRoot};
19use miden_protocol::transaction::{RawOutputNote, TransactionArgs, TransactionScript};
20use miden_tx::TransactionMastStore;
21use miden_tx::auth::BasicAuthenticator;
22
23use super::TransactionContext;
24use crate::MockChain;
25use crate::mock_chain::TxContextInput;
26
27pub struct MockTransactionBuilder<'chain> {
71 chain: &'chain MockChain,
72 input: TxContextInput,
73 reference_block: Option<BlockNumber>,
74 authenticated_notes: Vec<NoteId>,
75 unauthenticated_notes: Vec<Note>,
76 authenticator: Option<BasicAuthenticator>,
77 advice_inputs: AdviceInputs,
78 foreign_account_inputs: BTreeMap<AccountId, (Account, AccountWitness)>,
79 expected_output_notes: Vec<Note>,
80 tx_script: Option<TransactionScript>,
81 tx_script_args: Word,
82 auth_args: Word,
83 note_args: BTreeMap<NoteId, Word>,
84 signatures: Vec<(PublicKeyCommitment, Word, Signature)>,
85 note_scripts: BTreeMap<NoteScriptRoot, NoteScript>,
86 source_manager: Option<Arc<dyn SourceManagerSync>>,
87 is_lazy_loading_enabled: bool,
88}
89
90impl<'chain> MockTransactionBuilder<'chain> {
91 pub(crate) fn new(chain: &'chain MockChain, input: impl Into<TxContextInput>) -> Self {
95 let input = input.into();
96 let authenticator = chain.account_authenticator(input.id());
100
101 Self {
102 chain,
103 input,
104 reference_block: None,
105 authenticated_notes: Vec::new(),
106 unauthenticated_notes: Vec::new(),
107 authenticator,
108 advice_inputs: AdviceInputs::default(),
109 foreign_account_inputs: BTreeMap::new(),
110 expected_output_notes: Vec::new(),
111 tx_script: None,
112 tx_script_args: EMPTY_WORD,
113 auth_args: EMPTY_WORD,
114 note_args: BTreeMap::new(),
115 signatures: Vec::new(),
116 note_scripts: BTreeMap::new(),
117 source_manager: None,
118 is_lazy_loading_enabled: true,
119 }
120 }
121
122 pub fn authenticated_input_note(mut self, note_id: NoteId) -> Self {
127 self.authenticated_notes.push(note_id);
128 self
129 }
130
131 pub fn authenticated_input_notes(mut self, note_ids: impl IntoIterator<Item = NoteId>) -> Self {
135 self.authenticated_notes.extend(note_ids);
136 self
137 }
138
139 pub fn unauthenticated_input_note(mut self, note: Note) -> Self {
144 self.unauthenticated_notes.push(note);
145 self
146 }
147
148 pub fn unauthenticated_input_notes(mut self, notes: impl IntoIterator<Item = Note>) -> Self {
152 self.unauthenticated_notes.extend(notes);
153 self
154 }
155
156 pub fn reference_block(mut self, reference_block: impl Into<BlockNumber>) -> Self {
163 self.reference_block = Some(reference_block.into());
164 self
165 }
166
167 pub fn authenticator(mut self, authenticator: Option<BasicAuthenticator>) -> Self {
169 self.authenticator = authenticator;
170 self
171 }
172
173 pub fn extend_advice_inputs(mut self, advice_inputs: AdviceInputs) -> Self {
175 self.advice_inputs.extend(advice_inputs);
176 self
177 }
178
179 pub fn extend_advice_map(mut self, key: Word, value: Vec<Felt>) -> Self {
183 self.advice_inputs.map.insert(key, value);
184 self
185 }
186
187 pub fn foreign_accounts(
189 mut self,
190 inputs: impl IntoIterator<Item = (Account, AccountWitness)>,
191 ) -> Self {
192 self.foreign_account_inputs.extend(
193 inputs.into_iter().map(|(account, witness)| (account.id(), (account, witness))),
194 );
195 self
196 }
197
198 pub fn tx_script(mut self, tx_script: TransactionScript) -> Self {
200 self.tx_script = Some(tx_script);
201 self
202 }
203
204 pub fn tx_script_args(mut self, tx_script_args: Word) -> Self {
206 self.tx_script_args = tx_script_args;
207 self
208 }
209
210 pub fn auth_args(mut self, auth_args: Word) -> Self {
212 self.auth_args = auth_args;
213 self
214 }
215
216 pub fn extend_note_args(mut self, note_args: BTreeMap<NoteId, Word>) -> Self {
218 self.note_args.extend(note_args);
219 self
220 }
221
222 pub fn expected_output_note(mut self, output_note: RawOutputNote) -> Self {
227 if let RawOutputNote::Full(note) = output_note {
228 self.expected_output_notes.push(note);
229 }
230 self
231 }
232
233 pub fn expected_output_notes(mut self, output_notes: Vec<RawOutputNote>) -> Self {
237 let output_notes = output_notes.into_iter().filter_map(|note| match note {
238 RawOutputNote::Full(note) => Some(note),
239 RawOutputNote::Partial(_) => None,
240 });
241 self.expected_output_notes.extend(output_notes);
242 self
243 }
244
245 pub fn add_signature(
247 mut self,
248 pub_key: PublicKeyCommitment,
249 message: Word,
250 signature: Signature,
251 ) -> Self {
252 self.signatures.push((pub_key, message, signature));
253 self
254 }
255
256 pub fn add_note_script(mut self, script: NoteScript) -> Self {
258 self.note_scripts.insert(script.root(), script);
259 self
260 }
261
262 pub fn with_source_manager(mut self, source_manager: Arc<dyn SourceManagerSync>) -> Self {
267 self.source_manager = Some(source_manager);
268 self
269 }
270
271 pub fn disable_lazy_loading(mut self) -> Self {
276 self.is_lazy_loading_enabled = false;
277 self
278 }
279
280 pub fn build(self) -> anyhow::Result<TransactionContext> {
289 let account = self.chain.resolve_tx_account(self.input)?;
290
291 let mut tx_inputs = match self.reference_block {
292 Some(reference_block) => {
293 let latest_block = self.chain.latest_block_header().block_num();
294 anyhow::ensure!(
295 reference_block <= latest_block,
296 "reference block {reference_block} is out of range (latest {latest_block})",
297 );
298
299 self.chain.get_transaction_inputs_at(
300 reference_block,
301 &account,
302 &self.authenticated_notes,
303 &self.unauthenticated_notes,
304 )
305 },
306 None => self.chain.get_transaction_inputs(
307 &account,
308 &self.authenticated_notes,
309 &self.unauthenticated_notes,
310 ),
311 }
312 .context("failed to resolve transaction inputs from mock chain")?;
313
314 let mut tx_args = TransactionArgs::default().with_note_args(self.note_args);
315 if let Some(tx_script) = self.tx_script {
316 tx_args = tx_args.with_tx_script_and_args(tx_script, self.tx_script_args);
317 }
318 tx_args = tx_args.with_auth_args(self.auth_args);
319 tx_args.extend_advice_inputs(self.advice_inputs);
320 tx_args.extend_output_note_recipients(&self.expected_output_notes);
321 for (public_key_commitment, message, signature) in self.signatures {
322 tx_args.add_signature(public_key_commitment, message, signature);
323 }
324 tx_inputs.set_tx_args(tx_args);
325
326 let mast_store = TransactionMastStore::new();
327 mast_store.load_account_code(tx_inputs.account().code());
328 for (account, _) in self.foreign_account_inputs.values() {
329 mast_store.insert(account.code().mast());
330 }
331
332 let source_manager =
333 self.source_manager.unwrap_or_else(|| Arc::new(DefaultSourceManager::default()));
334
335 Ok(TransactionContext {
336 account,
337 expected_output_notes: self.expected_output_notes,
338 foreign_account_inputs: self.foreign_account_inputs,
339 tx_inputs,
340 mast_store,
341 authenticator: self.authenticator,
342 source_manager,
343 note_scripts: self.note_scripts,
344 is_lazy_loading_enabled: self.is_lazy_loading_enabled,
345 })
346 }
347}