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::transaction::{
19    RawOutputNote,
20    TransactionArgs,
21    TransactionInputs,
22    TransactionScript,
23};
24use miden_tx::TransactionMastStore;
25use miden_tx::auth::BasicAuthenticator;
26
27use super::TransactionContext;
28use crate::MockChain;
29
30// TRANSACTION CONTEXT BUILDER
31// ================================================================================================
32
33/// [TransactionContextBuilder] is a utility to construct [TransactionContext] for testing
34/// purposes. It allows users to build accounts, create notes, provide advice inputs, and
35/// execute code. The VM process can be inspected afterward.
36///
37/// # Examples
38///
39/// Create a transaction context for an existing account and execute code:
40/// ```
41/// # use anyhow::Result;
42/// # use miden_protocol::Felt;
43/// # use miden_testing::{Auth, MockChain};
44/// #
45/// # #[tokio::main(flavor = "current_thread")]
46/// # async fn main() -> Result<()> {
47/// let mut builder = MockChain::builder();
48/// let account = builder.add_existing_mock_account(Auth::IncrNonce)?;
49/// let mock_chain = builder.build()?;
50/// let tx_context = mock_chain.build_tx_context(account.id(), &[], &[])?.build()?;
51///
52/// let code = "
53/// use miden::tx_kernel_core::prologue
54///
55/// begin
56///     exec.prologue::prepare_transaction
57///     push.5
58///     swap drop
59/// end
60/// ";
61///
62/// let exec_output = tx_context.execute_code(code).await?;
63/// assert_eq!(exec_output.stack.get(0).unwrap(), &Felt::from(5u32));
64/// # Ok(())
65/// # }
66/// ```
67#[derive(Clone)]
68pub struct TransactionContextBuilder {
69    source_manager: Arc<dyn SourceManagerSync>,
70    account: Account,
71    advice_inputs: AdviceInputs,
72    authenticator: Option<BasicAuthenticator>,
73    expected_output_notes: Vec<Note>,
74    foreign_account_inputs: BTreeMap<AccountId, (Account, AccountWitness)>,
75    input_notes: Vec<Note>,
76    tx_script: Option<TransactionScript>,
77    tx_script_args: Word,
78    note_args: BTreeMap<NoteId, Word>,
79    tx_inputs: Option<TransactionInputs>,
80    auth_args: Word,
81    signatures: Vec<(PublicKeyCommitment, Word, Signature)>,
82    note_scripts: BTreeMap<NoteScriptRoot, NoteScript>,
83    is_lazy_loading_enabled: bool,
84}
85
86impl TransactionContextBuilder {
87    pub fn new(account: Account) -> Self {
88        Self {
89            source_manager: Arc::new(DefaultSourceManager::default()),
90            account,
91            input_notes: Vec::new(),
92            expected_output_notes: Vec::new(),
93            tx_script: None,
94            tx_script_args: EMPTY_WORD,
95            authenticator: None,
96            advice_inputs: Default::default(),
97            tx_inputs: None,
98            note_args: BTreeMap::new(),
99            foreign_account_inputs: BTreeMap::new(),
100            auth_args: EMPTY_WORD,
101            signatures: Vec::new(),
102            note_scripts: BTreeMap::new(),
103            is_lazy_loading_enabled: true,
104        }
105    }
106
107    /// Extend the advice inputs with the provided [AdviceInputs] instance.
108    pub fn extend_advice_inputs(mut self, advice_inputs: AdviceInputs) -> Self {
109        self.advice_inputs.extend(advice_inputs);
110        self
111    }
112
113    /// Extend the advice inputs map with the provided iterator.
114    pub fn extend_advice_map(
115        mut self,
116        map_entries: impl IntoIterator<Item = (Word, Vec<Felt>)>,
117    ) -> Self {
118        self.advice_inputs.map.extend(map_entries);
119        self
120    }
121
122    /// Set the authenticator for the transaction (if needed)
123    pub fn authenticator(mut self, authenticator: Option<BasicAuthenticator>) -> Self {
124        self.authenticator = authenticator;
125        self
126    }
127
128    /// Set foreign account codes that are used by the transaction
129    pub fn foreign_accounts(
130        mut self,
131        inputs: impl IntoIterator<Item = (Account, AccountWitness)>,
132    ) -> Self {
133        self.foreign_account_inputs.extend(
134            inputs.into_iter().map(|(account, witness)| (account.id(), (account, witness))),
135        );
136        self
137    }
138
139    /// Set the desired transaction script
140    pub fn tx_script(mut self, tx_script: TransactionScript) -> Self {
141        self.tx_script = Some(tx_script);
142        self
143    }
144
145    /// Set the transaction script arguments
146    pub fn tx_script_args(mut self, tx_script_args: Word) -> Self {
147        self.tx_script_args = tx_script_args;
148        self
149    }
150
151    /// Set the desired auth arguments
152    pub fn auth_args(mut self, auth_args: Word) -> Self {
153        self.auth_args = auth_args;
154        self
155    }
156
157    /// Set the desired transaction inputs
158    pub fn tx_inputs(mut self, tx_inputs: TransactionInputs) -> Self {
159        assert_eq!(
160            AccountHeader::from(&self.account),
161            tx_inputs.account().into(),
162            "account in context and account provided via tx inputs are not the same account"
163        );
164        self.tx_inputs = Some(tx_inputs);
165        self
166    }
167
168    /// Disables lazy loading.
169    ///
170    /// Only affects [`TransactionContext::execute_code`] and causes the host to _not_ handle lazy
171    /// loading events.
172    pub fn disable_lazy_loading(mut self) -> Self {
173        self.is_lazy_loading_enabled = false;
174        self
175    }
176
177    /// Extend the note arguments map with the provided one.
178    pub fn extend_note_args(mut self, note_args: BTreeMap<NoteId, Word>) -> Self {
179        self.note_args.extend(note_args);
180        self
181    }
182
183    /// Extend the expected output notes.
184    pub fn extend_expected_output_notes(mut self, output_notes: Vec<RawOutputNote>) -> Self {
185        let output_notes = output_notes.into_iter().filter_map(|n| match n {
186            RawOutputNote::Full(note) => Some(note),
187            RawOutputNote::Partial(_) => None,
188        });
189
190        self.expected_output_notes.extend(output_notes);
191        self
192    }
193
194    /// Sets the [`SourceManagerSync`] on the [`TransactionContext`] that will be built.
195    ///
196    /// This source manager should contain the sources of all involved scripts and account code in
197    /// order to provide better error messages if an error occurs.
198    pub fn with_source_manager(mut self, source_manager: Arc<dyn SourceManagerSync>) -> Self {
199        self.source_manager = source_manager.clone();
200        self
201    }
202
203    /// Add a new signature for the message and the public key.
204    pub fn add_signature(
205        mut self,
206        pub_key: PublicKeyCommitment,
207        message: Word,
208        signature: Signature,
209    ) -> Self {
210        self.signatures.push((pub_key, message, signature));
211        self
212    }
213
214    /// Add a note script to the context for testing.
215    pub fn add_note_script(mut self, script: NoteScript) -> Self {
216        self.note_scripts.insert(script.root(), script);
217        self
218    }
219
220    /// Builds the [TransactionContext].
221    ///
222    /// If no transaction inputs were provided manually, an ad-hoc MockChain is created in order
223    /// to generate valid block data for the required notes.
224    pub fn build(self) -> anyhow::Result<TransactionContext> {
225        let mut tx_inputs = match self.tx_inputs {
226            Some(tx_inputs) => tx_inputs,
227            None => {
228                // If no specific transaction inputs was provided, initialize an ad-hoc mockchain
229                // to generate valid block header/MMR data
230
231                let mut builder = MockChain::builder();
232
233                // Get the set of note IDs in the provided order.
234                let input_note_ids: Vec<NoteId> = self.input_notes.iter().map(Note::id).collect();
235
236                for input_note in self.input_notes {
237                    builder.add_output_note(RawOutputNote::Full(input_note));
238                }
239                let mut mock_chain = builder.build()?;
240
241                mock_chain.prove_next_block().context("failed to prove first block")?;
242                mock_chain.prove_next_block().context("failed to prove second block")?;
243
244                mock_chain
245                    .get_transaction_inputs(&self.account, &input_note_ids, &[])
246                    .context("failed to get transaction inputs from mock chain")?
247            },
248        };
249
250        let mut tx_args = TransactionArgs::default().with_note_args(self.note_args);
251
252        tx_args = if let Some(tx_script) = self.tx_script {
253            tx_args.with_tx_script_and_args(tx_script, self.tx_script_args)
254        } else {
255            tx_args
256        };
257        tx_args = tx_args.with_auth_args(self.auth_args);
258        tx_args.extend_advice_inputs(self.advice_inputs.clone());
259        tx_args.extend_output_note_recipients(self.expected_output_notes.clone());
260
261        for (public_key_commitment, message, signature) in self.signatures {
262            tx_args.add_signature(public_key_commitment, message, signature);
263        }
264
265        tx_inputs.set_tx_args(tx_args);
266
267        let mast_store = {
268            let mast_forest_store = TransactionMastStore::new();
269            mast_forest_store.load_account_code(tx_inputs.account().code());
270
271            for (account, _) in self.foreign_account_inputs.values() {
272                mast_forest_store.load_account_code(account.code());
273            }
274
275            mast_forest_store
276        };
277
278        Ok(TransactionContext {
279            account: self.account,
280            expected_output_notes: self.expected_output_notes,
281            foreign_account_inputs: self.foreign_account_inputs,
282            tx_inputs,
283            mast_store,
284            authenticator: self.authenticator,
285            source_manager: self.source_manager,
286            note_scripts: self.note_scripts,
287            is_lazy_loading_enabled: self.is_lazy_loading_enabled,
288        })
289    }
290}
miden_testing/tx_context/builder.rs

miden_testing/tx_context/
builder.rs
