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

miden_testing/tx_context/
builder.rs
