Skip to main content

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