Skip to main content

miden_testing/mock_transaction/
builder.rs

1// MOCK TRANSACTION BUILDER
2// ================================================================================================
3
4use alloc::collections::{BTreeMap, BTreeSet};
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, AccountCodeUpgrade, 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_standards::tx_script::SendNotesTransactionScript;
21use miden_tx::TransactionMastStore;
22use miden_tx::auth::BasicAuthenticator;
23
24use super::MockTransaction;
25use crate::MockChain;
26use crate::mock_chain::MockTransactionInput;
27
28// MOCK TRANSACTION BUILDER
29// ================================================================================================
30
31/// A builder for a [`MockTransaction`] that is coupled to a concrete [`MockChain`].
32///
33/// It is the public entry point for executing a transaction against a chain and is created through
34/// [`MockChain::build_transaction`]. Input notes are added explicitly through
35/// [`Self::authenticated_input_note`] and [`Self::unauthenticated_input_note`]. The transaction
36/// inputs are only resolved against the chain in [`Self::build`], once all input notes are known,
37/// by calling [`MockChain::get_transaction_inputs`].
38///
39/// # Examples
40///
41/// ```
42/// # use anyhow::Result;
43/// # use miden_protocol::{asset::FungibleAsset, note::NoteType};
44/// # use miden_testing::{Auth, MockChain};
45/// #
46/// # #[tokio::main(flavor = "current_thread")]
47/// # async fn main() -> Result<()> {
48/// let mut builder = MockChain::builder();
49/// let sender = builder.add_existing_mock_account(Auth::IncrNonce)?;
50/// let account = builder.add_existing_mock_account(Auth::IncrNonce)?;
51/// let note = builder.add_p2id_note(
52///     sender.id(),
53///     account.id(),
54///     &[FungibleAsset::mock(100)],
55///     NoteType::Public,
56/// )?;
57/// let chain = builder.build()?;
58///
59/// let executed = chain
60///     .build_transaction(account.id())
61///     .authenticated_input_note(note.id())
62///     .build()?
63///     .execute()
64///     .await?;
65///
66/// assert_eq!(executed.input_notes().num_notes(), 1);
67/// # Ok(())
68/// # }
69/// ```
70#[derive(Clone)]
71pub struct MockTransactionBuilder<'chain> {
72    chain: &'chain MockChain,
73    input: MockTransactionInput,
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    account_code_upgrade: Option<AccountCodeUpgrade>,
85    required_blocks: BTreeSet<BlockNumber>,
86    note_args: BTreeMap<NoteId, Word>,
87    signatures: Vec<(PublicKeyCommitment, Word, Signature)>,
88    note_scripts: BTreeMap<NoteScriptRoot, NoteScript>,
89    source_manager: Option<Arc<dyn SourceManagerSync>>,
90}
91
92impl<'chain> MockTransactionBuilder<'chain> {
93    /// Creates a new [`MockTransactionBuilder`] against the provided chain.
94    ///
95    /// Use [`MockChain::build_transaction`] instead of calling this directly.
96    pub(crate) fn new(chain: &'chain MockChain, input: impl Into<MockTransactionInput>) -> Self {
97        let input = input.into();
98        // Resolve the chain's authenticator for the account up front. The chain is borrowed
99        // immutably for the builder's lifetime, so this default cannot change before `build`; an
100        // explicit `authenticator` call may still override it.
101        let authenticator = chain.account_authenticator(input.id());
102
103        Self {
104            chain,
105            input,
106            reference_block: None,
107            authenticated_notes: Vec::new(),
108            unauthenticated_notes: Vec::new(),
109            authenticator,
110            advice_inputs: AdviceInputs::default(),
111            foreign_account_inputs: BTreeMap::new(),
112            expected_output_notes: Vec::new(),
113            tx_script: None,
114            tx_script_args: EMPTY_WORD,
115            auth_args: EMPTY_WORD,
116            account_code_upgrade: None,
117            required_blocks: BTreeSet::new(),
118            note_args: BTreeMap::new(),
119            signatures: Vec::new(),
120            note_scripts: BTreeMap::new(),
121            source_manager: None,
122        }
123    }
124
125    /// Adds an authenticated input note that the transaction consumes.
126    ///
127    /// The note must already be committed to the chain so that its inclusion proof can be resolved
128    /// in [`Self::build`].
129    pub fn authenticated_input_note(mut self, note_id: NoteId) -> Self {
130        self.authenticated_notes.push(note_id);
131        self
132    }
133
134    /// Adds multiple authenticated input notes that the transaction consumes.
135    ///
136    /// This is the iterator equivalent of [`Self::authenticated_input_note`].
137    pub fn authenticated_input_notes(mut self, note_ids: impl IntoIterator<Item = NoteId>) -> Self {
138        self.authenticated_notes.extend(note_ids);
139        self
140    }
141
142    /// Adds an unauthenticated input note that the transaction consumes.
143    ///
144    /// Contrary to [`Self::authenticated_input_note`], the note does not need to be committed to
145    /// the chain.
146    pub fn unauthenticated_input_note(mut self, note: Note) -> Self {
147        self.unauthenticated_notes.push(note);
148        self
149    }
150
151    /// Adds multiple unauthenticated input notes that the transaction consumes.
152    ///
153    /// This is the iterator equivalent of [`Self::unauthenticated_input_note`].
154    pub fn unauthenticated_input_notes(mut self, notes: impl IntoIterator<Item = Note>) -> Self {
155        self.unauthenticated_notes.extend(notes);
156        self
157    }
158
159    /// Sets the block the transaction executes against.
160    ///
161    /// By default the transaction is built against the chain's latest block. Use this to execute
162    /// against an earlier block instead, e.g. to test block-height-dependent script logic such as
163    /// timelocks or note expiration. All input notes must have been created at or before this
164    /// block.
165    pub fn reference_block(mut self, reference_block: impl Into<BlockNumber>) -> Self {
166        self.reference_block = Some(reference_block.into());
167        self
168    }
169
170    /// Set the authenticator for the transaction (if needed).
171    pub fn authenticator(mut self, authenticator: Option<BasicAuthenticator>) -> Self {
172        self.authenticator = authenticator;
173        self
174    }
175
176    /// Extends the advice inputs with the provided [`AdviceInputs`] instance.
177    pub fn extend_advice_inputs(mut self, advice_inputs: AdviceInputs) -> Self {
178        self.advice_inputs.extend(advice_inputs);
179        self
180    }
181
182    /// Inserts a single key-value pair into the advice inputs map.
183    ///
184    /// To add multiple entries, call this repeatedly or use [`Self::extend_advice_inputs`].
185    pub fn add_advice_map_entry(mut self, key: Word, value: Vec<Felt>) -> Self {
186        self.advice_inputs = self.advice_inputs.with_map([(key, value)]);
187        self
188    }
189
190    /// Sets foreign account inputs that are used by the transaction.
191    pub fn foreign_accounts(
192        mut self,
193        inputs: impl IntoIterator<Item = (Account, AccountWitness)>,
194    ) -> Self {
195        self.foreign_account_inputs.extend(
196            inputs.into_iter().map(|(account, witness)| (account.id(), (account, witness))),
197        );
198        self
199    }
200
201    /// Sets the desired transaction script.
202    pub fn tx_script(mut self, tx_script: TransactionScript) -> Self {
203        self.tx_script = Some(tx_script);
204        self
205    }
206
207    /// Sets the transaction script arguments.
208    pub fn tx_script_args(mut self, tx_script_args: Word) -> Self {
209        self.tx_script_args = tx_script_args;
210        self
211    }
212
213    /// Sets the transaction script and script arguments required to execute the provided
214    /// [`SendNotesTransactionScript`].
215    ///
216    /// The script's advice map entries are embedded in its MAST forest, so they load with the
217    /// script and need not be set here.
218    pub fn send_notes_script(self, script: &SendNotesTransactionScript) -> Self {
219        self.tx_script(script.tx_script().clone())
220            .tx_script_args(script.tx_script_args())
221    }
222
223    /// Sets the desired auth arguments.
224    pub fn auth_args(mut self, auth_args: Word) -> Self {
225        self.auth_args = auth_args;
226        self
227    }
228
229    /// Sets the code upgrade of the native account, which the transaction must provide if it
230    /// upgrades the account's code.
231    pub fn account_code_upgrade(mut self, account_code_upgrade: AccountCodeUpgrade) -> Self {
232        self.account_code_upgrade = Some(account_code_upgrade);
233        self
234    }
235
236    /// Requires the transaction's partial blockchain to track the provided block, so that the
237    /// executed code can read its commitment.
238    ///
239    /// The blocks the input notes were created in are tracked anyway. Blocks at or after the
240    /// reference block are ignored, see [`MockChain::get_transaction_inputs_at`].
241    pub fn required_block(mut self, block_num: BlockNumber) -> Self {
242        self.required_blocks.insert(block_num);
243        self
244    }
245
246    /// Extends the note arguments map with the provided one.
247    pub fn extend_note_args(mut self, note_args: BTreeMap<NoteId, Word>) -> Self {
248        self.note_args.extend(note_args);
249        self
250    }
251
252    /// Adds a single expected output note.
253    ///
254    /// A [`RawOutputNote::Partial`] note is ignored, since it does not carry the recipient details
255    /// required to reconstruct the note.
256    pub fn expected_output_note(mut self, output_note: RawOutputNote) -> Self {
257        if let RawOutputNote::Full(note) = output_note {
258            self.expected_output_notes.push(note);
259        }
260        self
261    }
262
263    /// Extends the expected output notes.
264    ///
265    /// This is the iterator equivalent of [`Self::expected_output_note`].
266    pub fn expected_output_notes(mut self, output_notes: Vec<RawOutputNote>) -> Self {
267        let output_notes = output_notes.into_iter().filter_map(|note| match note {
268            RawOutputNote::Full(note) => Some(note),
269            RawOutputNote::Partial(_) => None,
270        });
271        self.expected_output_notes.extend(output_notes);
272        self
273    }
274
275    /// Adds a new signature for the message and the public key.
276    pub fn add_signature(
277        mut self,
278        pub_key: PublicKeyCommitment,
279        message: Word,
280        signature: Signature,
281    ) -> Self {
282        self.signatures.push((pub_key, message, signature));
283        self
284    }
285
286    /// Adds a note script to the mock transaction for testing.
287    pub fn add_note_script(mut self, script: NoteScript) -> Self {
288        self.note_scripts.insert(script.root(), script);
289        self
290    }
291
292    /// Sets the [`SourceManagerSync`] on the [`MockTransaction`] that will be built.
293    ///
294    /// This source manager should contain the sources of all involved scripts and account code in
295    /// order to provide better error messages if an error occurs.
296    pub fn with_source_manager(mut self, source_manager: Arc<dyn SourceManagerSync>) -> Self {
297        self.source_manager = Some(source_manager);
298        self
299    }
300
301    /// Builds the [`MockTransaction`].
302    ///
303    /// The configured account and input notes are resolved into [`TransactionInputs`] against the
304    /// [`Self::reference_block`] (defaulting to the chain's latest block) through
305    /// [`MockChain::get_transaction_inputs`], and the remaining configuration is assembled into the
306    /// [`MockTransaction`].
307    ///
308    /// [`TransactionInputs`]: miden_protocol::transaction::TransactionInputs
309    pub fn build(self) -> anyhow::Result<MockTransaction> {
310        let account = self.chain.resolve_tx_account(self.input)?;
311
312        let latest_block = self.chain.latest_block_header().block_num();
313        let reference_block = self.reference_block.unwrap_or(latest_block);
314        anyhow::ensure!(
315            reference_block <= latest_block,
316            "reference block {reference_block} is out of range (latest {latest_block})",
317        );
318
319        let mut tx_inputs = self
320            .chain
321            .get_transaction_inputs_at(
322                reference_block,
323                &account,
324                &self.authenticated_notes,
325                &self.unauthenticated_notes,
326                self.required_blocks,
327            )
328            .context("failed to resolve transaction inputs from mock chain")?;
329
330        let mut tx_args = TransactionArgs::default().with_note_args(self.note_args);
331        if let Some(tx_script) = self.tx_script {
332            tx_args = tx_args.with_tx_script_and_args(tx_script, self.tx_script_args);
333        }
334        tx_args = tx_args.with_auth_args(self.auth_args);
335        if let Some(account_code_upgrade) = self.account_code_upgrade {
336            tx_args = tx_args.with_account_code_upgrade(account_code_upgrade);
337        }
338        tx_args.extend_advice_inputs(self.advice_inputs);
339        tx_args.extend_output_note_recipients(&self.expected_output_notes);
340        for (public_key_commitment, message, signature) in self.signatures {
341            tx_args.add_signature(public_key_commitment, message, signature);
342        }
343        tx_inputs.set_tx_args(tx_args);
344
345        let mast_store = TransactionMastStore::new();
346        mast_store.load_account_code(tx_inputs.account().code());
347        for (account, _) in self.foreign_account_inputs.values() {
348            mast_store.load_account_code(account.code());
349        }
350
351        let source_manager =
352            self.source_manager.unwrap_or_else(|| Arc::new(DefaultSourceManager::default()));
353
354        Ok(MockTransaction {
355            account,
356            expected_output_notes: self.expected_output_notes,
357            foreign_account_inputs: self.foreign_account_inputs,
358            tx_inputs,
359            mast_store,
360            authenticator: self.authenticator,
361            source_manager,
362            note_scripts: self.note_scripts,
363        })
364    }
365}