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
27#[derive(Clone)]
71pub struct MockTransactionBuilder<'chain> {
72 chain: &'chain MockChain,
73 input: TxContextInput,
74 reference_block: Option<BlockNumber>,
75 authenticated_notes: Vec<NoteId>,
76 unauthenticated_notes: Vec<Note>,
77 authenticator: Option<BasicAuthenticator>,
78 advice_inputs: AdviceInputs,
79 foreign_account_inputs: BTreeMap<AccountId, (Account, AccountWitness)>,
80 expected_output_notes: Vec<Note>,
81 tx_script: Option<TransactionScript>,
82 tx_script_args: Word,
83 auth_args: Word,
84 note_args: BTreeMap<NoteId, Word>,
85 signatures: Vec<(PublicKeyCommitment, Word, Signature)>,
86 note_scripts: BTreeMap<NoteScriptRoot, NoteScript>,
87 source_manager: Option<Arc<dyn SourceManagerSync>>,
88 is_lazy_loading_enabled: bool,
89}
90
91impl<'chain> MockTransactionBuilder<'chain> {
92 pub(crate) fn new(chain: &'chain MockChain, input: impl Into<TxContextInput>) -> Self {
96 let input = input.into();
97 let authenticator = chain.account_authenticator(input.id());
101
102 Self {
103 chain,
104 input,
105 reference_block: None,
106 authenticated_notes: Vec::new(),
107 unauthenticated_notes: Vec::new(),
108 authenticator,
109 advice_inputs: AdviceInputs::default(),
110 foreign_account_inputs: BTreeMap::new(),
111 expected_output_notes: Vec::new(),
112 tx_script: None,
113 tx_script_args: EMPTY_WORD,
114 auth_args: EMPTY_WORD,
115 note_args: BTreeMap::new(),
116 signatures: Vec::new(),
117 note_scripts: BTreeMap::new(),
118 source_manager: None,
119 is_lazy_loading_enabled: true,
120 }
121 }
122
123 pub fn authenticated_input_note(mut self, note_id: NoteId) -> Self {
128 self.authenticated_notes.push(note_id);
129 self
130 }
131
132 pub fn authenticated_input_notes(mut self, note_ids: impl IntoIterator<Item = NoteId>) -> Self {
136 self.authenticated_notes.extend(note_ids);
137 self
138 }
139
140 pub fn unauthenticated_input_note(mut self, note: Note) -> Self {
145 self.unauthenticated_notes.push(note);
146 self
147 }
148
149 pub fn unauthenticated_input_notes(mut self, notes: impl IntoIterator<Item = Note>) -> Self {
153 self.unauthenticated_notes.extend(notes);
154 self
155 }
156
157 pub fn reference_block(mut self, reference_block: impl Into<BlockNumber>) -> Self {
164 self.reference_block = Some(reference_block.into());
165 self
166 }
167
168 pub fn authenticator(mut self, authenticator: Option<BasicAuthenticator>) -> Self {
170 self.authenticator = authenticator;
171 self
172 }
173
174 pub fn extend_advice_inputs(mut self, advice_inputs: AdviceInputs) -> Self {
176 self.advice_inputs.extend(advice_inputs);
177 self
178 }
179
180 pub fn extend_advice_map(mut self, key: Word, value: Vec<Felt>) -> Self {
184 self.advice_inputs.map.insert(key, value);
185 self
186 }
187
188 pub fn foreign_accounts(
190 mut self,
191 inputs: impl IntoIterator<Item = (Account, AccountWitness)>,
192 ) -> Self {
193 self.foreign_account_inputs.extend(
194 inputs.into_iter().map(|(account, witness)| (account.id(), (account, witness))),
195 );
196 self
197 }
198
199 pub fn tx_script(mut self, tx_script: TransactionScript) -> Self {
201 self.tx_script = Some(tx_script);
202 self
203 }
204
205 pub fn tx_script_args(mut self, tx_script_args: Word) -> Self {
207 self.tx_script_args = tx_script_args;
208 self
209 }
210
211 pub fn auth_args(mut self, auth_args: Word) -> Self {
213 self.auth_args = auth_args;
214 self
215 }
216
217 pub fn extend_note_args(mut self, note_args: BTreeMap<NoteId, Word>) -> Self {
219 self.note_args.extend(note_args);
220 self
221 }
222
223 pub fn expected_output_note(mut self, output_note: RawOutputNote) -> Self {
228 if let RawOutputNote::Full(note) = output_note {
229 self.expected_output_notes.push(note);
230 }
231 self
232 }
233
234 pub fn expected_output_notes(mut self, output_notes: Vec<RawOutputNote>) -> Self {
238 let output_notes = output_notes.into_iter().filter_map(|note| match note {
239 RawOutputNote::Full(note) => Some(note),
240 RawOutputNote::Partial(_) => None,
241 });
242 self.expected_output_notes.extend(output_notes);
243 self
244 }
245
246 pub fn add_signature(
248 mut self,
249 pub_key: PublicKeyCommitment,
250 message: Word,
251 signature: Signature,
252 ) -> Self {
253 self.signatures.push((pub_key, message, signature));
254 self
255 }
256
257 pub fn add_note_script(mut self, script: NoteScript) -> Self {
259 self.note_scripts.insert(script.root(), script);
260 self
261 }
262
263 pub fn with_source_manager(mut self, source_manager: Arc<dyn SourceManagerSync>) -> Self {
268 self.source_manager = Some(source_manager);
269 self
270 }
271
272 pub fn disable_lazy_loading(mut self) -> Self {
277 self.is_lazy_loading_enabled = false;
278 self
279 }
280
281 pub fn build(self) -> anyhow::Result<TransactionContext> {
290 let account = self.chain.resolve_tx_account(self.input)?;
291
292 let mut tx_inputs = match self.reference_block {
293 Some(reference_block) => {
294 let latest_block = self.chain.latest_block_header().block_num();
295 anyhow::ensure!(
296 reference_block <= latest_block,
297 "reference block {reference_block} is out of range (latest {latest_block})",
298 );
299
300 self.chain.get_transaction_inputs_at(
301 reference_block,
302 &account,
303 &self.authenticated_notes,
304 &self.unauthenticated_notes,
305 )
306 },
307 None => self.chain.get_transaction_inputs(
308 &account,
309 &self.authenticated_notes,
310 &self.unauthenticated_notes,
311 ),
312 }
313 .context("failed to resolve transaction inputs from mock chain")?;
314
315 let mut tx_args = TransactionArgs::default().with_note_args(self.note_args);
316 if let Some(tx_script) = self.tx_script {
317 tx_args = tx_args.with_tx_script_and_args(tx_script, self.tx_script_args);
318 }
319 tx_args = tx_args.with_auth_args(self.auth_args);
320 tx_args.extend_advice_inputs(self.advice_inputs);
321 tx_args.extend_output_note_recipients(&self.expected_output_notes);
322 for (public_key_commitment, message, signature) in self.signatures {
323 tx_args.add_signature(public_key_commitment, message, signature);
324 }
325 tx_inputs.set_tx_args(tx_args);
326
327 let mast_store = TransactionMastStore::new();
328 mast_store.load_account_code(tx_inputs.account().code());
329 for (account, _) in self.foreign_account_inputs.values() {
330 mast_store.insert(account.code().mast());
331 }
332
333 let source_manager =
334 self.source_manager.unwrap_or_else(|| Arc::new(DefaultSourceManager::default()));
335
336 Ok(TransactionContext {
337 account,
338 expected_output_notes: self.expected_output_notes,
339 foreign_account_inputs: self.foreign_account_inputs,
340 tx_inputs,
341 mast_store,
342 authenticator: self.authenticator,
343 source_manager,
344 note_scripts: self.note_scripts,
345 is_lazy_loading_enabled: self.is_lazy_loading_enabled,
346 })
347 }
348}