miden_testing/tx_context/
builder.rs

1// TRANSACTION CONTEXT BUILDER
2// ================================================================================================
3
4use alloc::{collections::BTreeMap, vec::Vec};
5
6use anyhow::Context;
7use miden_lib::transaction::TransactionKernel;
8use miden_objects::{
9    EMPTY_WORD, FieldElement,
10    account::Account,
11    assembly::Assembler,
12    note::{Note, NoteId},
13    testing::{
14        account_component::{IncrNonceAuthComponent, NoopAuthComponent},
15        account_id::ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE,
16    },
17    transaction::{
18        AccountInputs, OutputNote, TransactionArgs, TransactionInputs, TransactionScript,
19    },
20    vm::AdviceMap,
21};
22use miden_tx::{TransactionMastStore, auth::BasicAuthenticator};
23use rand_chacha::ChaCha20Rng;
24use vm_processor::{AdviceInputs, Felt, Word};
25
26use super::TransactionContext;
27use crate::{MockChain, MockChainNote};
28
29pub type MockAuthenticator = BasicAuthenticator<ChaCha20Rng>;
30
31// TRANSACTION CONTEXT BUILDER
32// ================================================================================================
33
34/// [TransactionContextBuilder] is a utility to construct [TransactionContext] for testing
35/// purposes. It allows users to build accounts, create notes, provide advice inputs, and
36/// execute code. The VM process can be inspected afterward.
37///
38/// # Examples
39///
40/// Create a new account and execute code:
41/// ```
42/// # use miden_testing::TransactionContextBuilder;
43/// # use miden_objects::{account::AccountBuilder,Felt, FieldElement};
44/// # use miden_lib::transaction::TransactionKernel;
45/// let tx_context = TransactionContextBuilder::with_existing_mock_account().build().unwrap();
46///
47/// let code = "
48/// use.kernel::prologue
49/// use.test::account
50///
51/// begin
52///     exec.prologue::prepare_transaction
53///     push.5
54///     swap drop
55/// end
56/// ";
57///
58/// let process = tx_context.execute_code(code).unwrap();
59/// assert_eq!(process.stack.get(0), Felt::new(5),);
60/// ```
61pub struct TransactionContextBuilder {
62    assembler: Assembler,
63    account: Account,
64    account_seed: Option<Word>,
65    advice_inputs: AdviceInputs,
66    authenticator: Option<MockAuthenticator>,
67    expected_output_notes: Vec<Note>,
68    foreign_account_inputs: Vec<AccountInputs>,
69    input_notes: Vec<Note>,
70    tx_script: Option<TransactionScript>,
71    tx_script_arg: Word,
72    note_args: BTreeMap<NoteId, Word>,
73    transaction_inputs: Option<TransactionInputs>,
74}
75
76impl TransactionContextBuilder {
77    pub fn new(account: Account) -> Self {
78        Self {
79            assembler: TransactionKernel::testing_assembler_with_mock_account(),
80            account,
81            account_seed: None,
82            input_notes: Vec::new(),
83            expected_output_notes: Vec::new(),
84            tx_script: None,
85            tx_script_arg: EMPTY_WORD,
86            authenticator: None,
87            advice_inputs: Default::default(),
88            transaction_inputs: None,
89            note_args: BTreeMap::new(),
90            foreign_account_inputs: vec![],
91        }
92    }
93
94    /// Initializes a [TransactionContextBuilder] with a mock account.
95    ///
96    /// The wallet:
97    ///
98    /// - Includes a series of mocked assets ([miden_objects::asset::AssetVault::mock()]).
99    /// - Has a nonce of `1` (so it does not imply seed validation).
100    /// - Has an ID of [`ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE`].
101    /// - Has an account code based on an
102    ///   [miden_objects::testing::account_component::AccountMockComponent].
103    pub fn with_existing_mock_account() -> Self {
104        // Build standard account with normal assembler because the testing one already contains it
105        let assembler = TransactionKernel::testing_assembler();
106        let auth_component =
107            IncrNonceAuthComponent::new(assembler.clone()).expect("valid component");
108
109        let account = Account::mock(
110            ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE,
111            Felt::ONE,
112            auth_component,
113            assembler,
114        );
115
116        let assembler = TransactionKernel::testing_assembler_with_mock_account();
117
118        Self {
119            assembler: assembler.clone(),
120            account,
121            account_seed: None,
122            authenticator: None,
123            input_notes: Vec::new(),
124            expected_output_notes: Vec::new(),
125            advice_inputs: Default::default(),
126            tx_script: None,
127            tx_script_arg: EMPTY_WORD,
128            transaction_inputs: None,
129            note_args: BTreeMap::new(),
130            foreign_account_inputs: vec![],
131        }
132    }
133
134    pub fn with_noop_auth_account(nonce: Felt) -> Self {
135        let assembler = TransactionKernel::testing_assembler();
136        let auth_component = NoopAuthComponent::new(assembler.clone()).expect("valid component");
137
138        let account = Account::mock(
139            ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE,
140            nonce,
141            auth_component,
142            assembler,
143        );
144
145        Self::new(account)
146    }
147
148    /// Initializes a [TransactionContextBuilder] with a mocked fungible faucet.
149    pub fn with_fungible_faucet(acct_id: u128, nonce: Felt, initial_balance: Felt) -> Self {
150        let account = Account::mock_fungible_faucet(
151            acct_id,
152            nonce,
153            initial_balance,
154            TransactionKernel::testing_assembler(),
155        );
156
157        Self { account, ..Self::default() }
158    }
159
160    /// Initializes a [TransactionContextBuilder] with a mocked non-fungible faucet.
161    pub fn with_non_fungible_faucet(acct_id: u128, nonce: Felt, empty_reserved_slot: bool) -> Self {
162        let account = Account::mock_non_fungible_faucet(
163            acct_id,
164            nonce,
165            empty_reserved_slot,
166            TransactionKernel::testing_assembler(),
167        );
168
169        Self { account, ..Self::default() }
170    }
171
172    /// Returns a clone of the assembler.
173    ///
174    /// This is primarily useful to assemble a script whose source will end up in the source manager
175    /// that is passed to the processor. This will help generate better error messages.
176    pub fn assembler(&self) -> Assembler {
177        self.assembler.clone()
178    }
179
180    /// Override and set the account seed manually
181    pub fn account_seed(mut self, account_seed: Option<Word>) -> Self {
182        self.account_seed = account_seed;
183        self
184    }
185
186    /// Extend the advice inputs with the provided [AdviceInputs] instance.
187    pub fn extend_advice_inputs(mut self, advice_inputs: AdviceInputs) -> Self {
188        self.advice_inputs.extend(advice_inputs);
189        self
190    }
191
192    /// Extend the advice inputs map with the provided iterator.
193    pub fn extend_advice_map(
194        mut self,
195        map_entries: impl IntoIterator<Item = (Word, Vec<Felt>)>,
196    ) -> Self {
197        self.advice_inputs
198            .extend_map(map_entries.into_iter().map(|(hash, input)| (hash.into(), input)));
199        self
200    }
201
202    /// Set the authenticator for the transaction (if needed)
203    pub fn authenticator(mut self, authenticator: Option<MockAuthenticator>) -> Self {
204        self.authenticator = authenticator;
205        self
206    }
207
208    /// Set foreign account codes that are used by the transaction
209    pub fn foreign_accounts(mut self, inputs: Vec<AccountInputs>) -> Self {
210        self.foreign_account_inputs = inputs;
211        self
212    }
213
214    /// Extend the set of used input notes
215    pub fn extend_input_notes(mut self, input_notes: Vec<Note>) -> Self {
216        self.input_notes.extend(input_notes);
217        self
218    }
219
220    /// Set the desired transaction script
221    pub fn tx_script(mut self, tx_script: TransactionScript) -> Self {
222        self.tx_script = Some(tx_script);
223        self
224    }
225
226    /// Set the transaction script argument
227    pub fn tx_script_arg(mut self, tx_script_arg: Word) -> Self {
228        self.tx_script_arg = tx_script_arg;
229        self
230    }
231
232    /// Set the desired transaction inputs
233    pub fn tx_inputs(mut self, tx_inputs: TransactionInputs) -> Self {
234        self.transaction_inputs = Some(tx_inputs);
235        self
236    }
237
238    /// Extend the note arguments map with the provided one.
239    pub fn extend_note_args(mut self, note_args: BTreeMap<NoteId, Word>) -> Self {
240        self.note_args.extend(note_args);
241        self
242    }
243
244    /// Extend the expected output notes.
245    pub fn extend_expected_output_notes(mut self, output_notes: Vec<OutputNote>) -> Self {
246        let output_notes = output_notes.into_iter().filter_map(|n| match n {
247            OutputNote::Full(note) => Some(note),
248            OutputNote::Partial(_) => None,
249            OutputNote::Header(_) => None,
250        });
251
252        self.expected_output_notes.extend(output_notes);
253        self
254    }
255
256    /// Builds the [TransactionContext].
257    ///
258    /// If no transaction inputs were provided manually, an ad-hoc MockChain is created in order
259    /// to generate valid block data for the required notes.
260    pub fn build(self) -> anyhow::Result<TransactionContext> {
261        let source_manager = self.assembler.source_manager();
262
263        let tx_inputs = match self.transaction_inputs {
264            Some(tx_inputs) => tx_inputs,
265            None => {
266                // If no specific transaction inputs was provided, initialize an ad-hoc mockchain
267                // to generate valid block header/MMR data
268
269                let mut mock_chain = MockChain::default();
270                for i in self.input_notes {
271                    mock_chain.add_pending_note(OutputNote::Full(i));
272                }
273
274                mock_chain.prove_next_block().context("failed to prove first block")?;
275                mock_chain.prove_next_block().context("failed to prove second block")?;
276
277                let input_note_ids: Vec<NoteId> =
278                    mock_chain.committed_notes().values().map(MockChainNote::id).collect();
279
280                mock_chain
281                    .get_transaction_inputs(
282                        self.account.clone(),
283                        self.account_seed,
284                        &input_note_ids,
285                        &[],
286                    )
287                    .context("failed to get transaction inputs from mock chain")?
288            },
289        };
290
291        let tx_args = TransactionArgs::new(AdviceMap::default(), self.foreign_account_inputs)
292            .with_note_args(self.note_args);
293
294        let mut tx_args = if let Some(tx_script) = self.tx_script {
295            tx_args.with_tx_script_and_arg(tx_script, self.tx_script_arg)
296        } else {
297            tx_args
298        };
299
300        tx_args.extend_advice_inputs(self.advice_inputs.clone());
301        tx_args.extend_output_note_recipients(self.expected_output_notes.clone());
302
303        let mast_store = {
304            let mast_forest_store = TransactionMastStore::new();
305            mast_forest_store.load_account_code(tx_inputs.account().code());
306
307            for acc_inputs in tx_args.foreign_account_inputs() {
308                mast_forest_store.insert(acc_inputs.code().mast());
309            }
310
311            mast_forest_store
312        };
313
314        Ok(TransactionContext {
315            expected_output_notes: self.expected_output_notes,
316            tx_args,
317            tx_inputs,
318            mast_store,
319            authenticator: self.authenticator,
320            advice_inputs: self.advice_inputs,
321            source_manager,
322        })
323    }
324}
325
326impl Default for TransactionContextBuilder {
327    fn default() -> Self {
328        Self::with_existing_mock_account()
329    }
330}