miden_testing/tx_context/
builder.rs

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