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    is_lazy_loading_enabled: bool,
84    note_scripts: BTreeMap<Word, NoteScript>,
85}
86
87impl TransactionContextBuilder {
88    pub fn new(account: Account) -> Self {
89        Self {
90            source_manager: Arc::new(DefaultSourceManager::default()),
91            account,
92            input_notes: Vec::new(),
93            expected_output_notes: Vec::new(),
94            tx_script: None,
95            tx_script_args: EMPTY_WORD,
96            authenticator: None,
97            advice_inputs: Default::default(),
98            tx_inputs: None,
99            note_args: BTreeMap::new(),
100            foreign_account_inputs: BTreeMap::new(),
101            auth_args: EMPTY_WORD,
102            signatures: Vec::new(),
103            is_lazy_loading_enabled: true,
104            note_scripts: BTreeMap::new(),
105        }
106    }
107
108    /// Initializes a [TransactionContextBuilder] with a mock account.
109    ///
110    /// The wallet:
111    ///
112    /// - Includes a series of mocked assets ([miden_protocol::asset::AssetVault::mock()]).
113    /// - Has a nonce of `1` (so it does not imply seed validation).
114    /// - Has an ID of [`ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE`].
115    /// - Has an account code based on an
116    ///   [miden_standards::testing::account_component::MockAccountComponent].
117    pub fn with_existing_mock_account() -> Self {
118        Self::new(Account::mock(
119            ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE,
120            IncrNonceAuthComponent,
121        ))
122    }
123
124    /// Same as [`Self::with_existing_mock_account`] but with a [`NoopAuthComponent`].
125    pub fn with_noop_auth_account() -> Self {
126        let account =
127            Account::mock(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE, NoopAuthComponent);
128
129        Self::new(account)
130    }
131
132    /// Initializes a [TransactionContextBuilder] with a mocked fungible faucet.
133    pub fn with_fungible_faucet(acct_id: u128, initial_balance: Felt) -> Self {
134        let account = Account::mock_fungible_faucet(acct_id, initial_balance);
135        Self::new(account)
136    }
137
138    /// Initializes a [TransactionContextBuilder] with a mocked non-fungible faucet.
139    pub fn with_non_fungible_faucet(acct_id: u128) -> Self {
140        let account = Account::mock_non_fungible_faucet(acct_id);
141        Self::new(account)
142    }
143
144    /// Extend the advice inputs with the provided [AdviceInputs] instance.
145    pub fn extend_advice_inputs(mut self, advice_inputs: AdviceInputs) -> Self {
146        self.advice_inputs.extend(advice_inputs);
147        self
148    }
149
150    /// Extend the advice inputs map with the provided iterator.
151    pub fn extend_advice_map(
152        mut self,
153        map_entries: impl IntoIterator<Item = (Word, Vec<Felt>)>,
154    ) -> Self {
155        self.advice_inputs.map.extend(map_entries);
156        self
157    }
158
159    /// Set the authenticator for the transaction (if needed)
160    pub fn authenticator(mut self, authenticator: Option<BasicAuthenticator>) -> Self {
161        self.authenticator = authenticator;
162        self
163    }
164
165    /// Set foreign account codes that are used by the transaction
166    pub fn foreign_accounts(
167        mut self,
168        inputs: impl IntoIterator<Item = (Account, AccountWitness)>,
169    ) -> Self {
170        self.foreign_account_inputs.extend(
171            inputs.into_iter().map(|(account, witness)| (account.id(), (account, witness))),
172        );
173        self
174    }
175
176    /// Extend the set of used input notes
177    pub fn extend_input_notes(mut self, input_notes: Vec<Note>) -> Self {
178        self.input_notes.extend(input_notes);
179        self
180    }
181
182    /// Set the desired transaction script
183    pub fn tx_script(mut self, tx_script: TransactionScript) -> Self {
184        self.tx_script = Some(tx_script);
185        self
186    }
187
188    /// Set the transaction script arguments
189    pub fn tx_script_args(mut self, tx_script_args: Word) -> Self {
190        self.tx_script_args = tx_script_args;
191        self
192    }
193
194    /// Set the desired auth arguments
195    pub fn auth_args(mut self, auth_args: Word) -> Self {
196        self.auth_args = auth_args;
197        self
198    }
199
200    /// Set the desired transaction inputs
201    pub fn tx_inputs(mut self, tx_inputs: TransactionInputs) -> Self {
202        assert_eq!(
203            AccountHeader::from(&self.account),
204            tx_inputs.account().into(),
205            "account in context and account provided via tx inputs are not the same account"
206        );
207        self.tx_inputs = Some(tx_inputs);
208        self
209    }
210
211    /// Disables lazy loading.
212    ///
213    /// Only affects [`TransactionContext::execute_code`] and causes the host to _not_ handle lazy
214    /// loading events.
215    pub fn disable_lazy_loading(mut self) -> Self {
216        self.is_lazy_loading_enabled = false;
217        self
218    }
219
220    /// Extend the note arguments map with the provided one.
221    pub fn extend_note_args(mut self, note_args: BTreeMap<NoteId, Word>) -> Self {
222        self.note_args.extend(note_args);
223        self
224    }
225
226    /// Extend the expected output notes.
227    pub fn extend_expected_output_notes(mut self, output_notes: Vec<OutputNote>) -> Self {
228        let output_notes = output_notes.into_iter().filter_map(|n| match n {
229            OutputNote::Full(note) => Some(note),
230            OutputNote::Partial(_) => None,
231            OutputNote::Header(_) => None,
232        });
233
234        self.expected_output_notes.extend(output_notes);
235        self
236    }
237
238    /// Sets the [`SourceManagerSync`] on the [`TransactionContext`] that will be built.
239    ///
240    /// This source manager should contain the sources of all involved scripts and account code in
241    /// order to provide better error messages if an error occurs.
242    pub fn with_source_manager(mut self, source_manager: Arc<dyn SourceManagerSync>) -> Self {
243        self.source_manager = source_manager.clone();
244        self
245    }
246
247    /// Add a new signature for the message and the public key.
248    pub fn add_signature(
249        mut self,
250        pub_key: PublicKeyCommitment,
251        message: Word,
252        signature: Signature,
253    ) -> Self {
254        self.signatures.push((pub_key, message, signature));
255        self
256    }
257
258    /// Add a note script to the context for testing.
259    pub fn add_note_script(mut self, script: NoteScript) -> Self {
260        self.note_scripts.insert(script.root(), script);
261        self
262    }
263
264    /// Builds the [TransactionContext].
265    ///
266    /// If no transaction inputs were provided manually, an ad-hoc MockChain is created in order
267    /// to generate valid block data for the required notes.
268    pub fn build(self) -> anyhow::Result<TransactionContext> {
269        let mut tx_inputs = match self.tx_inputs {
270            Some(tx_inputs) => tx_inputs,
271            None => {
272                // If no specific transaction inputs was provided, initialize an ad-hoc mockchain
273                // to generate valid block header/MMR data
274
275                let mut builder = MockChain::builder();
276                for i in self.input_notes {
277                    builder.add_output_note(OutputNote::Full(i));
278                }
279                let mut mock_chain = builder.build()?;
280
281                mock_chain.prove_next_block().context("failed to prove first block")?;
282                mock_chain.prove_next_block().context("failed to prove second block")?;
283
284                let input_note_ids: Vec<NoteId> =
285                    mock_chain.committed_notes().values().map(MockChainNote::id).collect();
286
287                mock_chain
288                    .get_transaction_inputs(&self.account, &input_note_ids, &[])
289                    .context("failed to get transaction inputs from mock chain")?
290            },
291        };
292
293        let mut tx_args = TransactionArgs::default().with_note_args(self.note_args);
294
295        tx_args = if let Some(tx_script) = self.tx_script {
296            tx_args.with_tx_script_and_args(tx_script, self.tx_script_args)
297        } else {
298            tx_args
299        };
300        tx_args = tx_args.with_auth_args(self.auth_args);
301        tx_args.extend_advice_inputs(self.advice_inputs.clone());
302        tx_args.extend_output_note_recipients(self.expected_output_notes.clone());
303
304        for (public_key_commitment, message, signature) in self.signatures {
305            tx_args.add_signature(public_key_commitment, message, signature);
306        }
307
308        tx_inputs.set_tx_args(tx_args);
309
310        let mast_store = {
311            let mast_forest_store = TransactionMastStore::new();
312            mast_forest_store.load_account_code(tx_inputs.account().code());
313
314            for (account, _) in self.foreign_account_inputs.values() {
315                mast_forest_store.insert(account.code().mast());
316            }
317
318            mast_forest_store
319        };
320
321        Ok(TransactionContext {
322            account: self.account,
323            expected_output_notes: self.expected_output_notes,
324            foreign_account_inputs: self.foreign_account_inputs,
325            tx_inputs,
326            mast_store,
327            authenticator: self.authenticator,
328            source_manager: self.source_manager,
329            is_lazy_loading_enabled: self.is_lazy_loading_enabled,
330            note_scripts: self.note_scripts,
331        })
332    }
333}
334
335impl Default for TransactionContextBuilder {
336    fn default() -> Self {
337        Self::with_existing_mock_account()
338    }
339}
miden_testing/tx_context/builder.rs

miden_testing/tx_context/
builder.rs
